Disallow colocating tables when collations don't match

- No need to make any changes for the following UDFs:

  - update_distributed_table_colocation()/create_distributed_table_concurrently()

    These are already checking whether collations of distribution keys match.

  - ALTER TABLE .. ALTER COLUMN .. TYPE .. COLLATE ..

    We already don't allow changing the collation of a distribution column.

- The following UDFs are fixed with this commit:

  These UDFs were not explicitly checking whether collation of the distribution
  keys of two distributed tables match, however, we were catching this via
  EnsurePartitionMetadataIsSane() when syncing metadata. And this incorrectly
  caused allowing to colocate two such distributed tables in a single-node setup
  as we don't have any nodes to sync metadata there. In such cases, we would only
  become aware of the situation **when adding a new node to the cluster etc.**.

  Otherwise, i.e., if we already had more than one nodes in the cluster, we were
  seeing this error while Citus is syncing metadata for the new / altered distributed
  table, which wasn't too bad but not ideal too:
  ```sql
  ERROR:  cannot colocate tables test_a_tbl_2 and test_a_tbl_1
  DETAIL:  Distribution column collations don't match for test_a_tbl_2 and test_a_tbl_1.
  CONTEXT:  while executing command on localhost:9701
  ```

  - create_distributed_table()

    To fix that, now EnsureTableCanBeColocatedWith() calls EnsureColumnTypeEquality()
    to perform necessary checks, which already ensures collation match.

    And to do that, now EnsureTableCanBeColocatedWith() has to accept distributionColumn
    Var as an argument (to be passed into EnsureColumnTypeEquality()), so now it doesn't
    have to accept distributionColumnType separately, so refactor accordingly.

    And to provide distributionColumn Var to EnsureTableCanBeColocatedWith(), similarly
    the call sites of the function also need to have distributionColumn Var, like
    FindColocateWithColocationId(). And now that FindColocateWithColocationId() accepts
    distributionColumn Var, it doesn't have to accept distributionColumnType and
    distributionColumnCollation separately, so refactor accordingly.

    Plus, now this UDF also correctly records the colocation groups with the correct
    distribution column collations.

  - alter_distributed_table()

    Make sure to check column collations.

- Also, now citus_finish_citus_upgrade() automatically fixes such colocation groups when
  executed after upgrading to Citus 14.

  This also helps removing the last alternative test output for
  upgrade_citus_finish_citus_upgrade. This is because, the recent change made in
  citus_finish_citus_upgrade() to automatically fix such colocation groups results in
  updating last_upgrade_version in pg_dist_node_metadata when upgrading to Citus 14,
  so the notice message in the mentioned test is now gone, so we don't anymore need to
  have an alternative output for that test.

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

@@ -3,3 +3,5 @@
 
 #include "udfs/citus_prepare_pg_upgrade/14.0-1.sql"
 #include "udfs/citus_finish_pg_upgrade/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

@@ -3,3 +3,5 @@
 
 #include "../udfs/citus_prepare_pg_upgrade/13.0-1.sql"
 #include "../udfs/citus_finish_pg_upgrade/13.2-1.sql"
+#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,79 @@
+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)
+            -- ignore the table if its distribution column collation matches the collation saved for the colocation group
+            WHERE 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 ORDER BY logicalrelid) 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 id % 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 id % to colocate with table id %', 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';