Merge remote-tracking branch 'upstream/main' into issue/6694

src/backend/distributed/commands/foreign_constraint.c

@@ -221,7 +221,8 @@ ErrorIfUnsupportedForeignConstraintExists(Relation relation, char referencingDis
 		if (!referencedIsCitus && !selfReferencingTable)
 		{
 			if (IsCitusLocalTableByDistParams(referencingDistMethod,
-											  referencingReplicationModel))
+											  referencingReplicationModel,
+											  referencingColocationId))
 			{
 				ErrorOutForFKeyBetweenPostgresAndCitusLocalTable(referencedTableId);
 			}
@@ -245,8 +246,7 @@ ErrorIfUnsupportedForeignConstraintExists(Relation relation, char referencingDis
 		if (!selfReferencingTable)
 		{
 			referencedDistMethod = PartitionMethod(referencedTableId);
-			referencedDistKey = IsCitusTableType(referencedTableId,
+			referencedDistKey = !HasDistributionKey(referencedTableId) ?
-												 CITUS_TABLE_WITH_NO_DIST_KEY) ?
 								NULL :
 								DistPartitionKey(referencedTableId);
 			referencedColocationId = TableColocationId(referencedTableId);
@@ -278,9 +278,17 @@ ErrorIfUnsupportedForeignConstraintExists(Relation relation, char referencingDis
 		}
 
 		bool referencingIsCitusLocalOrRefTable =
-			(referencingDistMethod == DISTRIBUTE_BY_NONE);
+			IsCitusLocalTableByDistParams(referencingDistMethod,
+										  referencingReplicationModel,
+										  referencingColocationId) ||
+			IsReferenceTableByDistParams(referencingDistMethod,
+										 referencingReplicationModel);
 		bool referencedIsCitusLocalOrRefTable =
-			(referencedDistMethod == DISTRIBUTE_BY_NONE);
+			IsCitusLocalTableByDistParams(referencedDistMethod,
+										  referencedReplicationModel,
+										  referencedColocationId) ||
+			IsReferenceTableByDistParams(referencedDistMethod,
+										 referencedReplicationModel);
 		if (referencingIsCitusLocalOrRefTable && referencedIsCitusLocalOrRefTable)
 		{
 			EnsureSupportedFKeyBetweenCitusLocalAndRefTable(constraintForm,
@@ -313,7 +321,8 @@ ErrorIfUnsupportedForeignConstraintExists(Relation relation, char referencingDis
 		 * reference table is referenced.
 		 */
 		bool referencedIsReferenceTable =
-			(referencedReplicationModel == REPLICATION_MODEL_2PC);
+			IsReferenceTableByDistParams(referencedDistMethod,
+										 referencedReplicationModel);
 		if (!referencedIsReferenceTable && (
 				referencingColocationId == INVALID_COLOCATION_ID ||
 				referencingColocationId != referencedColocationId))

src/backend/distributed/commands/index.c

@@ -1190,7 +1190,7 @@ ErrorIfUnsupportedIndexStmt(IndexStmt *createIndexStatement)
 		 * Non-distributed tables do not have partition key, and unique constraints
 		 * are allowed for them. Thus, we added a short-circuit for non-distributed tables.
 		 */
-		if (IsCitusTableType(relationId, CITUS_TABLE_WITH_NO_DIST_KEY))
+		if (!HasDistributionKey(relationId))
 		{
 			return;
 		}

src/backend/distributed/commands/multi_copy.c

@@ -393,7 +393,7 @@ CitusCopyFrom(CopyStmt *copyStatement, QueryCompletion *completionTag)
 	if (IsCitusTableTypeCacheEntry(cacheEntry, HASH_DISTRIBUTED) ||
 		IsCitusTableTypeCacheEntry(cacheEntry, RANGE_DISTRIBUTED) ||
 		IsCitusTableTypeCacheEntry(cacheEntry, APPEND_DISTRIBUTED) ||
-		IsCitusTableTypeCacheEntry(cacheEntry, CITUS_TABLE_WITH_NO_DIST_KEY))
+		!HasDistributionKeyCacheEntry(cacheEntry))
 	{
 		CopyToExistingShards(copyStatement, completionTag);
 	}

src/backend/distributed/commands/table.c

@@ -75,7 +75,7 @@ static void DistributePartitionUsingParent(Oid parentRelationId,
 static void ErrorIfMultiLevelPartitioning(Oid parentRelationId, Oid partitionRelationId);
 static void ErrorIfAttachCitusTableToPgLocalTable(Oid parentRelationId,
 												  Oid partitionRelationId);
-static bool AlterTableDefinesFKeyBetweenPostgresAndNonDistTable(
+static bool ATDefinesFKeyBetweenPostgresAndCitusLocalOrRef(
 	AlterTableStmt *alterTableStatement);
 static bool ShouldMarkConnectedRelationsNotAutoConverted(Oid leftRelationId,
 														 Oid rightRelationId);
@@ -1119,7 +1119,7 @@ PreprocessAlterTableStmt(Node *node, const char *alterTableCommand,
 
 	if (ShouldEnableLocalReferenceForeignKeys() &&
 		processUtilityContext != PROCESS_UTILITY_SUBCOMMAND &&
-		AlterTableDefinesFKeyBetweenPostgresAndNonDistTable(alterTableStatement))
+		ATDefinesFKeyBetweenPostgresAndCitusLocalOrRef(alterTableStatement))
 	{
 		/*
 		 * We don't process subcommands generated by postgres.
@@ -1561,12 +1561,12 @@ PreprocessAlterTableStmt(Node *node, const char *alterTableCommand,
 
 
 /*
- * AlterTableDefinesFKeyBetweenPostgresAndNonDistTable returns true if given
+ * ATDefinesFKeyBetweenPostgresAndCitusLocalOrRef returns true if given
  * alter table command defines foreign key between a postgres table and a
  * reference or citus local table.
  */
 static bool
-AlterTableDefinesFKeyBetweenPostgresAndNonDistTable(AlterTableStmt *alterTableStatement)
+ATDefinesFKeyBetweenPostgresAndCitusLocalOrRef(AlterTableStmt *alterTableStatement)
 {
 	List *foreignKeyConstraintList =
 		GetAlterTableAddFKeyConstraintList(alterTableStatement);
@@ -1584,9 +1584,12 @@ AlterTableDefinesFKeyBetweenPostgresAndNonDistTable(AlterTableStmt *alterTableSt
 	if (!IsCitusTable(leftRelationId))
 	{
 		return RelationIdListContainsCitusTableType(rightRelationIdList,
-													CITUS_TABLE_WITH_NO_DIST_KEY);
+													CITUS_LOCAL_TABLE) ||
+			   RelationIdListContainsCitusTableType(rightRelationIdList,
+													REFERENCE_TABLE);
 	}
-	else if (IsCitusTableType(leftRelationId, CITUS_TABLE_WITH_NO_DIST_KEY))
+	else if (IsCitusTableType(leftRelationId, CITUS_LOCAL_TABLE) ||
+			 IsCitusTableType(leftRelationId, REFERENCE_TABLE))
 	{
 		return RelationIdListContainsPostgresTable(rightRelationIdList);
 	}
@@ -3626,7 +3629,7 @@ SetupExecutionModeForAlterTable(Oid relationId, AlterTableCmd *command)
 	 * sequential mode.
 	 */
 	if (executeSequentially &&
-		!IsCitusTableType(relationId, CITUS_TABLE_WITH_NO_DIST_KEY) &&
+		HasDistributionKey(relationId) &&
 		ParallelQueryExecutedInTransaction())
 	{
 		char *relationName = get_rel_name(relationId);

src/backend/distributed/commands/truncate.c

@@ -324,7 +324,7 @@ ExecuteTruncateStmtSequentialIfNecessary(TruncateStmt *command)
 	{
 		Oid relationId = RangeVarGetRelid(rangeVar, NoLock, failOK);
 
-		if (IsCitusTableType(relationId, CITUS_TABLE_WITH_NO_DIST_KEY) &&
+		if (IsCitusTable(relationId) && !HasDistributionKey(relationId) &&
 			TableReferenced(relationId))
 		{
 			char *relationName = get_rel_name(relationId);

src/backend/distributed/metadata/metadata_cache.c

@@ -311,7 +311,7 @@ static void InvalidateDistTableCache(void);
 static void InvalidateDistObjectCache(void);
 static bool InitializeTableCacheEntry(int64 shardId, bool missingOk);
 static bool IsCitusTableTypeInternal(char partitionMethod, char replicationModel,
-									 CitusTableType tableType);
+									 uint32 colocationId, CitusTableType tableType);
 static bool RefreshTableCacheEntryIfInvalid(ShardIdCacheEntry *shardEntry, bool
 											missingOk);
 
@@ -450,7 +450,36 @@ bool
 IsCitusTableTypeCacheEntry(CitusTableCacheEntry *tableEntry, CitusTableType tableType)
 {
 	return IsCitusTableTypeInternal(tableEntry->partitionMethod,
-									tableEntry->replicationModel, tableType);
+									tableEntry->replicationModel,
+									tableEntry->colocationId, tableType);
+}
+
+
+/*
+ * HasDistributionKey returs true if given Citus table doesn't have a
+ * distribution key.
+ */
+bool
+HasDistributionKey(Oid relationId)
+{
+	CitusTableCacheEntry *tableEntry = LookupCitusTableCacheEntry(relationId);
+	if (tableEntry == NULL)
+	{
+		ereport(ERROR, (errmsg("relation with oid %u is not a Citus table", relationId)));
+	}
+
+	return HasDistributionKeyCacheEntry(tableEntry);
+}
+
+
+/*
+ * HasDistributionKey returs true if given cache entry identifies a Citus
+ * table that doesn't have a distribution key.
+ */
+bool
+HasDistributionKeyCacheEntry(CitusTableCacheEntry *tableEntry)
+{
+	return tableEntry->partitionMethod != DISTRIBUTE_BY_NONE;
 }
 
 
@@ -460,7 +489,7 @@ IsCitusTableTypeCacheEntry(CitusTableCacheEntry *tableEntry, CitusTableType tabl
  */
 static bool
 IsCitusTableTypeInternal(char partitionMethod, char replicationModel,
-						 CitusTableType tableType)
+						 uint32 colocationId, CitusTableType tableType)
 {
 	switch (tableType)
 	{
@@ -501,12 +530,8 @@ IsCitusTableTypeInternal(char partitionMethod, char replicationModel,
 		case CITUS_LOCAL_TABLE:
 		{
 			return partitionMethod == DISTRIBUTE_BY_NONE &&
-				   replicationModel != REPLICATION_MODEL_2PC;
+				   replicationModel != REPLICATION_MODEL_2PC &&
-		}
+				   colocationId == INVALID_COLOCATION_ID;
-
-		case CITUS_TABLE_WITH_NO_DIST_KEY:
-		{
-			return partitionMethod == DISTRIBUTE_BY_NONE;
 		}
 
 		case ANY_CITUS_TABLE_TYPE:
@@ -529,33 +554,21 @@ IsCitusTableTypeInternal(char partitionMethod, char replicationModel,
 char *
 GetTableTypeName(Oid tableId)
 {
-	bool regularTable = false;
+	if (!IsCitusTable(tableId))
-	char partitionMethod = ' ';
-	char replicationModel = ' ';
-	if (IsCitusTable(tableId))
-	{
-		CitusTableCacheEntry *referencingCacheEntry = GetCitusTableCacheEntry(tableId);
-		partitionMethod = referencingCacheEntry->partitionMethod;
-		replicationModel = referencingCacheEntry->replicationModel;
-	}
-	else
-	{
-		regularTable = true;
-	}
-
-	if (regularTable)
 	{
 		return "regular table";
 	}
-	else if (partitionMethod == 'h')
+
+	CitusTableCacheEntry *tableCacheEntry = GetCitusTableCacheEntry(tableId);
+	if (IsCitusTableTypeCacheEntry(tableCacheEntry, HASH_DISTRIBUTED))
 	{
 		return "distributed table";
 	}
-	else if (partitionMethod == 'n' && replicationModel == 't')
+	else if (IsCitusTableTypeCacheEntry(tableCacheEntry, REFERENCE_TABLE))
 	{
 		return "reference table";
 	}
-	else if (partitionMethod == 'n' && replicationModel != 't')
+	else if (IsCitusTableTypeCacheEntry(tableCacheEntry, CITUS_LOCAL_TABLE))
 	{
 		return "citus local table";
 	}
@@ -765,14 +778,28 @@ PgDistPartitionTupleViaCatalog(Oid relationId)
 
 
 /*
- * IsCitusLocalTableByDistParams returns true if given partitionMethod and
+ * IsReferenceTableByDistParams returns true if given partitionMethod and
- * replicationModel would identify a citus local table.
+ * replicationModel would identify a reference table.
  */
 bool
-IsCitusLocalTableByDistParams(char partitionMethod, char replicationModel)
+IsReferenceTableByDistParams(char partitionMethod, char replicationModel)
 {
 	return partitionMethod == DISTRIBUTE_BY_NONE &&
-		   replicationModel != REPLICATION_MODEL_2PC;
+		   replicationModel == REPLICATION_MODEL_2PC;
+}
+
+
+/*
+ * IsCitusLocalTableByDistParams returns true if given partitionMethod,
+ * replicationModel and colocationId would identify a citus local table.
+ */
+bool
+IsCitusLocalTableByDistParams(char partitionMethod, char replicationModel,
+							  uint32 colocationId)
+{
+	return partitionMethod == DISTRIBUTE_BY_NONE &&
+		   replicationModel != REPLICATION_MODEL_2PC &&
+		   colocationId == INVALID_COLOCATION_ID;
 }
 
 
@@ -4837,11 +4864,14 @@ CitusTableTypeIdList(CitusTableType citusTableType)
 
 		Datum partMethodDatum = datumArray[Anum_pg_dist_partition_partmethod - 1];
 		Datum replicationModelDatum = datumArray[Anum_pg_dist_partition_repmodel - 1];
+		Datum colocationIdDatum = datumArray[Anum_pg_dist_partition_colocationid - 1];
 
 		Oid partitionMethod = DatumGetChar(partMethodDatum);
 		Oid replicationModel = DatumGetChar(replicationModelDatum);
+		uint32 colocationId = DatumGetUInt32(colocationIdDatum);
 
-		if (IsCitusTableTypeInternal(partitionMethod, replicationModel, citusTableType))
+		if (IsCitusTableTypeInternal(partitionMethod, replicationModel, colocationId,
+									 citusTableType))
 		{
 			Datum relationIdDatum = datumArray[Anum_pg_dist_partition_logicalrelid - 1];
 

src/backend/distributed/metadata/metadata_sync.c

@@ -535,7 +535,7 @@ ShouldSyncTableMetadata(Oid relationId)
 
 	bool hashDistributed = IsCitusTableTypeCacheEntry(tableEntry, HASH_DISTRIBUTED);
 	bool citusTableWithNoDistKey =
-		IsCitusTableTypeCacheEntry(tableEntry, CITUS_TABLE_WITH_NO_DIST_KEY);
+		!HasDistributionKeyCacheEntry(tableEntry);
 
 	return ShouldSyncTableMetadataInternal(hashDistributed, citusTableWithNoDistKey);
 }
@@ -1158,7 +1158,7 @@ DistributionCreateCommand(CitusTableCacheEntry *cacheEntry)
 	char replicationModel = cacheEntry->replicationModel;
 	StringInfo tablePartitionKeyNameString = makeStringInfo();
 
-	if (IsCitusTableTypeCacheEntry(cacheEntry, CITUS_TABLE_WITH_NO_DIST_KEY))
+	if (!HasDistributionKeyCacheEntry(cacheEntry))
 	{
 		appendStringInfo(tablePartitionKeyNameString, "NULL");
 	}

src/backend/distributed/metadata/node_metadata.c

@@ -1536,7 +1536,7 @@ get_shard_id_for_distribution_column(PG_FUNCTION_ARGS)
 						errmsg("relation is not distributed")));
 	}
 
-	if (IsCitusTableType(relationId, CITUS_TABLE_WITH_NO_DIST_KEY))
+	if (!HasDistributionKey(relationId))
 	{
 		List *shardIntervalList = LoadShardIntervalList(relationId);
 		if (shardIntervalList == NIL)

src/backend/distributed/planner/multi_join_order.c

@@ -1383,7 +1383,7 @@ DistPartitionKey(Oid relationId)
 	CitusTableCacheEntry *partitionEntry = GetCitusTableCacheEntry(relationId);
 
 	/* non-distributed tables do not have partition column */
-	if (IsCitusTableTypeCacheEntry(partitionEntry, CITUS_TABLE_WITH_NO_DIST_KEY))
+	if (!HasDistributionKeyCacheEntry(partitionEntry))
 	{
 		return NULL;
 	}

src/backend/distributed/planner/multi_logical_planner.c

@@ -228,7 +228,7 @@ TargetListOnPartitionColumn(Query *query, List *targetEntryList)
 		 * If the expression belongs to a non-distributed table continue searching for
 		 * other partition keys.
 		 */
-		if (IsCitusTableType(relationId, CITUS_TABLE_WITH_NO_DIST_KEY))
+		if (IsCitusTable(relationId) && !HasDistributionKey(relationId))
 		{
 			continue;
 		}

src/backend/distributed/planner/multi_physical_planner.c

@@ -2199,7 +2199,7 @@ QueryPushdownSqlTaskList(Query *query, uint64 jobId,
 		Oid relationId = relationRestriction->relationId;
 
 		CitusTableCacheEntry *cacheEntry = GetCitusTableCacheEntry(relationId);
-		if (IsCitusTableTypeCacheEntry(cacheEntry, CITUS_TABLE_WITH_NO_DIST_KEY))
+		if (!HasDistributionKeyCacheEntry(cacheEntry))
 		{
 			continue;
 		}
@@ -2377,7 +2377,7 @@ ErrorIfUnsupportedShardDistribution(Query *query)
 			nonReferenceRelations = lappend_oid(nonReferenceRelations,
 												relationId);
 		}
-		else if (IsCitusTableType(relationId, CITUS_TABLE_WITH_NO_DIST_KEY))
+		else if (IsCitusTable(relationId) && !HasDistributionKey(relationId))
 		{
 			/* do not need to handle non-distributed tables */
 			continue;
@@ -2482,7 +2482,7 @@ QueryPushdownTaskCreate(Query *originalQuery, int shardIndex,
 		ShardInterval *shardInterval = NULL;
 
 		CitusTableCacheEntry *cacheEntry = GetCitusTableCacheEntry(relationId);
-		if (IsCitusTableTypeCacheEntry(cacheEntry, CITUS_TABLE_WITH_NO_DIST_KEY))
+		if (!HasDistributionKeyCacheEntry(cacheEntry))
 		{
 			/* non-distributed tables have only one shard */
 			shardInterval = cacheEntry->sortedShardIntervalArray[0];
@@ -3697,7 +3697,7 @@ PartitionedOnColumn(Var *column, List *rangeTableList, List *dependentJobList)
 		Var *partitionColumn = PartitionColumn(relationId, rangeTableId);
 
 		/* non-distributed tables do not have partition columns */
-		if (IsCitusTableType(relationId, CITUS_TABLE_WITH_NO_DIST_KEY))
+		if (IsCitusTable(relationId) && !HasDistributionKey(relationId))
 		{
 			return false;
 		}

src/backend/distributed/planner/multi_router_planner.c

@@ -2675,7 +2675,7 @@ TargetShardIntervalForFastPathQuery(Query *query, bool *isMultiShardQuery,
 {
 	Oid relationId = ExtractFirstCitusTableId(query);
 
-	if (IsCitusTableType(relationId, CITUS_TABLE_WITH_NO_DIST_KEY))
+	if (!HasDistributionKey(relationId))
 	{
 		/* we don't need to do shard pruning for non-distributed tables */
 		return list_make1(LoadShardIntervalList(relationId));
@@ -2968,7 +2968,7 @@ BuildRoutesForInsert(Query *query, DeferredErrorMessage **planningError)
 	Assert(query->commandType == CMD_INSERT);
 
 	/* reference tables and citus local tables can only have one shard */
-	if (IsCitusTableTypeCacheEntry(cacheEntry, CITUS_TABLE_WITH_NO_DIST_KEY))
+	if (!HasDistributionKeyCacheEntry(cacheEntry))
 	{
 		List *shardIntervalList = LoadShardIntervalList(distributedTableId);
 
@@ -3509,7 +3509,7 @@ ExtractInsertPartitionKeyValue(Query *query)
 	uint32 rangeTableId = 1;
 	Const *singlePartitionValueConst = NULL;
 
-	if (IsCitusTableType(distributedTableId, CITUS_TABLE_WITH_NO_DIST_KEY))
+	if (!HasDistributionKey(distributedTableId))
 	{
 		return NULL;
 	}
@@ -3829,8 +3829,7 @@ ErrorIfQueryHasUnroutableModifyingCTE(Query *queryTree)
 			CitusTableCacheEntry *modificationTableCacheEntry =
 				GetCitusTableCacheEntry(distributedTableId);
 
-			if (IsCitusTableTypeCacheEntry(modificationTableCacheEntry,
+			if (!HasDistributionKeyCacheEntry(modificationTableCacheEntry))
-										   CITUS_TABLE_WITH_NO_DIST_KEY))
 			{
 				return DeferredError(ERRCODE_FEATURE_NOT_SUPPORTED,
 									 "cannot router plan modification of a non-distributed table",

src/backend/distributed/planner/query_colocation_checker.c

@@ -168,7 +168,7 @@ AnchorRte(Query *subquery)
 		{
 			Oid relationId = currentRte->relid;
 
-			if (IsCitusTableType(relationId, CITUS_TABLE_WITH_NO_DIST_KEY))
+			if (IsCitusTable(relationId) && !HasDistributionKey(relationId))
 			{
 				/*
 				 * Non-distributed tables should not be the anchor rte since they

src/backend/distributed/planner/relation_restriction_equivalence.c

@@ -703,8 +703,8 @@ EquivalenceListContainsRelationsEquality(List *attributeEquivalenceList,
 		int rteIdentity = GetRTEIdentity(relationRestriction->rte);
 
 		/* we shouldn't check for the equality of non-distributed tables */
-		if (IsCitusTableType(relationRestriction->relationId,
+		if (IsCitusTable(relationRestriction->relationId) &&
-							 CITUS_TABLE_WITH_NO_DIST_KEY))
+			!HasDistributionKey(relationRestriction->relationId))
 		{
 			continue;
 		}
@@ -1933,7 +1933,7 @@ AllRelationsInRestrictionContextColocated(RelationRestrictionContext *restrictio
 	{
 		Oid relationId = relationRestriction->relationId;
 
-		if (IsCitusTableType(relationId, CITUS_TABLE_WITH_NO_DIST_KEY))
+		if (IsCitusTable(relationId) && !HasDistributionKey(relationId))
 		{
 			continue;
 		}

src/backend/distributed/planner/shard_pruning.c

@@ -333,7 +333,7 @@ PruneShards(Oid relationId, Index rangeTableId, List *whereClauseList,
 	}
 
 	/* short circuit for non-distributed tables such as reference table */
-	if (IsCitusTableTypeCacheEntry(cacheEntry, CITUS_TABLE_WITH_NO_DIST_KEY))
+	if (!HasDistributionKeyCacheEntry(cacheEntry))
 	{
 		prunedList = ShardArrayToList(cacheEntry->sortedShardIntervalArray,
 									  cacheEntry->shardIntervalArrayLength);

src/backend/distributed/transaction/relation_access_tracking.c

@@ -195,7 +195,7 @@ RecordRelationAccessIfNonDistTable(Oid relationId, ShardPlacementAccessType acce
 	 * recursively calling RecordRelationAccessBase(), so becareful about
 	 * removing this check.
 	 */
-	if (!IsCitusTableType(relationId, CITUS_TABLE_WITH_NO_DIST_KEY))
+	if (IsCitusTable(relationId) && HasDistributionKey(relationId))
 	{
 		return;
 	}
@@ -732,8 +732,8 @@ CheckConflictingRelationAccesses(Oid relationId, ShardPlacementAccessType access
 
 	CitusTableCacheEntry *cacheEntry = GetCitusTableCacheEntry(relationId);
 
-	if (!(IsCitusTableTypeCacheEntry(cacheEntry, CITUS_TABLE_WITH_NO_DIST_KEY) &&
+	if (HasDistributionKeyCacheEntry(cacheEntry) ||
-		  cacheEntry->referencingRelationsViaForeignKey != NIL))
+		cacheEntry->referencingRelationsViaForeignKey == NIL)
 	{
 		return;
 	}
@@ -931,7 +931,7 @@ HoldsConflictingLockWithReferencedRelations(Oid relationId, ShardPlacementAccess
 		 * We're only interested in foreign keys to reference tables and citus
 		 * local tables.
 		 */
-		if (!IsCitusTableType(referencedRelation, CITUS_TABLE_WITH_NO_DIST_KEY))
+		if (IsCitusTable(referencedRelation) && HasDistributionKey(referencedRelation))
 		{
 			continue;
 		}
@@ -993,7 +993,7 @@ HoldsConflictingLockWithReferencingRelations(Oid relationId, ShardPlacementAcces
 	CitusTableCacheEntry *cacheEntry = GetCitusTableCacheEntry(relationId);
 	bool holdsConflictingLocks = false;
 
-	Assert(IsCitusTableTypeCacheEntry(cacheEntry, CITUS_TABLE_WITH_NO_DIST_KEY));
+	Assert(!HasDistributionKeyCacheEntry(cacheEntry));
 
 	Oid referencingRelation = InvalidOid;
 	foreach_oid(referencingRelation, cacheEntry->referencingRelationsViaForeignKey)

src/backend/distributed/utils/colocation_utils.c

@@ -442,8 +442,7 @@ ShardsIntervalsEqual(ShardInterval *leftShardInterval, ShardInterval *rightShard
 	{
 		return HashPartitionedShardIntervalsEqual(leftShardInterval, rightShardInterval);
 	}
-	else if (IsCitusTableType(leftShardInterval->relationId,
+	else if (!HasDistributionKey(leftShardInterval->relationId))
-							  CITUS_TABLE_WITH_NO_DIST_KEY))
 	{
 		/*
 		 * Reference tables has only a single shard and all reference tables

src/backend/distributed/utils/shardinterval_utils.c

@@ -223,8 +223,7 @@ ShardIndex(ShardInterval *shardInterval)
 	 * currently it is not required.
 	 */
 	if (!IsCitusTableTypeCacheEntry(cacheEntry, HASH_DISTRIBUTED) &&
-		!IsCitusTableTypeCacheEntry(
+		HasDistributionKeyCacheEntry(cacheEntry))
-			cacheEntry, CITUS_TABLE_WITH_NO_DIST_KEY))
 	{
 		ereport(ERROR, (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 						errmsg("finding index of a given shard is only supported for "
@@ -233,7 +232,7 @@ ShardIndex(ShardInterval *shardInterval)
 	}
 
 	/* short-circuit for reference tables */
-	if (IsCitusTableTypeCacheEntry(cacheEntry, CITUS_TABLE_WITH_NO_DIST_KEY))
+	if (!HasDistributionKeyCacheEntry(cacheEntry))
 	{
 		/*
 		 * Reference tables and citus local tables have only a single shard,
@@ -333,7 +332,7 @@ FindShardIntervalIndex(Datum searchedValue, CitusTableCacheEntry *cacheEntry)
 			shardIndex = CalculateUniformHashRangeIndex(hashedValue, shardCount);
 		}
 	}
-	else if (IsCitusTableTypeCacheEntry(cacheEntry, CITUS_TABLE_WITH_NO_DIST_KEY))
+	else if (!HasDistributionKeyCacheEntry(cacheEntry))
 	{
 		/* non-distributed tables have a single shard, all values mapped to that shard */
 		Assert(shardCount == 1);

src/include/distributed/metadata_cache.h

@@ -133,9 +133,6 @@ typedef enum
 	REFERENCE_TABLE,
 	CITUS_LOCAL_TABLE,
 
-	/* table without a dist key such as reference table */
-	CITUS_TABLE_WITH_NO_DIST_KEY,
-
 	ANY_CITUS_TABLE_TYPE
 } CitusTableType;
 
@@ -143,6 +140,8 @@ extern List * AllCitusTableIds(void);
 extern bool IsCitusTableType(Oid relationId, CitusTableType tableType);
 extern bool IsCitusTableTypeCacheEntry(CitusTableCacheEntry *tableEtnry,
 									   CitusTableType tableType);
+bool HasDistributionKey(Oid relationId);
+bool HasDistributionKeyCacheEntry(CitusTableCacheEntry *tableEntry);
 extern char * GetTableTypeName(Oid tableId);
 
 extern void SetCreateCitusTransactionLevel(int val);
@@ -154,7 +153,9 @@ 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);
+bool IsReferenceTableByDistParams(char partitionMethod, char replicationModel);
+extern bool IsCitusLocalTableByDistParams(char partitionMethod, char replicationModel,
+										  uint32 colocationId);
 extern List * CitusTableList(void);
 extern ShardInterval * LoadShardInterval(uint64 shardId);
 extern bool ShardExists(uint64 shardId);