Distribute PostgreSQL 17 with Citus 13

-- create, distribute and populate the table with a jsonb column
-- we will use that column as a regular postgres table through JSON_TABLE()
CREATE TABLE my_favorite_books(book_collection_id bigserial, jsonb_column jsonb);
SELECT create_distributed_table(' my_favorite_books ', 'book_collection_id');

-- these books will be inserted with the automatic book_collection_id of 1
INSERT INTO my_favorite_books (jsonb_column) VALUES (
'{ "favorites" : [
   { "kind" : "mystery", "books" : [
     { "title" : "The Count of Monte Cristo", "author" : "Alexandre Dumas"},
     { "title" : "Crime and Punishment", "author" : "Fyodor Dostoevsky" } ] },
   { "kind" : "drama", "books" : [
     { "title" : "Anna Karenina", "author" : "Leo Tolstoy" } ] },
   { "kind" : "poetry", "books" : [
     { "title" : "Masnavi", "author" : "Jalal al-Din Muhammad Rumi" } ] },
   { "kind" : "autobiography", "books" : [
     { "title" : "The Autobiography of Malcolm X", "author" : "Alex Haley" } ] }
  ] }');

-- these books will be inserted with the automatic book_collection_id of 2
INSERT INTO my_favorite_books (jsonb_column) VALUES (
'{ "favorites" : [
   { "kind" : "mystery", "books" : [
     { "title" : "To Kill a Mockingbird", "author" : "Harper Lee"},
     { "title" : "Our Mutual Friend", "author" : "Charles Dickens" } ] },
   { "kind" : "drama", "books" : [
     { "title" : "Pride and Prejudice", "author" : "Jane Austen" } ] },
   { "kind" : "poetry", "books" : [
     { "title" : "The Odyssey", "author" : "Homer" } ] },
   { "kind" : "autobiography", "books" : [
     { "title" : "The Diary of a Young Girl", "author" : "Anne Frank" } ] }
  ] }');

-- a simple router query, that outputs all the books under book_collection_id = 1
SELECT json_table_output.* FROM
my_favorite_books,
JSON_TABLE ( jsonb_column, '$.favorites[*]' COLUMNS (
   key FOR ORDINALITY, kind text PATH '$.kind',
   NESTED PATH '$.books[*]' COLUMNS (
     title text PATH '$.title', author text PATH '$.author'))) AS json_table_output
WHERE my_favorite_books. book_collection_id = 1
ORDER BY 1, 2, 3, 4;

 key |     kind      |             title              |           author
-----+---------------+--------------------------------+----------------------------
   1 | mystery       | The Count of Monte Cristo      | Alexandre Dumas
   1 | mystery       | Crime and Punishment           | Fyodor Dostoevsky
   2 | drama         | Anna Karenina                  | Leo Tolstoy
   3 | poetry        | Masnavi                        | Jalal al-Din Muhammad Rumi
   4 | autobiography | The Autobiography of Malcolm X | Alex Haley
(5 rows)

-- a simple multi-shard query, where we want to see all the books
SELECT json_table_output.* FROM
my_favorite_books,
JSON_TABLE ( jsonb_column, '$.favorites[*]' COLUMNS (
   key FOR ORDINALITY, kind text PATH '$.kind',
   NESTED PATH '$.books[*]' COLUMNS (
     title text PATH '$.title', author text PATH '$.author'))) AS json_table_output
ORDER BY 1, 2, 3, 4;

 key |     kind      |             title              |           author
-----+---------------+--------------------------------+----------------------------
   1 | mystery       | The Count of Monte Cristo      | Alexandre Dumas
   1 | mystery       | Crime and Punishment           | Fyodor Dostoevsky
   1 | mystery       | Our Mutual Friend              | Charles Dickens
   1 | mystery       | To Kill a Mockingbird          | Harper Lee
   2 | drama         | Anna Karenina                  | Leo Tolstoy
   2 | drama         | Pride and Prejudice            | Jane Austen
   3 | poetry        | Masnavi                        | Jalal al-Din Muhammad Rumi
   3 | poetry        | The Odyssey                    | Homer
   4 | autobiography | The Autobiography of Malcolm X | Alex Haley
   4 | autobiography | The Diary of a Young Girl      | Anne Frank
(10 rows)

-- more complex router query involving LATERAL and LIMIT
-- select two books under book_collection_id = 2
SELECT sub.*
FROM my_favorite_books,
lateral(SELECT * FROM JSON_TABLE (jsonb_column, '$.favorites[*]'
    COLUMNS (key FOR ORDINALITY, kind text PATH '$.kind',
        NESTED PATH '$.books[*]' COLUMNS
            (title text PATH '$.title', author text PATH '$.author')))
    AS json_table_output ORDER BY key DESC LIMIT 2) as sub
WHERE my_favorite_books.book_collection_id = 2;

 key |     kind      |           title           |   author
-----+---------------+---------------------------+------------
   4 | autobiography | The Diary of a Young Girl | Anne Frank
   3 | poetry        | The Odyssey               | Homer
(2 rows)

-- create and distribute the target and source tables
CREATE TABLE target_table (tid integer, balance float, val text);
CREATE TABLE source_table (sid integer, delta float);
SELECT create_distributed_table('target_table', 'tid');
SELECT create_distributed_table('source_table', 'sid');

-- populate the tables
INSERT INTO target_table SELECT id, id * 100, 'initial' FROM generate_series(1,5,2) AS id;
INSERT INTO source_table SELECT id, id * 10 FROM generate_series(1,4) AS id;

-- Use WHEN NOT MATCHED BY SOURCE
MERGE INTO target_table t
    USING source_table s
    ON t.tid = s.sid AND tid = 1
    WHEN MATCHED THEN
        UPDATE SET balance = balance + delta, val = val || ' updated by merge'
    WHEN NOT MATCHED BY TARGET THEN
        INSERT VALUES (sid, delta, 'inserted by merge')
    WHEN NOT MATCHED BY SOURCE THEN
        UPDATE SET val = val || ' not matched by source';

-- see the updated distributed target table
SELECT * FROM target_table ORDER BY tid;

 tid | balance |              val
-----+---------+-------------------------------
   1 |     110 | initial updated by merge
   2 |      20 | inserted by merge
   3 |      30 | inserted by merge
   3 |     300 | initial not matched by source
   4 |      40 | inserted by merge
   5 |     500 | initial not matched by source
(6 rows)

-- let's say we are at node 0
-- create a partitioned table
-- specify access method as columnar, use generated identity column
CREATE TABLE dist_partitioned_table
( id_test bigint GENERATED BY DEFAULT AS IDENTITY (START WITH 10 INCREMENT BY 10),
n int )
PARTITION BY RANGE (n)
USING columnar;

-- create a partition for the table
CREATE TABLE pt_1 PARTITION OF dist_partitioned_table FOR VALUES FROM (1) TO (50);

-- distribute the table, making it a distributed partitioned table
SELECT create_distributed_table('dist_partitioned_table', 'id_test');

-- create another partition of the table, it will be automatically distributed
CREATE TABLE pt_2 PARTITION OF dist_partitioned_table FOR VALUES FROM (50) TO (1000);

-- Altering an access method for a partitioned table lets the value be used
-- for all future partitions created under it.
-- Existing partitions are not modified.
ALTER TABLE dist_partitioned_table SET ACCESS METHOD heap;

-- Add an exclusion constraint, which will be part of current and future partitions as well
ALTER TABLE dist_partitioned_table ADD EXCLUDE USING btree (id_test WITH =, n WITH =);

-- Attaching a partition inherits the identity column from the parent table
CREATE TABLE pt_3 (id_test bigint not null, n int);
ALTER TABLE dist_partitioned_table ATTACH PARTITION pt_3 FOR VALUES FROM (1000) TO (2000);

-- verify that the identity column is inherited in all children
SELECT attrelid::regclass, attname, attidentity FROM pg_attribute
WHERE attname = 'id_test' AND attidentity = 'd' ORDER BY 1;

        attrelid        | attname | attidentity
------------------------+---------+-------------
 dist_partitioned_table | id_test | d
 pt_1                   | id_test | d
 pt_2                   | id_test | d
 pt_3                   | id_test | d
(4 rows)

-- the parent table and the new partition have the altered access method "heap"
-- whereas the old two partitions have the original access method "columnar"
SELECT relname, amname FROM pg_class c LEFT JOIN pg_am am ON (c.relam = am.oid)
WHERE relname IN ('dist_partitioned_table', 'pt_1', 'pt_2', 'pt_3') ORDER BY relname;

        relname         |  amname
------------------------+----------
 dist_partitioned_table | heap
 pt_1                   | columnar
 pt_2                   | columnar
 pt_3                   | heap
(4 rows)

-- verify that the distributed partitioned table and its distributed partitions
-- have exclude constraints
SELECT conname FROM pg_constraint
WHERE conname LIKE '%id_test%' ORDER BY 1;

           conname
------------------------
dist_partitioned_table_id_test_n_excl
 pt_1_id_test_n_excl
 pt_2_id_test_n_excl
 pt_3_id_test_n_excl
(4 rows)

-- now, verify correct propagation to all the nodes in the cluster
\c - - - :node_1_port
SELECT attrelid::regclass, attname, attidentity FROM pg_attribute
WHERE attname = 'id_test' AND attidentity = 'd' ORDER BY 1;

        attrelid        | attname | attidentity
------------------------+---------+-------------
 dist_partitioned_table | id_test | d
 pt_1                   | id_test | d
 pt_2                   | id_test | d
 pt_3                   | id_test | d
(4 rows)

SELECT relname, amname FROM pg_class c LEFT JOIN pg_am am ON (c.relam = am.oid)
WHERE relname IN ('dist_partitioned_table', 'pt_1', 'pt_2', 'pt_3') ORDER BY relname;

        relname         |  amname
------------------------+----------
 dist_partitioned_table | heap
 pt_1                   | columnar
 pt_2                   | columnar
 pt_3                   | heap
(4 rows)

-- this node is not a coordinator
-- so we can also see the exclusion constraints on the shards
SELECT conname FROM pg_constraint
WHERE conname LIKE '%id_test%' ORDER BY 1;

           conname
------------------------
dist_partitioned_table_id_test_n_excl
dist_partitioned_table_id_test_n_excl_102008
....
 pt_1_id_test_n_excl
 pt_1_id_test_n_excl_102040
...
 pt_2_id_test_n_excl
 pt_2_id_test_n_excl_102072
...
 pt_3_id_test_n_excl
 pt_3_id_test_n_excl_102104
...

-- create, distribute and populate a simple table
SET citus.shard_count TO 32;
CREATE TABLE dist_table(a int, b int);
SELECT create_distributed_table('dist_table', 'a');
INSERT INTO dist_table SELECT c, c * 10000 FROM generate_series(0, 1000) AS c;
-- explain a simple multi-shard query on the table using memory and serialize options
EXPLAIN (costs off, analyze, serialize, memory) SELECT * FROM dist_table;

                                      QUERY PLAN
--------------------------------------------------------------------------------------------
Custom Scan (Citus Adaptive) (actual time=18.490..18.519 rows=1001 loops=1)
Task Count: 32
Tuple data received from nodes: 8008 bytes
Tasks Shown: One of 32
-> Task
    Tuple data received from node: 272 bytes
     Node: host=localhost port=9702 dbname=Naisila
     ->  Seq Scan on dist_table_102141 dist_table (actual time=0.013..0.016 rows=34 loops=1)
         Planning:
           Memory: used=7kB  allocated=8kB
         Planning Time: 0.024 ms
         Serialization: time=0.000 ms  output=0kB  format=text
         Execution Time: 0.031 ms
Planning:
Memory: used=359kB allocated=512kB
Planning Time: 0.287 ms
Serialization: time=0.097 ms output=20kB format=text
Execution Time: 18.902 ms
(18 rows)

-- create, distribute and populate two simple tables
CREATE TABLE customer ( id int, name text, contact text, category text);
CREATE TABLE orders ( customer_id int, category text);
SELECT create_distributed_table('customer', 'id');
SELECT create_distributed_table('orders', 'customer_id');
INSERT INTO customer VALUES (1, 'Beana', 'beana1234@gmail.com', 'books'),
                            (2, 'Erida', 'erida1234@gmail.com', 'notebooks'),
                            (3, 'Redi', 'redi1234@gmail.com', 'pens');
INSERT INTO orders VALUES (1, 'books'), (2, 'notebooks'), (3, 'hats');

-- with Citus 13.0 in PG17 we are able to run a query
-- on the customer table that has a correlated subquery!
SELECT c.name, c.contact
FROM customer c
WHERE c.id in (SELECT customer_id FROM orders o WHERE o.category = c.category);

name  |       contact
------+---------------------
Beana | beana1234@gmail.com
Erida | erida1234@gmail.com
(2 rows)

-- pre Citus 13 or pre PG17 this query would fail with the following
ERROR:  complex joins are only supported when all distributed tables
        are co-located and joined on their distribution columns

EXPLAIN (costs off)
SELECT c.name, c.contact
FROM customer c
WHERE c.id in (SELECT customer_id FROM orders o WHERE o.category = c.category);

                              QUERY PLAN
--------------------------------------------------------------------------
Custom Scan (Citus Adaptive)
   Task Count: 32
   Tasks Shown: One of 32
   -> Task
    Node: host=localhost port=9701 dbname=citus
    -> Hash Join
        Hash Cond: ((c.category = o.category) AND (c.id = o.customer_id))
        -> Seq Scan on customer_105861 c
        -> Hash
            -> HashAggregate
        Group Key: o.category, o.customer_id
            -> Seq Scan on orders_105893 o

-- Pre-13.0 Citus:
EXPLAIN (costs off)
SELECT c.name, c.contact
FROM customer c
WHERE c.id in (SELECT customer_id FROM orders  o WHERE o.category = c.category);
DEBUG:  skipping recursive planning for the subquery since it contains
        references to outer queries
ERROR:  complex joins are only supported when all distributed tables
        are co-located and joined on their distribution columns