diff --git a/src/backend/distributed/planner/multi_logical_optimizer.c b/src/backend/distributed/planner/multi_logical_optimizer.c index c5028ba9c..00183f201 100644 --- a/src/backend/distributed/planner/multi_logical_optimizer.c +++ b/src/backend/distributed/planner/multi_logical_optimizer.c @@ -27,6 +27,7 @@ #include "commands/extension.h" #include "distributed/citus_nodes.h" #include "distributed/citus_ruleutils.h" +#include "distributed/colocation_utils.h" #include "distributed/metadata_cache.h" #include "distributed/multi_logical_optimizer.h" #include "distributed/multi_logical_planner.h" @@ -3960,10 +3961,8 @@ RelationIdList(Query *query) /* - * CoPartitionedTables checks if given two distributed tables have 1-to-1 shard - * partitioning. It uses shard interval array that are sorted on interval minimum - * values. Then it compares every shard interval in order and if any pair of - * shard intervals are not equal it returns false. + * CoPartitionedTables checks if given two distributed tables have 1-to-1 + * shard partitioning. */ static bool CoPartitionedTables(Oid firstRelationId, Oid secondRelationId) @@ -3992,6 +3991,22 @@ CoPartitionedTables(Oid firstRelationId, Oid secondRelationId) Assert(comparisonFunction != NULL); + /* + * Check if the tables have the same colocation ID - if so, we know + * they're colocated. + */ + if (firstTableCache->colocationId != INVALID_COLOCATION_ID && + firstTableCache->colocationId == secondTableCache->colocationId) + { + return true; + } + + /* + * If not known to be colocated check if the remaining shards are + * anyway. Do so by comparing the shard interval arrays that are sorted on + * interval minimum values. Then it compares every shard interval in order + * and if any pair of shard intervals are not equal it returns false. + */ for (intervalIndex = 0; intervalIndex < firstListShardCount; intervalIndex++) { ShardInterval *firstInterval = sortedFirstIntervalArray[intervalIndex];