Merge pull request #5310 from citusdata/teja_colocate_partitions

Parition shards to be colocated with parent shards

src/backend/distributed/commands/create_distributed_table.c

@@ -556,11 +556,16 @@ CreateDistributedTable(Oid relationId, Var *distributionColumn, char distributio
 	{
 		List *partitionList = PartitionList(relationId);
 		Oid partitionRelationId = InvalidOid;
+		Oid namespaceId = get_rel_namespace(relationId);
+		char *schemaName = get_namespace_name(namespaceId);
+		char *relationName = get_rel_name(relationId);
+		char *parentRelationName = quote_qualified_identifier(schemaName, relationName);
+
 		foreach_oid(partitionRelationId, partitionList)
 		{
 			CreateDistributedTable(partitionRelationId, distributionColumn,
 								   distributionMethod, shardCount, false,
-								   colocateWithTableName, viaDeprecatedAPI);
+								   parentRelationName, viaDeprecatedAPI);
 		}
 	}
 

src/test/regress/expected/multi_partitioning.out

@@ -1028,6 +1028,46 @@ CREATE TABLE partitioned_users_table_2009 PARTITION OF partitioned_users_table F
 CREATE TABLE partitioned_events_table_2009 PARTITION OF partitioned_events_table FOR VALUES FROM ('2017-01-01') TO ('2018-01-01');
 INSERT INTO partitioned_events_table SELECT * FROM events_table;
 INSERT INTO partitioned_users_table_2009 SELECT * FROM users_table;
+-- test distributed partitions are indeed colocated with the parent table
+CREATE TABLE sensors(measureid integer, eventdatetime date, measure_data jsonb, PRIMARY KEY (measureid, eventdatetime, measure_data))
+PARTITION BY RANGE(eventdatetime);
+CREATE TABLE sensors_old PARTITION OF sensors FOR VALUES FROM ('2000-01-01') TO ('2020-01-01');
+CREATE TABLE sensors_2020_01_01 PARTITION OF sensors FOR VALUES FROM ('2020-01-01') TO ('2020-02-01');
+CREATE TABLE sensors_new PARTITION OF sensors DEFAULT;
+SELECT create_distributed_table('sensors', 'measureid', colocate_with:='none');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT count(DISTINCT colocationid) FROM pg_dist_partition
+WHERE logicalrelid IN ('sensors'::regclass, 'sensors_old'::regclass, 'sensors_2020_01_01'::regclass, 'sensors_new'::regclass);
+ count
+---------------------------------------------------------------------
+     1
+(1 row)
+
+CREATE TABLE local_sensors(measureid integer, eventdatetime date, measure_data jsonb, PRIMARY KEY (measureid, eventdatetime, measure_data))
+PARTITION BY RANGE(eventdatetime);
+CREATE TABLE local_sensors_old PARTITION OF local_sensors FOR VALUES FROM ('2000-01-01') TO ('2020-01-01');
+CREATE TABLE local_sensors_2020_01_01 PARTITION OF local_sensors FOR VALUES FROM ('2020-01-01') TO ('2020-02-01');
+CREATE TABLE local_sensors_new PARTITION OF local_sensors DEFAULT;
+SELECT create_distributed_table('local_sensors', 'measureid', colocate_with:='sensors');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+SELECT count(DISTINCT colocationid) FROM pg_dist_partition
+WHERE logicalrelid IN ('sensors'::regclass, 'sensors_old'::regclass, 'sensors_2020_01_01'::regclass, 'sensors_new'::regclass,
+'local_sensors'::regclass, 'local_sensors_old'::regclass, 'local_sensors_2020_01_01'::regclass, 'local_sensors_new'::regclass);
+ count
+---------------------------------------------------------------------
+     1
+(1 row)
+
+DROP TABLE sensors;
+DROP TABLE local_sensors;
 --
 -- Complex JOINs, subqueries, UNIONs etc...
 --
@@ -1298,7 +1338,7 @@ INSERT INTO multi_column_partitioning VALUES(1, 1);
 INSERT INTO multi_column_partitioning_0_0_10_0 VALUES(5, -5);
 -- test INSERT to multi-column partitioned table where no suitable partition exists
 INSERT INTO multi_column_partitioning VALUES(10, 1);
-ERROR:  no partition of relation "multi_column_partitioning_1660101" found for row
+ERROR:  no partition of relation "multi_column_partitioning_1660133" found for row
 DETAIL:  Partition key of the failing row contains (c1, c2) = (10, 1).
 CONTEXT:  while executing command on localhost:xxxxx
 -- test with MINVALUE/MAXVALUE
@@ -1308,7 +1348,7 @@ INSERT INTO multi_column_partitioning VALUES(11, -11);
 INSERT INTO multi_column_partitioning_10_max_20_min VALUES(19, -19);
 -- test INSERT to multi-column partitioned table where no suitable partition exists
 INSERT INTO multi_column_partitioning VALUES(20, -20);
-ERROR:  no partition of relation "multi_column_partitioning_1660101" found for row
+ERROR:  no partition of relation "multi_column_partitioning_1660133" found for row
 DETAIL:  Partition key of the failing row contains (c1, c2) = (20, -20).
 CONTEXT:  while executing command on localhost:xxxxx
 -- see data is loaded to multi-column partitioned table

src/test/regress/sql/multi_partitioning.sql

@@ -638,6 +638,35 @@ CREATE TABLE partitioned_events_table_2009 PARTITION OF partitioned_events_table
 INSERT INTO partitioned_events_table SELECT * FROM events_table;
 INSERT INTO partitioned_users_table_2009 SELECT * FROM users_table;
 
+-- test distributed partitions are indeed colocated with the parent table
+CREATE TABLE sensors(measureid integer, eventdatetime date, measure_data jsonb, PRIMARY KEY (measureid, eventdatetime, measure_data))
+PARTITION BY RANGE(eventdatetime);
+
+CREATE TABLE sensors_old PARTITION OF sensors FOR VALUES FROM ('2000-01-01') TO ('2020-01-01');
+CREATE TABLE sensors_2020_01_01 PARTITION OF sensors FOR VALUES FROM ('2020-01-01') TO ('2020-02-01');
+CREATE TABLE sensors_new PARTITION OF sensors DEFAULT;
+
+SELECT create_distributed_table('sensors', 'measureid', colocate_with:='none');
+
+SELECT count(DISTINCT colocationid) FROM pg_dist_partition
+WHERE logicalrelid IN ('sensors'::regclass, 'sensors_old'::regclass, 'sensors_2020_01_01'::regclass, 'sensors_new'::regclass);
+
+CREATE TABLE local_sensors(measureid integer, eventdatetime date, measure_data jsonb, PRIMARY KEY (measureid, eventdatetime, measure_data))
+PARTITION BY RANGE(eventdatetime);
+
+CREATE TABLE local_sensors_old PARTITION OF local_sensors FOR VALUES FROM ('2000-01-01') TO ('2020-01-01');
+CREATE TABLE local_sensors_2020_01_01 PARTITION OF local_sensors FOR VALUES FROM ('2020-01-01') TO ('2020-02-01');
+CREATE TABLE local_sensors_new PARTITION OF local_sensors DEFAULT;
+
+SELECT create_distributed_table('local_sensors', 'measureid', colocate_with:='sensors');
+
+SELECT count(DISTINCT colocationid) FROM pg_dist_partition
+WHERE logicalrelid IN ('sensors'::regclass, 'sensors_old'::regclass, 'sensors_2020_01_01'::regclass, 'sensors_new'::regclass,
+'local_sensors'::regclass, 'local_sensors_old'::regclass, 'local_sensors_2020_01_01'::regclass, 'local_sensors_new'::regclass);
+
+DROP TABLE sensors;
+DROP TABLE local_sensors;
+
 --
 -- Complex JOINs, subqueries, UNIONs etc...
 --