implement shardgroups on reference tables;

src/backend/distributed/commands/create_distributed_table.c

@@ -1185,7 +1185,7 @@ CreateCitusTable(Oid relationId, CitusTableType tableType,
 	}
 	else if (tableType == REFERENCE_TABLE)
 	{
-		CreateReferenceTableShard(relationId);
+		CreateReferenceTableShard(relationId, colocatedTableId, colocationId);
 	}
 
 	if (ShouldSyncTableMetadata(relationId))

src/backend/distributed/operations/create_shards.c

@@ -350,7 +350,8 @@ CreateColocatedShards(Oid targetRelationId, Oid sourceRelationId, bool
  * Also, the shard is replicated to the all active nodes in the cluster.
  */
 void
-CreateReferenceTableShard(Oid distributedTableId)
+CreateReferenceTableShard(Oid distributedTableId, Oid colocatedTableId,
+						  uint32 colocationId)
 {
 	int workerStartIndex = 0;
 	text *shardMinValue = NULL;
@@ -382,6 +383,11 @@ CreateReferenceTableShard(Oid distributedTableId)
 							   tableName)));
 	}
 
+	List *insertedShardPlacements = NIL;
+	if (!OidIsValid(colocatedTableId))
+	{
+		/* create first reference table, place them on all active nodes */
+
 		/*
 		 * load and sort the worker node list for deterministic placements
 		 * create_reference_table has already acquired pg_dist_node lock
@@ -393,13 +399,80 @@ CreateReferenceTableShard(Oid distributedTableId)
 
 		/* get the next shard id */
 		uint64 shardId = GetNextShardId();
+		uint64 shardGroupId = shardId;
 
-	InsertShardRow(distributedTableId, shardId, shardStorageType, shardMinValue,
+		StringInfoData shardgroupQuery = { 0 };
-				   shardMaxValue, NULL);
+		initStringInfo(&shardgroupQuery);
 
-	List *insertedShardPlacements = InsertShardPlacementRows(distributedTableId, shardId,
+		appendStringInfoString(&shardgroupQuery,
+							   "WITH shardgroup_data(shardgroupid, colocationid, "
+							   "shardminvalue, shardmaxvalue) AS (VALUES ");
+
+		InsertShardGroupRow(shardGroupId, colocationId, shardMinValue, shardMaxValue);
+		appendStringInfo(&shardgroupQuery, "(%ld, %d, %s, %s)",
+						 shardGroupId,
+						 colocationId,
+						 TextToSQLLiteral(shardMinValue),
+						 TextToSQLLiteral(shardMaxValue));
+
+		appendStringInfoString(&shardgroupQuery, ") ");
+		appendStringInfoString(&shardgroupQuery,
+							   "SELECT pg_catalog.citus_internal_add_shardgroup_metadata("
+							   "shardgroupid, colocationid, shardminvalue, shardmaxvalue)"
+							   "FROM shardgroup_data;");
+
+		SendCommandToWorkersWithMetadata(shardgroupQuery.data);
+
+		InsertShardRow(distributedTableId, shardId, shardStorageType,
+					   shardMinValue, shardMaxValue, &shardGroupId);
+
+		insertedShardPlacements = InsertShardPlacementRows(distributedTableId, shardId,
 														   nodeList, workerStartIndex,
 														   replicationFactor);
+	}
+	else
+	{
+		/* add reference table as colocated table to already existing reference table */
+
+		/* prevent placement changes of the source relation until we colocate with them */
+		List *sourceShardIntervalList = LoadShardIntervalList(colocatedTableId);
+		LockShardListMetadata(sourceShardIntervalList, ShareLock);
+
+		if (list_length(sourceShardIntervalList) != 1)
+		{
+			elog(ERROR, "colocating a reference table to a table with shardcount > 1");
+		}
+
+		ShardInterval *sourceInterval =
+			(ShardInterval *) linitial(sourceShardIntervalList);
+		uint64 newShardId = GetNextShardId();
+
+		InsertShardRow(distributedTableId, newShardId, shardStorageType,
+					   shardMinValue, shardMaxValue, &sourceInterval->shardGroupId);
+
+		List *sourceShardPlacementList =
+			ShardPlacementList(sourceInterval->shardId);
+
+		ShardPlacement *sourceShardPlacement = NULL;
+		foreach_ptr(sourceShardPlacement, sourceShardPlacementList)
+		{
+			int32 groupId = sourceShardPlacement->groupId;
+			const uint64 shardSize = 0;
+
+			/*
+			 * Optimistically add shard placement row the pg_dist_shard_placement, in case
+			 * of any error it will be roll-backed.
+			 */
+			uint64 shardPlacementId = InsertShardPlacementRow(newShardId,
+															  INVALID_PLACEMENT_ID,
+															  shardSize,
+															  groupId);
+
+			ShardPlacement *shardPlacement =
+				LoadShardPlacement(newShardId, shardPlacementId);
+			insertedShardPlacements = lappend(insertedShardPlacements, shardPlacement);
+		}
+	}
 
 	CreateShardsOnWorkers(distributedTableId, insertedShardPlacements,
 						  useExclusiveConnection, colocatedShard);

src/include/distributed/coordinator_protocol.h

@@ -262,7 +262,8 @@ extern void CreateShardsWithRoundRobinPolicy(Oid distributedTableId, int32 shard
 											 bool useExclusiveConnections);
 extern void CreateColocatedShards(Oid targetRelationId, Oid sourceRelationId,
 								  bool useExclusiveConnections);
-extern void CreateReferenceTableShard(Oid distributedTableId);
+extern void CreateReferenceTableShard(Oid distributedTableId, Oid colocatedTableId,
+									  uint32 colocationId);
 extern List * WorkerCreateShardCommandList(Oid relationId, int shardIndex, uint64 shardId,
 										   List *ddlCommandList,
 										   List *foreignConstraintCommandList);