Don't allow tables with different replication models to be colocated

src/backend/distributed/commands/create_distributed_table.c

@@ -889,7 +889,7 @@ CreateHashDistributedTable(Oid relationId, char *distributionColumnName,
 	Relation distributedRelation = NULL;
 	Relation pgDistColocation = NULL;
 	uint32 colocationId = INVALID_COLOCATION_ID;
-	Oid colocationTableId = InvalidOid;
+	Oid sourceRelationId = InvalidOid;
 	Oid distributionColumnType = InvalidOid;
 
 	/* get an access lock on the relation to prevent DROP TABLE and ALTER TABLE */
@@ -917,7 +917,7 @@ CreateHashDistributedTable(Oid relationId, char *distributionColumnName,
 		}
 		else
 		{
-			colocationTableId = ColocatedTableId(colocationId);
+			sourceRelationId = ColocatedTableId(colocationId);
 		}
 	}
 	else if (pg_strncasecmp(colocateWithTableName, "none", NAMEDATALEN) == 0)
@@ -931,11 +931,11 @@ CreateHashDistributedTable(Oid relationId, char *distributionColumnName,
 
 		/* get colocation group of the target table */
 		text *colocateWithTableNameText = cstring_to_text(colocateWithTableName);
-		colocationTableId = ResolveRelationId(colocateWithTableNameText);
+		sourceRelationId = ResolveRelationId(colocateWithTableNameText);
 
-		colocationId = TableColocationId(colocationTableId);
+		colocationId = TableColocationId(sourceRelationId);
 
-		colocationTablePartitionColumn = PartitionKey(colocationTableId);
+		colocationTablePartitionColumn = PartitionKey(sourceRelationId);
 		colocationTablePartitionColumnType = colocationTablePartitionColumn->vartype;
 
 		if (colocationTablePartitionColumnType != distributionColumnType)
@@ -951,9 +951,11 @@ CreateHashDistributedTable(Oid relationId, char *distributionColumnName,
 							  colocationId);
 
 	/* create shards */
-	if (colocationTableId != InvalidOid)
+	if (sourceRelationId != InvalidOid)
 	{
-		CreateColocatedShards(relationId, colocationTableId);
+		CheckReplicationModel(sourceRelationId, relationId);
+
+		CreateColocatedShards(relationId, sourceRelationId);
 	}
 	else
 	{

src/backend/distributed/utils/colocation_utils.c

@@ -75,6 +75,9 @@ mark_tables_colocated(PG_FUNCTION_ARGS)
 	for (relationIndex = 0; relationIndex < relationCount; relationIndex++)
 	{
 		Oid nextRelationOid = DatumGetObjectId(relationIdDatumArray[relationIndex]);
+
+		CheckReplicationModel(sourceRelationId, nextRelationOid);
+
 		MarkTablesColocated(sourceRelationId, nextRelationOid);
 	}
 
@@ -483,6 +486,34 @@ GetNextColocationId()
 }
 
 
+/*
+ * CheckReplicationModel checks if given relation and colocation group are from
+ * the same replication model. Otherwise, it errors out.
+ */
+void
+CheckReplicationModel(Oid sourceRelationId, Oid targetRelationId)
+{
+	DistTableCacheEntry *sourceTableEntry = NULL;
+	DistTableCacheEntry *targetTableEntry = NULL;
+	char sourceReplicationModel = 0;
+	char targetReplicationModel = 0;
+
+	sourceTableEntry = DistributedTableCacheEntry(sourceRelationId);
+	sourceReplicationModel = sourceTableEntry->replicationModel;
+
+	targetTableEntry = DistributedTableCacheEntry(targetRelationId);
+	targetReplicationModel = targetTableEntry->replicationModel;
+
+	if (sourceReplicationModel != targetReplicationModel)
+	{
+		ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+						errmsg("cannot create colocation"),
+						errdetail("Colocating tables with different replication "
+								  "models is not supported.")));
+	}
+}
+
+
 /*
  * UpdateRelationColocationGroup updates colocation group in pg_dist_partition
  * for the given relation.

src/include/distributed/colocation_utils.h

@@ -29,6 +29,7 @@ uint32 ColocationId(int shardCount, int replicationFactor, Oid distributionColum
 extern uint32 CreateColocationGroup(int shardCount, int replicationFactor,
 									Oid distributionColumnType);
 extern uint32 GetNextColocationId(void);
+extern void CheckReplicationModel(Oid sourceRelationId, Oid targetRelationId);
 
 
 #endif /* COLOCATION_UTILS_H_ */

src/test/regress/expected/multi_colocation_utils.out

@@ -611,8 +611,8 @@ SELECT create_distributed_table('table_failing', 'id', colocate_with => NULL);
 -- 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 with table1_groupE
+ERROR:  cannot colocate tables table1_groupe and table_bigint
-DETAIL:  Distribution column types are different.
+DETAIL:  Distribution column types don't match for table1_groupe and table_bigint.
 -- check worker table schemas
 \c - - - :worker_1_port
 \d table3_groupE_1300050
@@ -921,8 +921,8 @@ SELECT * FROM pg_dist_colocation
 (5 rows)
 
 SELECT logicalrelid, colocationid FROM pg_dist_partition
-    WHERE colocationid >= 1 AND colocationid < 1000 
+    WHERE colocationid >= 1 AND colocationid < 1000
-    ORDER BY logicalrelid;
+    ORDER BY colocationid, logicalrelid;
    logicalrelid    | colocationid 
 -------------------+--------------
  table1_groupb     |            2
@@ -986,3 +986,49 @@ SELECT logicalrelid, colocationid FROM pg_dist_partition
  table2_group_none |            7
 (13 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;

src/test/regress/sql/multi_colocation_utils.sql

@@ -399,8 +399,8 @@ SELECT * FROM pg_dist_colocation
     ORDER BY colocationid;
 
 SELECT logicalrelid, colocationid FROM pg_dist_partition
-    WHERE colocationid >= 1 AND colocationid < 1000 
+    WHERE colocationid >= 1 AND colocationid < 1000
-    ORDER BY logicalrelid;
+    ORDER BY colocationid, logicalrelid;
 
 -- 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']);
@@ -416,3 +416,40 @@ SELECT * FROM pg_dist_colocation
 SELECT logicalrelid, colocationid FROM pg_dist_partition
     WHERE colocationid >= 1 AND colocationid < 1000
     ORDER BY colocationid, logicalrelid;
+
+-- try to colocate different replication models
+CREATE TABLE table1_groupG ( id int );
+SELECT create_distributed_table('table1_groupG', 'id');
+
+-- 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');
+
+CREATE TABLE table2_groupG ( id int );
+SELECT create_distributed_table('table2_groupG', 'id', colocate_with => 'NONE');
+
+SELECT mark_tables_colocated('table1_groupG', ARRAY['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;