Remove colocation group if no table exists in it

src/backend/distributed/commands/drop_distributed_table.c

@@ -12,6 +12,7 @@
 #include "miscadmin.h"
 
 
+#include "distributed/colocation_utils.h"
 #include "distributed/commands/utility_hook.h"
 #include "distributed/commands.h"
 #include "distributed/metadata_utility.h"
@@ -70,6 +71,8 @@ master_remove_partition_metadata(PG_FUNCTION_ARGS)
 	char *schemaName = text_to_cstring(schemaNameText);
 	char *tableName = text_to_cstring(tableNameText);
 
+	uint32 colocationId = ColocationIdViaCatalog(relationId);
+
 	/*
 	 * The SQL_DROP trigger calls this function even for tables that are
 	 * not distributed. In that case, silently ignore. This is not very
@@ -87,6 +90,8 @@ master_remove_partition_metadata(PG_FUNCTION_ARGS)
 
 	DeletePartitionRow(relationId);
 
+	DeleteColocationGroupIfNoTablesBelong(colocationId);
+
 	PG_RETURN_VOID();
 }
 

src/backend/distributed/metadata/metadata_cache.c

@@ -594,6 +594,45 @@ PartitionColumnViaCatalog(Oid relationId)
 }
 
 
+/*
+ * ColocationIdViaCatalog gets a relationId and returns the colocation
+ * id column from pg_dist_partition via reading from catalog.
+ */
+uint32
+ColocationIdViaCatalog(Oid relationId)
+{
+	HeapTuple partitionTuple = PgDistPartitionTupleViaCatalog(relationId);
+	if (!HeapTupleIsValid(partitionTuple))
+	{
+		return INVALID_COLOCATION_ID;
+	}
+
+	Datum datumArray[Natts_pg_dist_partition];
+	bool isNullArray[Natts_pg_dist_partition];
+
+	Relation pgDistPartition = table_open(DistPartitionRelationId(), AccessShareLock);
+
+	TupleDesc tupleDescriptor = RelationGetDescr(pgDistPartition);
+	heap_deform_tuple(partitionTuple, tupleDescriptor, datumArray, isNullArray);
+
+	if (isNullArray[Anum_pg_dist_partition_colocationid - 1])
+	{
+		/* colocation id cannot be NULL, still let's make sure */
+		heap_freetuple(partitionTuple);
+		table_close(pgDistPartition, NoLock);
+		return INVALID_COLOCATION_ID;
+	}
+
+	Datum colocationIdDatum = datumArray[Anum_pg_dist_partition_colocationid - 1];
+	uint32 colocationId = DatumGetUInt32(colocationIdDatum);
+
+	heap_freetuple(partitionTuple);
+	table_close(pgDistPartition, NoLock);
+
+	return colocationId;
+}
+
+
 /*
  * PgDistPartitionTupleViaCatalog is a helper function that searches
  * pg_dist_partition for the given relationId. The caller is responsible

src/include/distributed/metadata_cache.h

@@ -150,6 +150,7 @@ extern char PgDistPartitionViaCatalog(Oid relationId);
 extern List * LookupDistShardTuples(Oid relationId);
 extern char PartitionMethodViaCatalog(Oid relationId);
 extern Var * PartitionColumnViaCatalog(Oid relationId);
+extern uint32 ColocationIdViaCatalog(Oid relationId);
 extern bool IsCitusLocalTableByDistParams(char partitionMethod, char replicationModel);
 extern List * CitusTableList(void);
 extern ShardInterval * LoadShardInterval(uint64 shardId);