SET citus.next_shard_id TO 1300000;
ALTER SEQUENCE pg_catalog.pg_dist_colocationid_seq RESTART 4;
-- ===================================================================
-- create test utility function
-- ===================================================================
CREATE SEQUENCE colocation_test_seq
    MINVALUE 1000
    NO CYCLE;
/* a very simple UDF that only sets the colocation ids the same 
 * DO NOT USE THIS FUNCTION IN PRODUCTION. It manually sets colocationid column of
 * pg_dist_partition and it does not check anything about pyshical state about shards.
 */
CREATE OR REPLACE FUNCTION colocation_test_colocate_tables(source_table regclass, target_table regclass)
    RETURNS BOOL
    LANGUAGE plpgsql
    AS $colocate_tables$
DECLARE nextid INTEGER;
BEGIN
    SELECT nextval('colocation_test_seq') INTO nextid;

    UPDATE pg_dist_partition SET colocationId = nextid
    WHERE logicalrelid IN
    (
        (SELECT p1.logicalrelid
            FROM pg_dist_partition p1, pg_dist_partition p2
            WHERE
                p2.logicalrelid = source_table AND
                (p1.logicalrelid = source_table OR 
                (p1.colocationId = p2.colocationId AND p1.colocationId != 0)))
        UNION
        (SELECT target_table)
    );
    RETURN TRUE;
END;
$colocate_tables$;
-- ===================================================================
-- create test functions
-- ===================================================================
CREATE FUNCTION get_table_colocation_id(regclass)
    RETURNS INTEGER
    AS 'citus'
    LANGUAGE C STRICT;
CREATE FUNCTION tables_colocated(regclass, regclass)
    RETURNS bool
    AS 'citus'
    LANGUAGE C;
CREATE FUNCTION shards_colocated(bigint, bigint)
    RETURNS bool
    AS 'citus'
    LANGUAGE C STRICT;
CREATE FUNCTION find_shard_interval_index(bigint)
    RETURNS int
    AS 'citus'
    LANGUAGE C STRICT;
-- ===================================================================
-- test co-location util functions
-- ===================================================================
-- create distributed table observe shard pruning
CREATE TABLE table1_group1 ( id int );
SELECT master_create_distributed_table('table1_group1', 'id', 'hash');
 master_create_distributed_table 
---------------------------------
 
(1 row)

SELECT master_create_worker_shards('table1_group1', 4, 2);
 master_create_worker_shards 
-----------------------------
 
(1 row)

CREATE TABLE table2_group1 ( id int );
SELECT master_create_distributed_table('table2_group1', 'id', 'hash');
 master_create_distributed_table 
---------------------------------
 
(1 row)

SELECT master_create_worker_shards('table2_group1', 4, 2);
 master_create_worker_shards 
-----------------------------
 
(1 row)

CREATE TABLE table3_group2 ( id int );
SELECT master_create_distributed_table('table3_group2', 'id', 'hash');
 master_create_distributed_table 
---------------------------------
 
(1 row)

SELECT master_create_worker_shards('table3_group2', 4, 2);
 master_create_worker_shards 
-----------------------------
 
(1 row)

CREATE TABLE table4_group2 ( id int );
SELECT master_create_distributed_table('table4_group2', 'id', 'hash');
 master_create_distributed_table 
---------------------------------
 
(1 row)

SELECT master_create_worker_shards('table4_group2', 4, 2);
 master_create_worker_shards 
-----------------------------
 
(1 row)

CREATE TABLE table5_groupX ( id int );
SELECT master_create_distributed_table('table5_groupX', 'id', 'hash');
 master_create_distributed_table 
---------------------------------
 
(1 row)

SELECT master_create_worker_shards('table5_groupX', 4, 2);
 master_create_worker_shards 
-----------------------------
 
(1 row)

CREATE TABLE table6_append ( id int );
SELECT master_create_distributed_table('table6_append', 'id', 'append');
 master_create_distributed_table 
---------------------------------
 
(1 row)

SELECT master_create_empty_shard('table6_append');
 master_create_empty_shard 
---------------------------
                   1300020
(1 row)

SELECT master_create_empty_shard('table6_append');
 master_create_empty_shard 
---------------------------
                   1300021
(1 row)

-- make table1_group1 and table2_group1 co-located manually
SELECT colocation_test_colocate_tables('table1_group1', 'table2_group1');
 colocation_test_colocate_tables 
---------------------------------
 t
(1 row)

-- check co-location id
SELECT get_table_colocation_id('table1_group1');
 get_table_colocation_id 
-------------------------
                    1000
(1 row)

SELECT get_table_colocation_id('table5_groupX');
 get_table_colocation_id 
-------------------------
                       0
(1 row)

SELECT get_table_colocation_id('table6_append');
 get_table_colocation_id 
-------------------------
                       0
(1 row)

-- check self table co-location
SELECT tables_colocated('table1_group1', 'table1_group1');
 tables_colocated 
------------------
 t
(1 row)

SELECT tables_colocated('table5_groupX', 'table5_groupX');
 tables_colocated 
------------------
 t
(1 row)

SELECT tables_colocated('table6_append', 'table6_append');
 tables_colocated 
------------------
 t
(1 row)

-- check table co-location with same co-location group
SELECT tables_colocated('table1_group1', 'table2_group1');
 tables_colocated 
------------------
 t
(1 row)

-- check table co-location with different co-location group
SELECT tables_colocated('table1_group1', 'table3_group2');
 tables_colocated 
------------------
 f
(1 row)

-- check table co-location with invalid co-location group
SELECT tables_colocated('table1_group1', 'table5_groupX');
 tables_colocated 
------------------
 f
(1 row)

SELECT tables_colocated('table1_group1', 'table6_append');
 tables_colocated 
------------------
 f
(1 row)

-- check self shard co-location
SELECT shards_colocated(1300000, 1300000);
 shards_colocated 
------------------
 t
(1 row)

SELECT shards_colocated(1300016, 1300016);
 shards_colocated 
------------------
 t
(1 row)

SELECT shards_colocated(1300020, 1300020);
 shards_colocated 
------------------
 t
(1 row)

-- check shard co-location with same co-location group
SELECT shards_colocated(1300000, 1300004);
 shards_colocated 
------------------
 t
(1 row)

-- check shard co-location with same table different co-location group
SELECT shards_colocated(1300000, 1300001);
 shards_colocated 
------------------
 f
(1 row)

-- check shard co-location with different co-location group
SELECT shards_colocated(1300000, 1300005);
 shards_colocated 
------------------
 f
(1 row)

-- check shard co-location with invalid co-location group
SELECT shards_colocated(1300000, 1300016);
 shards_colocated 
------------------
 f
(1 row)

SELECT shards_colocated(1300000, 1300020);
 shards_colocated 
------------------
 f
(1 row)

-- check co-located table list
SELECT UNNEST(get_colocated_table_array('table1_group1'))::regclass ORDER BY 1;
    unnest     
---------------
 table1_group1
 table2_group1
(2 rows)

SELECT UNNEST(get_colocated_table_array('table5_groupX'))::regclass ORDER BY 1;
    unnest     
---------------
 table5_groupx
(1 row)

SELECT UNNEST(get_colocated_table_array('table6_append'))::regclass ORDER BY 1;
    unnest     
---------------
 table6_append
(1 row)

-- check co-located shard list
SELECT UNNEST(get_colocated_shard_array(1300000))::regclass ORDER BY 1;
 unnest  
---------
 1300000
 1300004
(2 rows)

SELECT UNNEST(get_colocated_shard_array(1300016))::regclass ORDER BY 1;
 unnest  
---------
 1300016
(1 row)

SELECT UNNEST(get_colocated_shard_array(1300020))::regclass ORDER BY 1;
 unnest  
---------
 1300020
(1 row)

-- check FindShardIntervalIndex function
SELECT find_shard_interval_index(1300000);
 find_shard_interval_index 
---------------------------
                         0
(1 row)

SELECT find_shard_interval_index(1300001);
 find_shard_interval_index 
---------------------------
                         1
(1 row)

SELECT find_shard_interval_index(1300002);
 find_shard_interval_index 
---------------------------
                         2
(1 row)

SELECT find_shard_interval_index(1300003);
 find_shard_interval_index 
---------------------------
                         3
(1 row)

SELECT find_shard_interval_index(1300016);
 find_shard_interval_index 
---------------------------
                         0
(1 row)

-- check external colocation API
SET citus.shard_count = 2;
CREATE TABLE table1_groupA ( id int );
SELECT create_distributed_table('table1_groupA', 'id');
 create_distributed_table 
--------------------------
 
(1 row)

CREATE TABLE table2_groupA ( id int );
SELECT create_distributed_table('table2_groupA', 'id');
 create_distributed_table 
--------------------------
 
(1 row)

-- change shard replication factor
SET citus.shard_replication_factor = 1;
CREATE TABLE table1_groupB ( id int );
SELECT create_distributed_table('table1_groupB', 'id');
 create_distributed_table 
--------------------------
 
(1 row)

CREATE TABLE table2_groupB ( id int );
SELECT create_distributed_table('table2_groupB', 'id');
 create_distributed_table 
--------------------------
 
(1 row)

UPDATE pg_dist_partition SET repmodel='c' WHERE logicalrelid='table1_groupB'::regclass;
UPDATE pg_dist_partition SET repmodel='c' WHERE logicalrelid='table2_groupB'::regclass;
-- revert back to default shard replication factor
SET citus.shard_replication_factor to DEFAULT;
-- change partition column type
CREATE TABLE table1_groupC ( id text );
SELECT create_distributed_table('table1_groupC', 'id');
 create_distributed_table 
--------------------------
 
(1 row)

CREATE TABLE table2_groupC ( id text );
SELECT create_distributed_table('table2_groupC', 'id');
 create_distributed_table 
--------------------------
 
(1 row)

-- change shard count
SET citus.shard_count = 8;
CREATE TABLE table1_groupD ( id int );
SELECT create_distributed_table('table1_groupD', 'id');
 create_distributed_table 
--------------------------
 
(1 row)

CREATE TABLE table2_groupD ( id int );
SELECT create_distributed_table('table2_groupD', 'id');
 create_distributed_table 
--------------------------
 
(1 row)

-- try other distribution methods
CREATE TABLE table_append ( id int );
SELECT create_distributed_table('table_append', 'id', 'append');
 create_distributed_table 
--------------------------
 
(1 row)

CREATE TABLE table_range ( id int );
SELECT create_distributed_table('table_range', 'id', 'range');
 create_distributed_table 
--------------------------
 
(1 row)

-- test foreign table creation
CREATE FOREIGN TABLE table3_groupD ( id int ) SERVER fake_fdw_server;
SELECT create_distributed_table('table3_groupD', 'id');
NOTICE:  foreign-data wrapper "fake_fdw" does not have an extension defined
 create_distributed_table 
--------------------------
 
(1 row)

-- check metadata
SELECT * FROM pg_dist_colocation 
    WHERE colocationid >= 1 AND colocationid < 1000 
    ORDER BY colocationid;
 colocationid | shardcount | replicationfactor | distributioncolumntype 
--------------+------------+-------------------+------------------------
            3 |          4 |                 2 |                     23
            4 |          2 |                 2 |                     23
            5 |          2 |                 1 |                     23
            6 |          2 |                 2 |                     25
            7 |          8 |                 2 |                     23
(5 rows)

SELECT logicalrelid, colocationid FROM pg_dist_partition
    WHERE colocationid >= 1 AND colocationid < 1000 
    ORDER BY logicalrelid;
 logicalrelid  | colocationid 
---------------+--------------
 table1_groupa |            4
 table2_groupa |            4
 table1_groupb |            5
 table2_groupb |            5
 table1_groupc |            6
 table2_groupc |            6
 table1_groupd |            7
 table2_groupd |            7
 table3_groupd |            7
(9 rows)

-- check effects of dropping tables
DROP TABLE table1_groupA;
SELECT * FROM pg_dist_colocation WHERE colocationid = 4;
 colocationid | shardcount | replicationfactor | distributioncolumntype 
--------------+------------+-------------------+------------------------
            4 |          2 |                 2 |                     23
(1 row)

-- dropping all tables in a colocation group also deletes the colocation group
DROP TABLE table2_groupA;
SELECT * FROM pg_dist_colocation WHERE colocationid = 4;
 colocationid | shardcount | replicationfactor | distributioncolumntype 
--------------+------------+-------------------+------------------------
            4 |          2 |                 2 |                     23
(1 row)

-- create dropped colocation group again
SET citus.shard_count = 2;
CREATE TABLE table1_groupE ( id int );
SELECT create_distributed_table('table1_groupE', 'id');
 create_distributed_table 
--------------------------
 
(1 row)

CREATE TABLE table2_groupE ( id int );
SELECT create_distributed_table('table2_groupE', 'id');
 create_distributed_table 
--------------------------
 
(1 row)

-- test different table DDL
CREATE TABLE table3_groupE ( dummy_column text, id int );
SELECT create_distributed_table('table3_groupE', 'id');
 create_distributed_table 
--------------------------
 
(1 row)

-- test different schema
CREATE SCHEMA schema_collocation;
CREATE TABLE schema_collocation.table4_groupE ( id int );
SELECT create_distributed_table('schema_collocation.table4_groupE', 'id');
 create_distributed_table 
--------------------------
 
(1 row)

-- test colocate_with option
CREATE TABLE table1_group_none_1 ( id int );
SELECT create_distributed_table('table1_group_none_1', 'id', colocate_with => 'none');
 create_distributed_table 
--------------------------
 
(1 row)

CREATE TABLE table2_group_none_1 ( id int );
SELECT create_distributed_table('table2_group_none_1', 'id', colocate_with => 'table1_group_none_1');
 create_distributed_table 
--------------------------
 
(1 row)

CREATE TABLE table1_group_none_2 ( id int );
SELECT create_distributed_table('table1_group_none_2', 'id', colocate_with => 'none');
 create_distributed_table 
--------------------------
 
(1 row)

CREATE TABLE table4_groupE ( id int );
SELECT create_distributed_table('table4_groupE', 'id', colocate_with => 'default');
 create_distributed_table 
--------------------------
 
(1 row)

SET citus.shard_count = 3;
-- check that this new configuration does not have a default group
CREATE TABLE table1_group_none_3 ( id int );
SELECT create_distributed_table('table1_group_none_3', 'id', colocate_with => 'NONE');
 create_distributed_table 
--------------------------
 
(1 row)

-- a new table does not use a non-default group
CREATE TABLE table1_group_default ( id int );
SELECT create_distributed_table('table1_group_default', 'id', colocate_with => 'DEFAULT');
 create_distributed_table 
--------------------------
 
(1 row)

-- check metadata
SELECT * FROM pg_dist_colocation
    WHERE colocationid >= 1 AND colocationid < 1000
    ORDER BY colocationid;
 colocationid | shardcount | replicationfactor | distributioncolumntype 
--------------+------------+-------------------+------------------------
            3 |          4 |                 2 |                     23
            4 |          2 |                 2 |                     23
            5 |          2 |                 1 |                     23
            6 |          2 |                 2 |                     25
            7 |          8 |                 2 |                     23
           11 |          3 |                 2 |                     23
(6 rows)

SELECT logicalrelid, colocationid FROM pg_dist_partition
    WHERE colocationid >= 1 AND colocationid < 1000
    ORDER BY colocationid, logicalrelid;
           logicalrelid           | colocationid 
----------------------------------+--------------
 table1_groupe                    |            4
 table2_groupe                    |            4
 table3_groupe                    |            4
 schema_collocation.table4_groupe |            4
 table4_groupe                    |            4
 table1_groupb                    |            5
 table2_groupb                    |            5
 table1_groupc                    |            6
 table2_groupc                    |            6
 table1_groupd                    |            7
 table2_groupd                    |            7
 table3_groupd                    |            7
 table1_group_none_1              |            8
 table2_group_none_1              |            8
 table1_group_none_2              |            9
 table1_group_none_3              |           10
 table1_group_default             |           11
(17 rows)

-- check failing colocate_with options
CREATE TABLE table_postgresql( id int );
CREATE TABLE table_failing ( id int );
SELECT create_distributed_table('table_failing', 'id', colocate_with => 'table_append');
ERROR:  cannot distribute relation
DETAIL:  Currently, colocate_with option is only supported for hash distributed tables.
SELECT create_distributed_table('table_failing', 'id', 'append', 'table1_groupE');
ERROR:  cannot distribute relation
DETAIL:  Currently, colocate_with option is only supported for hash distributed tables.
SELECT create_distributed_table('table_failing', 'id', colocate_with => 'table_postgresql');
ERROR:  relation table_postgresql is not distributed
SELECT create_distributed_table('table_failing', 'id', colocate_with => 'no_table');
ERROR:  relation "no_table" does not exist
SELECT create_distributed_table('table_failing', 'id', colocate_with => '');
ERROR:  invalid name syntax
SELECT create_distributed_table('table_failing', 'id', colocate_with => NULL);
 create_distributed_table 
--------------------------
 
(1 row)

-- check with different distribution column types
CREATE TABLE table_bigint ( id bigint );
SELECT create_distributed_table('table_bigint', 'id', colocate_with => 'table1_groupE');
ERROR:  cannot colocate tables table1_groupe and table_bigint
DETAIL:  Distribution column types don't match for table1_groupe and table_bigint.
-- check worker table schemas
\c - - - :worker_1_port
SELECT "Column", "Type", "Modifiers" FROM table_desc WHERE relid='public.table3_groupE_1300062'::regclass;
    Column    |  Type   | Modifiers 
--------------+---------+-----------
 dummy_column | text    | 
 id           | integer | 
(2 rows)

SELECT "Column", "Type", "Modifiers" FROM table_desc WHERE relid='schema_collocation.table4_groupE_1300064'::regclass;
 Column |  Type   | Modifiers 
--------+---------+-----------
 id     | integer | 
(1 row)

\c - - - :master_port
SET citus.next_shard_id TO 1300080;
CREATE TABLE table1_groupF ( id int );
SELECT create_reference_table('table1_groupF');
 create_reference_table 
------------------------
 
(1 row)

CREATE TABLE table2_groupF ( id int );
SELECT create_reference_table('table2_groupF');
 create_reference_table 
------------------------
 
(1 row)

-- check metadata
SELECT * FROM pg_dist_colocation
    WHERE colocationid >= 1 AND colocationid < 1000
    ORDER BY colocationid;
 colocationid | shardcount | replicationfactor | distributioncolumntype 
--------------+------------+-------------------+------------------------
            3 |          4 |                 2 |                     23
            4 |          2 |                 2 |                     23
            5 |          2 |                 1 |                     23
            6 |          2 |                 2 |                     25
            7 |          8 |                 2 |                     23
           11 |          3 |                 2 |                     23
(6 rows)

-- cross check with internal colocation API
SELECT 
    p1.logicalrelid::regclass AS table1,
    p2.logicalrelid::regclass AS table2,
    tables_colocated(p1.logicalrelid , p2.logicalrelid) AS colocated
FROM
    pg_dist_partition p1,
    pg_dist_partition p2
WHERE
    p1.logicalrelid < p2.logicalrelid AND
    p1.colocationid != 0 AND
    p2.colocationid != 0 AND
    tables_colocated(p1.logicalrelid , p2.logicalrelid) is TRUE
ORDER BY
    table1,
    table2;
              table1              |              table2              | colocated 
----------------------------------+----------------------------------+-----------
 table1_group1                    | table2_group1                    | t
 table1_groupb                    | table2_groupb                    | t
 table1_groupc                    | table2_groupc                    | t
 table1_groupd                    | table2_groupd                    | t
 table1_groupd                    | table3_groupd                    | t
 table2_groupd                    | table3_groupd                    | t
 table1_groupe                    | table2_groupe                    | t
 table1_groupe                    | table3_groupe                    | t
 table1_groupe                    | schema_collocation.table4_groupe | t
 table1_groupe                    | table4_groupe                    | t
 table2_groupe                    | table3_groupe                    | t
 table2_groupe                    | schema_collocation.table4_groupe | t
 table2_groupe                    | table4_groupe                    | t
 table3_groupe                    | schema_collocation.table4_groupe | t
 table3_groupe                    | table4_groupe                    | t
 schema_collocation.table4_groupe | table4_groupe                    | t
 table1_group_none_1              | table2_group_none_1              | t
 table1_groupf                    | table2_groupf                    | t
(18 rows)

-- check created shards
SELECT
    logicalrelid,
    pg_dist_shard.shardid AS shardid,
    shardstorage,
    nodeport,
    shardminvalue,
    shardmaxvalue
FROM
    pg_dist_shard,
    pg_dist_shard_placement
WHERE
    pg_dist_shard.shardid = pg_dist_shard_placement.shardid AND
    pg_dist_shard.shardid >= 1300026
ORDER BY
    logicalrelid,
    shardmaxvalue::integer,
    shardid,
    nodeport;
           logicalrelid           | shardid | shardstorage | nodeport | shardminvalue | shardmaxvalue 
----------------------------------+---------+--------------+----------+---------------+---------------
 table1_groupb                    | 1300026 | t            |    57637 | -2147483648   | -1
 table1_groupb                    | 1300027 | t            |    57638 | 0             | 2147483647
 table2_groupb                    | 1300028 | t            |    57637 | -2147483648   | -1
 table2_groupb                    | 1300029 | t            |    57638 | 0             | 2147483647
 table1_groupc                    | 1300030 | t            |    57637 | -2147483648   | -1
 table1_groupc                    | 1300030 | t            |    57638 | -2147483648   | -1
 table1_groupc                    | 1300031 | t            |    57637 | 0             | 2147483647
 table1_groupc                    | 1300031 | t            |    57638 | 0             | 2147483647
 table2_groupc                    | 1300032 | t            |    57637 | -2147483648   | -1
 table2_groupc                    | 1300032 | t            |    57638 | -2147483648   | -1
 table2_groupc                    | 1300033 | t            |    57637 | 0             | 2147483647
 table2_groupc                    | 1300033 | t            |    57638 | 0             | 2147483647
 table1_groupd                    | 1300034 | t            |    57637 | -2147483648   | -1610612737
 table1_groupd                    | 1300034 | t            |    57638 | -2147483648   | -1610612737
 table1_groupd                    | 1300035 | t            |    57637 | -1610612736   | -1073741825
 table1_groupd                    | 1300035 | t            |    57638 | -1610612736   | -1073741825
 table1_groupd                    | 1300036 | t            |    57637 | -1073741824   | -536870913
 table1_groupd                    | 1300036 | t            |    57638 | -1073741824   | -536870913
 table1_groupd                    | 1300037 | t            |    57637 | -536870912    | -1
 table1_groupd                    | 1300037 | t            |    57638 | -536870912    | -1
 table1_groupd                    | 1300038 | t            |    57637 | 0             | 536870911
 table1_groupd                    | 1300038 | t            |    57638 | 0             | 536870911
 table1_groupd                    | 1300039 | t            |    57637 | 536870912     | 1073741823
 table1_groupd                    | 1300039 | t            |    57638 | 536870912     | 1073741823
 table1_groupd                    | 1300040 | t            |    57637 | 1073741824    | 1610612735
 table1_groupd                    | 1300040 | t            |    57638 | 1073741824    | 1610612735
 table1_groupd                    | 1300041 | t            |    57637 | 1610612736    | 2147483647
 table1_groupd                    | 1300041 | t            |    57638 | 1610612736    | 2147483647
 table2_groupd                    | 1300042 | t            |    57637 | -2147483648   | -1610612737
 table2_groupd                    | 1300042 | t            |    57638 | -2147483648   | -1610612737
 table2_groupd                    | 1300043 | t            |    57637 | -1610612736   | -1073741825
 table2_groupd                    | 1300043 | t            |    57638 | -1610612736   | -1073741825
 table2_groupd                    | 1300044 | t            |    57637 | -1073741824   | -536870913
 table2_groupd                    | 1300044 | t            |    57638 | -1073741824   | -536870913
 table2_groupd                    | 1300045 | t            |    57637 | -536870912    | -1
 table2_groupd                    | 1300045 | t            |    57638 | -536870912    | -1
 table2_groupd                    | 1300046 | t            |    57637 | 0             | 536870911
 table2_groupd                    | 1300046 | t            |    57638 | 0             | 536870911
 table2_groupd                    | 1300047 | t            |    57637 | 536870912     | 1073741823
 table2_groupd                    | 1300047 | t            |    57638 | 536870912     | 1073741823
 table2_groupd                    | 1300048 | t            |    57637 | 1073741824    | 1610612735
 table2_groupd                    | 1300048 | t            |    57638 | 1073741824    | 1610612735
 table2_groupd                    | 1300049 | t            |    57637 | 1610612736    | 2147483647
 table2_groupd                    | 1300049 | t            |    57638 | 1610612736    | 2147483647
 table3_groupd                    | 1300050 | f            |    57637 | -2147483648   | -1610612737
 table3_groupd                    | 1300050 | f            |    57638 | -2147483648   | -1610612737
 table3_groupd                    | 1300051 | f            |    57637 | -1610612736   | -1073741825
 table3_groupd                    | 1300051 | f            |    57638 | -1610612736   | -1073741825
 table3_groupd                    | 1300052 | f            |    57637 | -1073741824   | -536870913
 table3_groupd                    | 1300052 | f            |    57638 | -1073741824   | -536870913
 table3_groupd                    | 1300053 | f            |    57637 | -536870912    | -1
 table3_groupd                    | 1300053 | f            |    57638 | -536870912    | -1
 table3_groupd                    | 1300054 | f            |    57637 | 0             | 536870911
 table3_groupd                    | 1300054 | f            |    57638 | 0             | 536870911
 table3_groupd                    | 1300055 | f            |    57637 | 536870912     | 1073741823
 table3_groupd                    | 1300055 | f            |    57638 | 536870912     | 1073741823
 table3_groupd                    | 1300056 | f            |    57637 | 1073741824    | 1610612735
 table3_groupd                    | 1300056 | f            |    57638 | 1073741824    | 1610612735
 table3_groupd                    | 1300057 | f            |    57637 | 1610612736    | 2147483647
 table3_groupd                    | 1300057 | f            |    57638 | 1610612736    | 2147483647
 table1_groupe                    | 1300058 | t            |    57637 | -2147483648   | -1
 table1_groupe                    | 1300058 | t            |    57638 | -2147483648   | -1
 table1_groupe                    | 1300059 | t            |    57637 | 0             | 2147483647
 table1_groupe                    | 1300059 | t            |    57638 | 0             | 2147483647
 table2_groupe                    | 1300060 | t            |    57637 | -2147483648   | -1
 table2_groupe                    | 1300060 | t            |    57638 | -2147483648   | -1
 table2_groupe                    | 1300061 | t            |    57637 | 0             | 2147483647
 table2_groupe                    | 1300061 | t            |    57638 | 0             | 2147483647
 table3_groupe                    | 1300062 | t            |    57637 | -2147483648   | -1
 table3_groupe                    | 1300062 | t            |    57638 | -2147483648   | -1
 table3_groupe                    | 1300063 | t            |    57637 | 0             | 2147483647
 table3_groupe                    | 1300063 | t            |    57638 | 0             | 2147483647
 schema_collocation.table4_groupe | 1300064 | t            |    57637 | -2147483648   | -1
 schema_collocation.table4_groupe | 1300064 | t            |    57638 | -2147483648   | -1
 schema_collocation.table4_groupe | 1300065 | t            |    57637 | 0             | 2147483647
 schema_collocation.table4_groupe | 1300065 | t            |    57638 | 0             | 2147483647
 table1_group_none_1              | 1300066 | t            |    57637 | -2147483648   | -1
 table1_group_none_1              | 1300066 | t            |    57638 | -2147483648   | -1
 table1_group_none_1              | 1300067 | t            |    57637 | 0             | 2147483647
 table1_group_none_1              | 1300067 | t            |    57638 | 0             | 2147483647
 table2_group_none_1              | 1300068 | t            |    57637 | -2147483648   | -1
 table2_group_none_1              | 1300068 | t            |    57638 | -2147483648   | -1
 table2_group_none_1              | 1300069 | t            |    57637 | 0             | 2147483647
 table2_group_none_1              | 1300069 | t            |    57638 | 0             | 2147483647
 table1_group_none_2              | 1300070 | t            |    57637 | -2147483648   | -1
 table1_group_none_2              | 1300070 | t            |    57638 | -2147483648   | -1
 table1_group_none_2              | 1300071 | t            |    57637 | 0             | 2147483647
 table1_group_none_2              | 1300071 | t            |    57638 | 0             | 2147483647
 table4_groupe                    | 1300072 | t            |    57637 | -2147483648   | -1
 table4_groupe                    | 1300072 | t            |    57638 | -2147483648   | -1
 table4_groupe                    | 1300073 | t            |    57637 | 0             | 2147483647
 table4_groupe                    | 1300073 | t            |    57638 | 0             | 2147483647
 table1_group_none_3              | 1300074 | t            |    57637 | -2147483648   | -715827884
 table1_group_none_3              | 1300074 | t            |    57638 | -2147483648   | -715827884
 table1_group_none_3              | 1300075 | t            |    57637 | -715827883    | 715827881
 table1_group_none_3              | 1300075 | t            |    57638 | -715827883    | 715827881
 table1_group_none_3              | 1300076 | t            |    57637 | 715827882     | 2147483647
 table1_group_none_3              | 1300076 | t            |    57638 | 715827882     | 2147483647
 table1_group_default             | 1300077 | t            |    57637 | -2147483648   | -715827884
 table1_group_default             | 1300077 | t            |    57638 | -2147483648   | -715827884
 table1_group_default             | 1300078 | t            |    57637 | -715827883    | 715827881
 table1_group_default             | 1300078 | t            |    57638 | -715827883    | 715827881
 table1_group_default             | 1300079 | t            |    57637 | 715827882     | 2147483647
 table1_group_default             | 1300079 | t            |    57638 | 715827882     | 2147483647
 table1_groupf                    | 1300080 | t            |    57637 |               | 
 table1_groupf                    | 1300080 | t            |    57638 |               | 
 table2_groupf                    | 1300081 | t            |    57637 |               | 
 table2_groupf                    | 1300081 | t            |    57638 |               | 
(108 rows)

-- reset colocation ids to test mark_tables_colocated
ALTER SEQUENCE pg_catalog.pg_dist_colocationid_seq RESTART 1;
DELETE FROM pg_dist_colocation
    WHERE colocationid >= 1 AND colocationid < 1000;
UPDATE pg_dist_partition SET colocationid = 0
    WHERE colocationid >= 1 AND colocationid < 1000;
-- check metadata
SELECT * FROM pg_dist_colocation
    WHERE colocationid >= 1 AND colocationid < 1000
    ORDER BY colocationid;
 colocationid | shardcount | replicationfactor | distributioncolumntype 
--------------+------------+-------------------+------------------------
(0 rows)

SELECT logicalrelid, colocationid FROM pg_dist_partition
    WHERE colocationid >= 1 AND colocationid < 1000
    ORDER BY colocationid, logicalrelid;
 logicalrelid | colocationid 
--------------+--------------
(0 rows)

-- first check failing cases
SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupC']);
ERROR:  cannot colocate tables table1_groupb and table1_groupc
DETAIL:  Distribution column types don't match for table1_groupb and table1_groupc.
SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupD']);
ERROR:  cannot colocate tables table1_groupb and table1_groupd
DETAIL:  Shard counts don't match for table1_groupb and table1_groupd.
SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupE']);
ERROR:  cannot colocate tables table1_groupb and table1_groupe
DETAIL:  Shard 1300026 of table1_groupb and shard 1300058 of table1_groupe have different number of shard placements.
SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupF']);
ERROR:  cannot colocate tables table1_groupb and table1_groupf
DETAIL:  Replication models don't match for table1_groupb and table1_groupf.
SELECT mark_tables_colocated('table1_groupB', ARRAY['table2_groupB', 'table1_groupD']);
ERROR:  cannot colocate tables table1_groupb and table1_groupd
DETAIL:  Shard counts don't match for table1_groupb and table1_groupd.
-- check metadata to see failing calls didn't have any side effects
SELECT * FROM pg_dist_colocation
    WHERE colocationid >= 1 AND colocationid < 1000
    ORDER BY colocationid;
 colocationid | shardcount | replicationfactor | distributioncolumntype 
--------------+------------+-------------------+------------------------
(0 rows)

SELECT logicalrelid, colocationid FROM pg_dist_partition
    WHERE colocationid >= 1 AND colocationid < 1000
    ORDER BY colocationid, logicalrelid;
 logicalrelid | colocationid 
--------------+--------------
(0 rows)

-- check successfully cololated tables
SELECT mark_tables_colocated('table1_groupB', ARRAY['table2_groupB']);
 mark_tables_colocated 
-----------------------
 
(1 row)

SELECT mark_tables_colocated('table1_groupC', ARRAY['table2_groupC']);
 mark_tables_colocated 
-----------------------
 
(1 row)

SELECT mark_tables_colocated('table1_groupD', ARRAY['table2_groupD']);
 mark_tables_colocated 
-----------------------
 
(1 row)

SELECT mark_tables_colocated('table1_groupE', ARRAY['table2_groupE', 'table3_groupE']);
 mark_tables_colocated 
-----------------------
 
(1 row)

SELECT mark_tables_colocated('table1_groupF', ARRAY['table2_groupF']);
 mark_tables_colocated 
-----------------------
 
(1 row)

-- check to colocate with itself
SELECT mark_tables_colocated('table1_groupB', ARRAY['table1_groupB']);
 mark_tables_colocated 
-----------------------
 
(1 row)

SET citus.shard_count = 2;
CREATE TABLE table1_group_none ( id int );
SELECT create_distributed_table('table1_group_none', 'id', colocate_with => 'NONE');
 create_distributed_table 
--------------------------
 
(1 row)

CREATE TABLE table2_group_none ( id int );
SELECT create_distributed_table('table2_group_none', 'id', colocate_with => 'NONE');
 create_distributed_table 
--------------------------
 
(1 row)

-- check metadata to see colocation groups are created successfully
SELECT * FROM pg_dist_colocation 
    WHERE colocationid >= 1 AND colocationid < 1000 
    ORDER BY colocationid;
 colocationid | shardcount | replicationfactor | distributioncolumntype 
--------------+------------+-------------------+------------------------
            2 |          2 |                 1 |                     23
            3 |          2 |                 2 |                     25
            4 |          8 |                 2 |                     23
            5 |          2 |                 2 |                     23
(4 rows)

SELECT logicalrelid, colocationid FROM pg_dist_partition
    WHERE colocationid >= 1 AND colocationid < 1000
    ORDER BY colocationid, logicalrelid;
   logicalrelid    | colocationid 
-------------------+--------------
 table1_groupb     |            2
 table2_groupb     |            2
 table1_groupc     |            3
 table2_groupc     |            3
 table1_groupd     |            4
 table2_groupd     |            4
 table1_groupe     |            5
 table2_groupe     |            5
 table3_groupe     |            5
 table1_group_none |            6
 table2_group_none |            7
(11 rows)

-- move the all tables in colocation group 5 to colocation group 7
SELECT mark_tables_colocated('table1_group_none', ARRAY['table1_groupE', 'table2_groupE', 'table3_groupE']);
 mark_tables_colocated 
-----------------------
 
(1 row)

-- move a table with a colocation id which is already not in pg_dist_colocation
SELECT mark_tables_colocated('table1_group_none', ARRAY['table2_group_none']);
 mark_tables_colocated 
-----------------------
 
(1 row)

-- check metadata to see that unused colocation group is deleted
SELECT * FROM pg_dist_colocation
    WHERE colocationid >= 1 AND colocationid < 1000
    ORDER BY colocationid;
 colocationid | shardcount | replicationfactor | distributioncolumntype 
--------------+------------+-------------------+------------------------
            2 |          2 |                 1 |                     23
            3 |          2 |                 2 |                     25
            4 |          8 |                 2 |                     23
(3 rows)

SELECT logicalrelid, colocationid FROM pg_dist_partition
    WHERE colocationid >= 1 AND colocationid < 1000
    ORDER BY colocationid, logicalrelid;
   logicalrelid    | colocationid 
-------------------+--------------
 table1_groupb     |            2
 table2_groupb     |            2
 table1_groupc     |            3
 table2_groupc     |            3
 table1_groupd     |            4
 table2_groupd     |            4
 table1_groupe     |            6
 table2_groupe     |            6
 table3_groupe     |            6
 table1_group_none |            6
 table2_group_none |            6
(11 rows)

-- try to colocate different replication models
CREATE TABLE table1_groupG ( id int );
SELECT create_distributed_table('table1_groupG', 'id');
 create_distributed_table 
--------------------------
 
(1 row)

-- update replication model
UPDATE pg_dist_partition SET repmodel = 's' WHERE logicalrelid = 'table1_groupG'::regclass;
CREATE TABLE table2_groupG ( id int );
SELECT create_distributed_table('table2_groupG', 'id', colocate_with => 'table1_groupG');
ERROR:  cannot colocate tables table1_groupg and table2_groupg
DETAIL:  Replication models don't match for table1_groupg and table2_groupg.
CREATE TABLE table2_groupG ( id int );
ERROR:  relation "table2_groupg" already exists
SELECT create_distributed_table('table2_groupG', 'id', colocate_with => 'NONE');
 create_distributed_table 
--------------------------
 
(1 row)

SELECT mark_tables_colocated('table1_groupG', ARRAY['table2_groupG']);
ERROR:  cannot colocate tables table1_groupg and table2_groupg
DETAIL:  Replication models don't match for table1_groupg and table2_groupg.
-- drop tables to clean test space
DROP TABLE table1_groupb;
DROP TABLE table2_groupb;
DROP TABLE table1_groupc;
DROP TABLE table2_groupc;
DROP TABLE table1_groupd;
DROP TABLE table2_groupd;
DROP TABLE table1_groupf;
DROP TABLE table2_groupf;
DROP TABLE table1_groupe;
DROP TABLE table2_groupe;
DROP TABLE table3_groupe;
DROP TABLE table4_groupe;
DROP TABLE schema_collocation.table4_groupe;
DROP TABLE table1_group_none_1;
DROP TABLE table2_group_none_1;
DROP TABLE table1_group_none_2;
DROP TABLE table1_group_none_3;
DROP TABLE table1_group_none;
DROP TABLE table2_group_none;
DROP TABLE table1_group_default;