Disallow colocating tables when collations don't match (#8257)

DESCRIPTION: Disallows creating a distributed table or altering it to be
colocated with another table if distribution key collations don't match.

Fixes #7683.

Author: onurctirtir

src/backend/distributed/commands/alter_table.c

@@ -2111,7 +2111,7 @@ CheckAlterDistributedTableConversionParameters(TableConversionState *con)
 		{
 			ereport(ERROR, (errmsg("cannot colocate with %s and change distribution "
 								   "column to %s because data type of column %s is "
-								   "different then the distribution column of the %s",
+								   "different than the distribution column of the %s",
 								   con->colocateWith, con->distributionColumn,
 								   con->distributionColumn, con->colocateWith)));
 		}
@@ -2122,6 +2122,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 than 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

@@ -531,8 +531,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))
@@ -547,8 +545,7 @@ CreateDistributedTableConcurrently(Oid relationId, char *distributionColumnName,
 	 */
 	uint32 colocationId = FindColocateWithColocationId(relationId,
 													   replicationModel,
-													   distributionColumnType,
-													   distributionColumnCollation,
+													   distributionColumn,
 													   shardCount,
 													   shardCountIsStrict,
 													   colocateWithTableName);
@@ -697,19 +694,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);
 }
 
 
@@ -1993,10 +1985,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);
@@ -2011,8 +2004,7 @@ ColocationIdForNewTable(Oid relationId, CitusTableType tableType,
 
 		colocationId = FindColocateWithColocationId(relationId,
 													citusTableParams.replicationModel,
-													distributionColumnType,
-													distributionColumnCollation,
+													distributionColumn,
 													distributedTableParams->shardCount,
 													distributedTableParams->
 													shardCountIsStrict,

src/backend/distributed/sql/citus--13.2-1--14.0-1.sql

@@ -9,3 +9,6 @@
 #include "udfs/coord_binary_combine_agg_ffunc/14.0-1.sql"
 #include "udfs/worker_binary_partial_agg/14.0-1.sql"
 #include "udfs/coord_binary_combine_agg/14.0-1.sql"
+
+#include "udfs/fix_pre_citus14_colocation_group_collation_mismatches/14.0-1.sql"
+#include "udfs/citus_finish_citus_upgrade/14.0-1.sql"

src/backend/distributed/sql/downgrades/citus--14.0-1--13.2-1.sql

@@ -9,3 +9,6 @@ DROP AGGREGATE IF EXISTS pg_catalog.coord_binary_combine_agg(oid, bytea, anyelem
 DROP FUNCTION IF EXISTS pg_catalog.worker_binary_partial_agg_ffunc(internal);
 DROP FUNCTION IF EXISTS pg_catalog.coord_binary_combine_agg_sfunc(internal, oid, bytea, anyelement);
 DROP FUNCTION IF EXISTS pg_catalog.coord_binary_combine_agg_ffunc(internal, oid, bytea, anyelement);
+
+#include "../udfs/citus_finish_citus_upgrade/11.0-2.sql"
+DROP FUNCTION IF EXISTS pg_catalog.fix_pre_citus14_colocation_group_collation_mismatches();

src/backend/distributed/sql/udfs/citus_finish_citus_upgrade/14.0-1.sql

@@ -33,6 +33,11 @@ BEGIN
 		performed_upgrade := true;
 	END IF;
 
+	IF last_upgrade_major_version < 14 THEN
+		PERFORM fix_pre_citus14_colocation_group_collation_mismatches();
+		performed_upgrade := true;
+	END IF;
+
 	-- add new upgrade steps here
 
 	IF NOT performed_upgrade THEN

src/backend/distributed/sql/udfs/citus_finish_citus_upgrade/latest.sql

@@ -0,0 +1,83 @@
+CREATE OR REPLACE FUNCTION pg_catalog.fix_pre_citus14_colocation_group_collation_mismatches()
+RETURNS VOID AS $func$
+DECLARE
+    v_colocationid oid;
+    v_tables_to_move_out_grouped_by_collation json;
+    v_collationid oid;
+    v_tables_to_move_out oid[];
+    v_table_to_move_out oid;
+    v_first_table oid;
+BEGIN
+    SET LOCAL search_path TO pg_catalog;
+
+    FOR v_colocationid, v_tables_to_move_out_grouped_by_collation
+    IN
+       WITH colocation_groups_and_tables_with_collation_mismatches AS (
+           SELECT pdc.colocationid, pa.attcollation as distkeycollation, pdp.logicalrelid
+             FROM pg_dist_colocation pdc
+             JOIN pg_dist_partition pdp
+               ON pdc.colocationid = pdp.colocationid
+             JOIN pg_attribute pa
+               ON pa.attrelid = pdp.logicalrelid
+              AND pa.attname = column_to_column_name(pdp.logicalrelid, pdp.partkey)
+                  -- column_to_column_name() returns NULL if partkey is NULL, so we're already
+                  -- implicitly ignoring the tables that don't have a distribution column, such
+                  -- as reference tables, but let's still explicitly discard such tables below.
+            WHERE pdp.partkey IS NOT NULL
+                  -- ignore the table if its distribution column collation matches the collation saved for the colocation group
+              AND pdc.distributioncolumncollation != pa.attcollation
+       )
+       SELECT
+           colocationid,
+           json_object_agg(distkeycollation, rels) AS tables_to_move_out_grouped_by_collation
+       FROM (
+               SELECT colocationid,
+                      distkeycollation,
+                      array_agg(logicalrelid::oid) AS rels
+                 FROM colocation_groups_and_tables_with_collation_mismatches
+             GROUP BY colocationid, distkeycollation
+       ) q
+       GROUP BY colocationid
+    LOOP
+        RAISE DEBUG 'Processing colocation group with id %', v_colocationid;
+
+        FOR v_collationid, v_tables_to_move_out
+        IN
+              SELECT key::oid AS collationid,
+                     array_agg(elem::oid) AS tables_to_move_out
+                FROM json_each(v_tables_to_move_out_grouped_by_collation) AS e(key, value),
+             LATERAL json_array_elements_text(e.value) AS elem
+            GROUP BY key
+        LOOP
+            RAISE DEBUG 'Moving out tables with collation id % from colocation group %', v_collationid, v_colocationid;
+
+            v_first_table := NULL;
+
+            FOR v_table_to_move_out IN SELECT unnest(v_tables_to_move_out)
+            LOOP
+                IF v_first_table IS NULL then
+                    -- Move the first table out to start a new colocation group.
+                    --
+                    -- Could check if there is an appropriate colocation group to move to instead of 'none',
+                    -- but this won't be super easy. Plus, even if we had such a colocation group, the user
+                    -- was anyways okay with having this in a different colocation group in the first place.
+
+                    RAISE DEBUG 'Moving out table with oid % to a new colocation group', v_table_to_move_out;
+                    PERFORM update_distributed_table_colocation(v_table_to_move_out, colocate_with => 'none');
+
+                    -- save the first table to colocate the rest of the tables with it
+                    v_first_table := v_table_to_move_out;
+                ELSE
+                    -- Move the rest of the tables to colocate with the first table.
+
+                    RAISE DEBUG 'Moving out table with oid % to colocate with table with oid %', v_table_to_move_out, v_first_table;
+                    PERFORM update_distributed_table_colocation(v_table_to_move_out, colocate_with => v_first_table::regclass::text);
+                END IF;
+            END LOOP;
+        END LOOP;
+    END LOOP;
+END;
+$func$
+LANGUAGE plpgsql;
+COMMENT ON FUNCTION pg_catalog.fix_pre_citus14_colocation_group_collation_mismatches()
+  IS 'Fix distributed tables whose colocation group collations do not match their distribution columns by moving them to new colocation groups';