Merge cda0fd0118 into 662b7248db

src/backend/distributed/commands/alter_table.c

@@ -2125,6 +2125,23 @@ CheckAlterDistributedTableConversionParameters(TableConversionState *con)
 								   "distribution column is different than %s",
 								   con->colocateWith, con->relationName)));
 		}
+		else if (con->distributionColumn &&
+				 colocateWithPartKey->varcollid != con->distributionKey->varcollid)
+		{
+			ereport(ERROR, (errmsg("cannot colocate with %s and change distribution "
+								   "column to %s because collation of column %s is "
+								   "different then the distribution column of the %s",
+								   con->colocateWith, con->distributionColumn,
+								   con->distributionColumn, con->colocateWith)));
+		}
+		else if (!con->distributionColumn &&
+				 colocateWithPartKey->varcollid !=
+				 con->originalDistributionKey->varcollid)
+		{
+			ereport(ERROR, (errmsg("cannot colocate with %s because collation of its "
+								   "distribution column is different than %s",
+								   con->colocateWith, con->relationName)));
+		}
 	}
 
 	if (!con->suppressNoticeMessages)

src/backend/distributed/commands/create_distributed_table.c

@@ -530,8 +530,6 @@ CreateDistributedTableConcurrently(Oid relationId, char *distributionColumnName,
 	Var *distributionColumn = BuildDistributionKeyFromColumnName(relationId,
 																 distributionColumnName,
 																 NoLock);
-	Oid distributionColumnType = distributionColumn->vartype;
-	Oid distributionColumnCollation = distributionColumn->varcollid;
 
 	/* get an advisory lock to serialize concurrent default group creations */
 	if (IsColocateWithDefault(colocateWithTableName))
@@ -546,8 +544,7 @@ CreateDistributedTableConcurrently(Oid relationId, char *distributionColumnName,
 	 */
 	uint32 colocationId = FindColocateWithColocationId(relationId,
 													   replicationModel,
-													   distributionColumnType,
+													   distributionColumn,
-													   distributionColumnCollation,
 													   shardCount,
 													   shardCountIsStrict,
 													   colocateWithTableName);
@@ -696,18 +693,14 @@ EnsureColocateWithTableIsValid(Oid relationId, char distributionMethod,
 	char replicationModel = DecideDistTableReplicationModel(distributionMethod,
 															colocateWithTableName);
 
-	/*
-	 * we fail transaction before local table conversion if the table could not be colocated with
-	 * given table. We should make those checks after local table conversion by acquiring locks to
-	 * the relation because the distribution column can be modified in that period.
-	 */
-	Oid distributionColumnType = ColumnTypeIdForRelationColumnName(relationId,
-																   distributionColumnName);
-
 	text *colocateWithTableNameText = cstring_to_text(colocateWithTableName);
 	Oid colocateWithTableId = ResolveRelationId(colocateWithTableNameText, false);
+
+	Var *distributionColumn = BuildDistributionKeyFromColumnName(relationId,
+																 distributionColumnName,
+																 NoLock);
 	EnsureTableCanBeColocatedWith(relationId, replicationModel,
-								  distributionColumnType, colocateWithTableId);
+								  distributionColumn, colocateWithTableId);
 }
 
 
@@ -1998,10 +1991,11 @@ ColocationIdForNewTable(Oid relationId, CitusTableType tableType,
 		 * until this transaction is committed.
 		 */
 
+		/* distributionColumn can only be null for single-shard tables */
 		Oid distributionColumnType =
 			distributionColumn ? distributionColumn->vartype : InvalidOid;
 		Oid distributionColumnCollation =
-			distributionColumn ? get_typcollation(distributionColumnType) : InvalidOid;
+			distributionColumn ? distributionColumn->varcollid : InvalidOid;
 
 		Assert(distributedTableParams->colocationParam.colocationParamType ==
 			   COLOCATE_WITH_TABLE_LIKE_OPT);
@@ -2016,8 +2010,7 @@ ColocationIdForNewTable(Oid relationId, CitusTableType tableType,
 
 		colocationId = FindColocateWithColocationId(relationId,
 													citusTableParams.replicationModel,
-													distributionColumnType,
+													distributionColumn,
-													distributionColumnCollation,
 													distributedTableParams->shardCount,
 													distributedTableParams->
 													shardCountIsStrict,

src/backend/distributed/utils/colocation_utils.c

@@ -1378,8 +1378,7 @@ DeleteColocationGroupLocally(uint32 colocationId)
  */
 uint32
 FindColocateWithColocationId(Oid relationId, char replicationModel,
-							 Oid distributionColumnType,
+							 Var *distributionColumn,
-							 Oid distributionColumnCollation,
 							 int shardCount, bool shardCountIsStrict,
 							 char *colocateWithTableName)
 {
@@ -1387,6 +1386,12 @@ FindColocateWithColocationId(Oid relationId, char replicationModel,
 
 	if (IsColocateWithDefault(colocateWithTableName))
 	{
+		/* distributionColumn can only be null for single-shard tables */
+		Oid distributionColumnType =
+			distributionColumn ? distributionColumn->vartype : InvalidOid;
+		Oid distributionColumnCollation =
+			distributionColumn ? distributionColumn->varcollid : InvalidOid;
+
 		/* check for default colocation group */
 		colocationId = ColocationId(shardCount, ShardReplicationFactor,
 									distributionColumnType,
@@ -1419,7 +1424,7 @@ FindColocateWithColocationId(Oid relationId, char replicationModel,
 		Oid sourceRelationId = ResolveRelationId(colocateWithTableNameText, false);
 
 		EnsureTableCanBeColocatedWith(relationId, replicationModel,
-									  distributionColumnType, sourceRelationId);
+									  distributionColumn, sourceRelationId);
 
 		colocationId = TableColocationId(sourceRelationId);
 	}
@@ -1437,7 +1442,7 @@ FindColocateWithColocationId(Oid relationId, char replicationModel,
  */
 void
 EnsureTableCanBeColocatedWith(Oid relationId, char replicationModel,
-							  Oid distributionColumnType, Oid sourceRelationId)
+							  Var *distributionColumn, Oid sourceRelationId)
 {
 	CitusTableCacheEntry *sourceTableEntry = GetCitusTableCacheEntry(sourceRelationId);
 
@@ -1465,19 +1470,8 @@ EnsureTableCanBeColocatedWith(Oid relationId, char replicationModel,
 	}
 
 	Var *sourceDistributionColumn = DistPartitionKey(sourceRelationId);
-	Oid sourceDistributionColumnType = !sourceDistributionColumn ? InvalidOid :
+	EnsureColumnTypeEquality(sourceRelationId, relationId,
-									   sourceDistributionColumn->vartype;
+							 sourceDistributionColumn, distributionColumn);
-	if (sourceDistributionColumnType != distributionColumnType)
-	{
-		char *relationName = get_rel_name(relationId);
-		char *sourceRelationName = get_rel_name(sourceRelationId);
-
-		ereport(ERROR, (errmsg("cannot colocate tables %s and %s",
-							   sourceRelationName, relationName),
-						errdetail("Distribution column types don't match for "
-								  "%s and %s.", sourceRelationName,
-								  relationName)));
-	}
 
 	/* prevent colocating regular tables with tenant tables */
 	Oid sourceRelationSchemaId = get_rel_namespace(sourceRelationId);

src/include/distributed/colocation_utils.h

@@ -53,13 +53,11 @@ extern List * ColocationGroupTableList(uint32 colocationId, uint32 count);
 extern void DeleteColocationGroup(uint32 colocationId);
 extern void DeleteColocationGroupLocally(uint32 colocationId);
 extern uint32 FindColocateWithColocationId(Oid relationId, char replicationModel,
-										   Oid distributionColumnType,
+										   Var *distributionColumn,
-										   Oid distributionColumnCollation,
 										   int shardCount, bool shardCountIsStrict,
 										   char *colocateWithTableName);
 extern void EnsureTableCanBeColocatedWith(Oid relationId, char replicationModel,
-										  Oid distributionColumnType,
+										  Var *distributionColumn, Oid sourceRelationId);
-										  Oid sourceRelationId);
 extern void AcquireColocationDefaultLock(void);
 extern void ReleaseColocationDefaultLock(void);
 

src/test/regress/expected/single_node.out

@@ -278,6 +278,107 @@ SELECT
  t
 (1 row)
 
+create table test_a_tbl_1(text_col text collate "C" unique);
+create table test_a_tbl_2(text_col text collate "en-x-icu" unique);
+select create_distributed_table('test_a_tbl_1', 'text_col');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select create_distributed_table('test_a_tbl_2', 'text_col');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- make sure we assing them to different colocation groups
+select count(*) = 2 as colocationids_different from (
+	select distinct(colocationid)
+	from pg_dist_partition
+	join pg_class on (logicalrelid = pg_class.oid)
+	join pg_namespace on (relnamespace = pg_namespace.oid)
+	join pg_dist_colocation using (colocationid)
+	where pg_namespace.nspname = 'single_node'
+		and pg_class.relname in ('test_a_tbl_1', 'test_a_tbl_2')
+) q;
+ colocationids_different
+---------------------------------------------------------------------
+ t
+(1 row)
+
+DROP TABLE test_a_tbl_1, test_a_tbl_2;
+create table test_d_tbl_1(text_col text collate "C" unique);
+create table test_d_tbl_2(text_col text collate "en-x-icu" unique);
+select create_distributed_table('test_d_tbl_1', 'text_col', shard_count=>4);
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select create_distributed_table('test_d_tbl_2', 'text_col', shard_count=>6);
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select alter_distributed_table('test_d_tbl_2', shard_count=>4);
+NOTICE:  creating a new table for single_node.test_d_tbl_2
+NOTICE:  moving the data of single_node.test_d_tbl_2
+NOTICE:  dropping the old single_node.test_d_tbl_2
+NOTICE:  renaming the new table to single_node.test_d_tbl_2
+ alter_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- make sure we assing them to different colocation groups
+select count(*) = 2 as colocationids_different from (
+	select distinct(colocationid)
+	from pg_dist_partition
+	join pg_class on (logicalrelid = pg_class.oid)
+	join pg_namespace on (relnamespace = pg_namespace.oid)
+	join pg_dist_colocation using (colocationid)
+	where pg_namespace.nspname = 'single_node'
+		and pg_class.relname in ('test_d_tbl_1', 'test_d_tbl_2')
+) q;
+ colocationids_different
+---------------------------------------------------------------------
+ t
+(1 row)
+
+DROP TABLE test_d_tbl_1, test_d_tbl_2;
+create table test_b_tbl_1(text_col text collate "C" unique);
+create table test_b_tbl_2(text_col text collate "en-x-icu" unique);
+select create_distributed_table('test_b_tbl_1', 'text_col');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- errors
+select create_distributed_table('test_b_tbl_2', 'text_col', colocate_with => 'test_b_tbl_1');
+ERROR:  cannot colocate tables test_b_tbl_1 and test_b_tbl_2
+DETAIL:  Distribution column collations don't match for test_b_tbl_1 and test_b_tbl_2.
+DROP TABLE test_b_tbl_1, test_b_tbl_2;
+create table test_c_tbl_1(text_col text collate "C" unique);
+create table test_c_tbl_2(text_col text collate "en-x-icu" unique);
+select create_distributed_table('test_c_tbl_1', 'text_col');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select create_distributed_table('test_c_tbl_2', 'text_col', colocate_with => 'none');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- errors
+select alter_distributed_table('test_c_tbl_2', colocate_with=>'test_c_tbl_1');
+ERROR:  cannot colocate with test_c_tbl_1 because collation of its distribution column is different than test_c_tbl_2
+DROP TABLE test_c_tbl_1, test_c_tbl_2;
 SET client_min_messages TO WARNING;
 DROP TABLE failover_to_local, single_node_nullkey_c1, single_node_nullkey_c2, ref_table_conversion_test, single_shard_conversion_test_1, single_shard_conversion_test_2;
 DROP SCHEMA tenant_1 CASCADE;

src/test/regress/expected/update_colocation_mx.out

@@ -95,6 +95,125 @@ SELECT count(DISTINCT colocationid) FROM pg_dist_partition WHERE logicalrelid IN
      2
 (1 row)
 
+\c - - - :master_port
+SET search_path TO "Update Colocation";
+create table test_a_tbl_1(text_col text collate "C" unique);
+create table test_a_tbl_2(text_col text collate "en-x-icu" unique);
+select create_distributed_table('test_a_tbl_1', 'text_col');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select create_distributed_table('test_a_tbl_2', 'text_col');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- make sure we assing them to different colocation groups
+select result as colocationids_different from run_command_on_all_nodes($$
+    select count(*) = 2 from (
+        select distinct(colocationid)
+        from pg_dist_partition
+        join pg_class on (logicalrelid = pg_class.oid)
+        join pg_namespace on (relnamespace = pg_namespace.oid)
+        join pg_dist_colocation using (colocationid)
+        where pg_namespace.nspname = 'Update Colocation'
+          and pg_class.relname in ('test_a_tbl_1', 'test_a_tbl_2')
+    ) q;
+$$);
+ colocationids_different
+---------------------------------------------------------------------
+ t
+ t
+ t
+(3 rows)
+
+DROP TABLE test_a_tbl_1, test_a_tbl_2;
+create table test_d_tbl_1(text_col text collate "C" unique);
+create table test_d_tbl_2(text_col text collate "en-x-icu" unique);
+select create_distributed_table('test_d_tbl_1', 'text_col', shard_count=>4);
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select create_distributed_table('test_d_tbl_2', 'text_col', shard_count=>6);
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select alter_distributed_table('test_d_tbl_2', shard_count=>4);
+NOTICE:  creating a new table for "Update Colocation".test_d_tbl_2
+NOTICE:  moving the data of "Update Colocation".test_d_tbl_2
+NOTICE:  dropping the old "Update Colocation".test_d_tbl_2
+NOTICE:  renaming the new table to "Update Colocation".test_d_tbl_2
+ alter_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- make sure we assing them to different colocation groups
+select result as colocationids_different from run_command_on_all_nodes($$
+    select count(*) = 2 from (
+        select distinct(colocationid)
+        from pg_dist_partition
+        join pg_class on (logicalrelid = pg_class.oid)
+        join pg_namespace on (relnamespace = pg_namespace.oid)
+        join pg_dist_colocation using (colocationid)
+        where pg_namespace.nspname = 'Update Colocation'
+          and pg_class.relname in ('test_d_tbl_1', 'test_d_tbl_2')
+    ) q;
+$$);
+ colocationids_different
+---------------------------------------------------------------------
+ t
+ t
+ t
+(3 rows)
+
+DROP TABLE test_d_tbl_1, test_d_tbl_2;
+create table test_b_tbl_1(text_col text collate "C" unique);
+create table test_b_tbl_2(text_col text collate "en-x-icu" unique);
+select create_distributed_table('test_b_tbl_1', 'text_col');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- errors
+select create_distributed_table('test_b_tbl_2', 'text_col', colocate_with => 'test_b_tbl_1');
+ERROR:  cannot colocate tables test_b_tbl_1 and test_b_tbl_2
+DETAIL:  Distribution column collations don't match for test_b_tbl_1 and test_b_tbl_2.
+DROP TABLE test_b_tbl_1, test_b_tbl_2;
+create table test_c_tbl_1(text_col text collate "C" unique);
+create table test_c_tbl_2(text_col text collate "en-x-icu" unique);
+select create_distributed_table('test_c_tbl_1', 'text_col');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+select create_distributed_table('test_c_tbl_2', 'text_col', colocate_with => 'none');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- errors
+select alter_distributed_table('test_c_tbl_2', colocate_with=>'test_c_tbl_1');
+ERROR:  cannot colocate with test_c_tbl_1 because collation of its distribution column is different than test_c_tbl_2
+DROP TABLE test_c_tbl_1, test_c_tbl_2;
 \c - postgres - :master_port
 SET client_min_messages TO ERROR;
 DROP SCHEMA "Update Colocation" cascade;
+SET citus.enable_ddl_propagation TO OFF;
+DROP ROLE mx_update_colocation;
+\c - postgres - :worker_1_port
+SET citus.enable_ddl_propagation TO OFF;
+DROP ROLE mx_update_colocation;
+\c - postgres - :worker_2_port
+SET citus.enable_ddl_propagation TO OFF;
+DROP ROLE mx_update_colocation;

src/test/regress/sql/single_node.sql

@@ -168,6 +168,67 @@ SELECT
     WHERE logicalrelid = 'single_node.single_shard_conversion_test_2'::regclass
 );
 
+create table test_a_tbl_1(text_col text collate "C" unique);
+create table test_a_tbl_2(text_col text collate "en-x-icu" unique);
+
+select create_distributed_table('test_a_tbl_1', 'text_col');
+select create_distributed_table('test_a_tbl_2', 'text_col');
+
+-- make sure we assing them to different colocation groups
+select count(*) = 2 as colocationids_different from (
+	select distinct(colocationid)
+	from pg_dist_partition
+	join pg_class on (logicalrelid = pg_class.oid)
+	join pg_namespace on (relnamespace = pg_namespace.oid)
+	join pg_dist_colocation using (colocationid)
+	where pg_namespace.nspname = 'single_node'
+		and pg_class.relname in ('test_a_tbl_1', 'test_a_tbl_2')
+) q;
+
+DROP TABLE test_a_tbl_1, test_a_tbl_2;
+
+create table test_d_tbl_1(text_col text collate "C" unique);
+create table test_d_tbl_2(text_col text collate "en-x-icu" unique);
+
+select create_distributed_table('test_d_tbl_1', 'text_col', shard_count=>4);
+select create_distributed_table('test_d_tbl_2', 'text_col', shard_count=>6);
+select alter_distributed_table('test_d_tbl_2', shard_count=>4);
+
+-- make sure we assing them to different colocation groups
+select count(*) = 2 as colocationids_different from (
+	select distinct(colocationid)
+	from pg_dist_partition
+	join pg_class on (logicalrelid = pg_class.oid)
+	join pg_namespace on (relnamespace = pg_namespace.oid)
+	join pg_dist_colocation using (colocationid)
+	where pg_namespace.nspname = 'single_node'
+		and pg_class.relname in ('test_d_tbl_1', 'test_d_tbl_2')
+) q;
+
+
+DROP TABLE test_d_tbl_1, test_d_tbl_2;
+
+create table test_b_tbl_1(text_col text collate "C" unique);
+create table test_b_tbl_2(text_col text collate "en-x-icu" unique);
+
+select create_distributed_table('test_b_tbl_1', 'text_col');
+
+-- errors
+select create_distributed_table('test_b_tbl_2', 'text_col', colocate_with => 'test_b_tbl_1');
+
+DROP TABLE test_b_tbl_1, test_b_tbl_2;
+
+create table test_c_tbl_1(text_col text collate "C" unique);
+create table test_c_tbl_2(text_col text collate "en-x-icu" unique);
+
+select create_distributed_table('test_c_tbl_1', 'text_col');
+select create_distributed_table('test_c_tbl_2', 'text_col', colocate_with => 'none');
+
+-- errors
+select alter_distributed_table('test_c_tbl_2', colocate_with=>'test_c_tbl_1');
+
+DROP TABLE test_c_tbl_1, test_c_tbl_2;
+
 SET client_min_messages TO WARNING;
 DROP TABLE failover_to_local, single_node_nullkey_c1, single_node_nullkey_c2, ref_table_conversion_test, single_shard_conversion_test_1, single_shard_conversion_test_2;
 DROP SCHEMA tenant_1 CASCADE;

src/test/regress/sql/update_colocation_mx.sql

@@ -53,6 +53,84 @@ SELECT count(DISTINCT colocationid) FROM pg_dist_partition WHERE logicalrelid IN
 SET search_path TO "Update Colocation";
 SELECT count(DISTINCT colocationid) FROM pg_dist_partition WHERE logicalrelid IN ('t1'::regclass, 't2'::regclass);
 
+\c - - - :master_port
+SET search_path TO "Update Colocation";
+
+create table test_a_tbl_1(text_col text collate "C" unique);
+create table test_a_tbl_2(text_col text collate "en-x-icu" unique);
+
+select create_distributed_table('test_a_tbl_1', 'text_col');
+select create_distributed_table('test_a_tbl_2', 'text_col');
+
+-- make sure we assing them to different colocation groups
+select result as colocationids_different from run_command_on_all_nodes($$
+    select count(*) = 2 from (
+        select distinct(colocationid)
+        from pg_dist_partition
+        join pg_class on (logicalrelid = pg_class.oid)
+        join pg_namespace on (relnamespace = pg_namespace.oid)
+        join pg_dist_colocation using (colocationid)
+        where pg_namespace.nspname = 'Update Colocation'
+          and pg_class.relname in ('test_a_tbl_1', 'test_a_tbl_2')
+    ) q;
+$$);
+
+DROP TABLE test_a_tbl_1, test_a_tbl_2;
+
+create table test_d_tbl_1(text_col text collate "C" unique);
+create table test_d_tbl_2(text_col text collate "en-x-icu" unique);
+
+select create_distributed_table('test_d_tbl_1', 'text_col', shard_count=>4);
+select create_distributed_table('test_d_tbl_2', 'text_col', shard_count=>6);
+select alter_distributed_table('test_d_tbl_2', shard_count=>4);
+
+-- make sure we assing them to different colocation groups
+select result as colocationids_different from run_command_on_all_nodes($$
+    select count(*) = 2 from (
+        select distinct(colocationid)
+        from pg_dist_partition
+        join pg_class on (logicalrelid = pg_class.oid)
+        join pg_namespace on (relnamespace = pg_namespace.oid)
+        join pg_dist_colocation using (colocationid)
+        where pg_namespace.nspname = 'Update Colocation'
+          and pg_class.relname in ('test_d_tbl_1', 'test_d_tbl_2')
+    ) q;
+$$);
+
+DROP TABLE test_d_tbl_1, test_d_tbl_2;
+
+create table test_b_tbl_1(text_col text collate "C" unique);
+create table test_b_tbl_2(text_col text collate "en-x-icu" unique);
+
+select create_distributed_table('test_b_tbl_1', 'text_col');
+
+-- errors
+select create_distributed_table('test_b_tbl_2', 'text_col', colocate_with => 'test_b_tbl_1');
+
+DROP TABLE test_b_tbl_1, test_b_tbl_2;
+
+create table test_c_tbl_1(text_col text collate "C" unique);
+create table test_c_tbl_2(text_col text collate "en-x-icu" unique);
+
+select create_distributed_table('test_c_tbl_1', 'text_col');
+select create_distributed_table('test_c_tbl_2', 'text_col', colocate_with => 'none');
+
+-- errors
+select alter_distributed_table('test_c_tbl_2', colocate_with=>'test_c_tbl_1');
+
+DROP TABLE test_c_tbl_1, test_c_tbl_2;
+
 \c - postgres - :master_port
 SET client_min_messages TO ERROR;
 DROP SCHEMA "Update Colocation" cascade;
+
+SET citus.enable_ddl_propagation TO OFF;
+DROP ROLE mx_update_colocation;
+
+\c - postgres - :worker_1_port
+SET citus.enable_ddl_propagation TO OFF;
+DROP ROLE mx_update_colocation;
+
+\c - postgres - :worker_2_port
+SET citus.enable_ddl_propagation TO OFF;
+DROP ROLE mx_update_colocation;