diff --git a/src/backend/distributed/planner/distributed_planner.c b/src/backend/distributed/planner/distributed_planner.c index 3b6a8f9f7..dfce411ad 100644 --- a/src/backend/distributed/planner/distributed_planner.c +++ b/src/backend/distributed/planner/distributed_planner.c @@ -108,6 +108,7 @@ static int AssignRTEIdentities(List *rangeTableList, int rteIdCounter); static void AssignRTEIdentity(RangeTblEntry *rangeTableEntry, int rteIdentifier); static void AdjustPartitioningForDistributedPlanning(List *rangeTableList, bool setPartitionedTablesInherited); +static bool RTEWentThroughAdjustPartitioning(RangeTblEntry *rangeTableEntry); static PlannedStmt * FinalizeNonRouterPlan(PlannedStmt *localPlan, DistributedPlan *distributedPlan, CustomScan *customScan); @@ -494,6 +495,20 @@ AdjustPartitioningForDistributedPlanning(List *rangeTableList, } +/* + * RTEWentThroughAdjustPartitioning returns true if the given rangetableentry + * has been modified through AdjustPartitioningForDistributedPlanning + * function, false otherwise. + */ +static bool +RTEWentThroughAdjustPartitioning(RangeTblEntry *rangeTableEntry) +{ + return (rangeTableEntry->rtekind == RTE_RELATION && + PartitionedTable(rangeTableEntry->relid) && + rangeTableEntry->inh == false); +} + + /* * AssignRTEIdentity assigns the given rteIdentifier to the given range table * entry. @@ -1976,6 +1991,62 @@ multi_relation_restriction_hook(PlannerInfo *root, RelOptInfo *relOptInfo, } +/* + * multi_get_relation_info_hook modifies the relation's indexlist + * if necessary, to avoid a crash in PG16 caused by our + * Citus function AdjustPartitioningForDistributedPlanning(). + * + * AdjustPartitioningForDistributedPlanning() is a hack that we use + * to prevent Postgres' standard_planner() to expand all the partitions + * for the distributed planning when a distributed partitioned table + * is queried. It is required for both correctness and performance + * reasons. Although we can eliminate the use of the function for + * the correctness (e.g., make sure that rest of the planner can handle + * partitions), it's performance implication is hard to avoid. Certain + * planning logic of Citus (such as router or query pushdown) relies + * heavily on the relationRestrictionList. If + * AdjustPartitioningForDistributedPlanning() is removed, all the + * partitions show up in the relationRestrictionList, causing high + * planning times for such queries. + */ +void +multi_get_relation_info_hook(PlannerInfo *root, Oid relationObjectId, bool inhparent, + RelOptInfo *rel) +{ + if (!CitusHasBeenLoaded()) + { + return; + } + + Index varno = rel->relid; + RangeTblEntry *rangeTableEntry = planner_rt_fetch(varno, root); + + if (RTEWentThroughAdjustPartitioning(rangeTableEntry)) + { + ListCell *lc = NULL; + foreach(lc, rel->indexlist) + { + IndexOptInfo *indexOptInfo = (IndexOptInfo *) lfirst(lc); + if (get_rel_relkind(indexOptInfo->indexoid) == RELKIND_PARTITIONED_INDEX) + { + /* + * Normally, we should not need this. However, the combination of + * Postgres commit 3c569049b7b502bb4952483d19ce622ff0af5fd6 and + * Citus function AdjustPartitioningForDistributedPlanning() + * forces us to do this. The commit expects partitioned indexes + * to belong to relations with "inh" flag set properly. Whereas, the + * function overrides "inh" flag. To avoid a crash, + * we go over the list of indexinfos and remove all partitioned indexes. + * Partitioned indexes were ignored pre PG16 anyway, we are essentially + * not breaking any logic. + */ + rel->indexlist = foreach_delete_current(rel->indexlist, lc); + } + } + } +} + + /* * TranslatedVars deep copies the translated vars for the given relation index * if there is any append rel list. diff --git a/src/backend/distributed/shared_library_init.c b/src/backend/distributed/shared_library_init.c index a2b1811b8..e84aa282a 100644 --- a/src/backend/distributed/shared_library_init.c +++ b/src/backend/distributed/shared_library_init.c @@ -104,6 +104,7 @@ #include "replication/walsender.h" #include "storage/ipc.h" #include "optimizer/planner.h" +#include "optimizer/plancat.h" #include "optimizer/paths.h" #include "tcop/tcopprot.h" #include "utils/guc.h" @@ -452,6 +453,7 @@ _PG_init(void) /* register for planner hook */ set_rel_pathlist_hook = multi_relation_restriction_hook; + get_relation_info_hook = multi_get_relation_info_hook; set_join_pathlist_hook = multi_join_restriction_hook; ExecutorStart_hook = CitusExecutorStart; ExecutorRun_hook = CitusExecutorRun; diff --git a/src/include/distributed/distributed_planner.h b/src/include/distributed/distributed_planner.h index aac936a98..d46fbf2e6 100644 --- a/src/include/distributed/distributed_planner.h +++ b/src/include/distributed/distributed_planner.h @@ -234,6 +234,8 @@ extern List * TranslatedVarsForRteIdentity(int rteIdentity); extern struct DistributedPlan * GetDistributedPlan(CustomScan *node); extern void multi_relation_restriction_hook(PlannerInfo *root, RelOptInfo *relOptInfo, Index restrictionIndex, RangeTblEntry *rte); +extern void multi_get_relation_info_hook(PlannerInfo *root, Oid relationObjectId, bool + inhparent, RelOptInfo *rel); extern void multi_join_restriction_hook(PlannerInfo *root, RelOptInfo *joinrel, RelOptInfo *outerrel,