From b1a38e0bca9a51af64544cf28ed494a1c21cdfde Mon Sep 17 00:00:00 2001 From: Mehmet Yilmaz Date: Thu, 17 Apr 2025 07:42:44 +0000 Subject: [PATCH] Refactor CTE level shifting logic in insert_select_planner.c to improve clarity and functionality --- .../planner/insert_select_planner.c | 166 ++++++++++++------ 1 file changed, 116 insertions(+), 50 deletions(-) diff --git a/src/backend/distributed/planner/insert_select_planner.c b/src/backend/distributed/planner/insert_select_planner.c index f44c9938e..7e41d7c87 100644 --- a/src/backend/distributed/planner/insert_select_planner.c +++ b/src/backend/distributed/planner/insert_select_planner.c @@ -108,16 +108,25 @@ static void ProcessEntryPair(TargetEntry *insertEntry, TargetEntry *selectEntry, Form_pg_attribute attr, int targetEntryIndex, List **projectedEntries, List **nonProjectedEntries); -typedef struct ShiftAllLevelsContext +typedef struct ShiftReferencesWalkerContext { - int baseLevel; /* we shift ctelevelsup if it’s >= this number */ -} ShiftAllLevelsContext; - + /* current nesting depth, 0 for top-level query, increments each time we go one Query deeper */ + int levelsup; + + /* how much to shift ctelevelsup by if it matches levelsup+1 (typically -1) */ + int offset; + + /* optional: if we want to check ctename is in top-level list */ + List *topLevelCteList; +} ShiftReferencesWalkerContext; static bool -ShiftAllLevelsWalker(Node *node, ShiftAllLevelsContext *ctx); +inline_cte_walker(Node *node, ShiftReferencesWalkerContext *context); + static void -ShiftAllCteLevelsInQuery(Query *subquery, int baseLevel); +DecrementInsertLevelReferences(Query *subquery, + int offset, /* typically -1 */ + List *topLevelCteList); /* depth of current insert/select planner. */ @@ -550,6 +559,8 @@ PrepareInsertSelectForCitusPlanner(Query *insertSelectQuery) */ elog(DEBUG1, "Unifying top-level cteList with subquery cteList"); + List *topLevelCteListCopy = copyObject(insertSelectQuery->cteList); + selectRte->subquery->cteList = list_concat(selectRte->subquery->cteList, copyObject(insertSelectQuery->cteList)); @@ -558,75 +569,130 @@ PrepareInsertSelectForCitusPlanner(Query *insertSelectQuery) /* Suppose we physically appended the top-level cteList into the subquery, so references are at ctelevelsup=1, 2, etc. We want them all to shift by -1. */ - ShiftAllCteLevelsInQuery(selectRte->subquery, - 1 /* baseLevel => anything >= 1 is shifted */); + DecrementInsertLevelReferences(selectRte->subquery, -1, topLevelCteListCopy /* for ctename check */); elog(DEBUG1, "Done shifting ctelevelsup X->X-1 for subquery references"); } } /* - * ShiftAllLevelsWalker: - * Recursively visits each Query using QTW_EXAMINE_RTES_BEFORE so it - * receives RangeTblEntry nodes. If an RTE_CTE has ctelevelsup >= baseLevel, - * do ctelevelsup -= 1. + * inline_cte_walker + * + * If node is a Query, we increase context->levelsup by 1, + * recursively walk the Query, then restore it. + * + * If node is a RangeTblEntry with RTE_CTE and ctelevelsup == (levelsup + 1), + * we do ctelevelsup += offset (e.g. -1 => so k+1 → k). + * + * We do expression_tree_walker for fallback on expressions. */ static bool -ShiftAllLevelsWalker(Node *node, ShiftAllLevelsContext *ctx) +inline_cte_walker(Node *node, ShiftReferencesWalkerContext *context) { - if (node == NULL) - return false; + if (node == NULL) + return false; - if (IsA(node, RangeTblEntry)) - { - RangeTblEntry *rte = (RangeTblEntry *) node; + if (IsA(node, Query)) + { + Query *query = (Query *) node; - if (rte->rtekind == RTE_CTE && - rte->ctelevelsup >= ctx->baseLevel) - { - elog(DEBUG2, "Shifting ctelevelsup from %d to %d", - rte->ctelevelsup, rte->ctelevelsup - 1); - rte->ctelevelsup -= 1; - } + /* descend one query level deeper */ + context->levelsup++; - return false; - } - else if (IsA(node, Query)) - { - /* For each Query, we call query_tree_walker with QTW_EXAMINE_RTES_BEFORE */ - Query *query = (Query *) node; + /* + * Use QTW_EXAMINE_RTES_AFTER or QTW_EXAMINE_RTES_BEFORE – your snippet + * used AFTER, so let's keep that. This means we'll handle the rtable + * after the rest, which is okay. + * + * Or we can do QTW_EXAMINE_RTES_BEFORE so we see RangeTblEntry first. + * The key is being consistent with your scenario. + */ + query_tree_walker(query, + inline_cte_walker, + context, + QTW_EXAMINE_RTES_AFTER); - query_tree_walker(query, - ShiftAllLevelsWalker, - (void *) ctx, - QTW_EXAMINE_RTES_BEFORE); + context->levelsup--; - return false; - } + return false; + } + else if (IsA(node, RangeTblEntry)) + { + RangeTblEntry *rte = (RangeTblEntry *) node; - /* fallback for expressions, etc. */ - return expression_tree_walker(node, ShiftAllLevelsWalker, (void *) ctx); + if (rte->rtekind == RTE_CTE) + { + /* + * If RTE_CTE's ctelevelsup == (current nesting level + 1), + * we do ctelevelsup += offset. + * e.g. if offset=-1, and ctelevelsup= (levelsup + 1), + * that effectively does "k+1 → k". + */ + if (rte->ctelevelsup == (context->levelsup + 1)) + { + /* optionally verify ctename is in topLevelCteList */ + bool foundName = false; + if (context->topLevelCteList != NULL) + { + ListCell *lc = NULL; + foreach(lc, context->topLevelCteList) + { + CommonTableExpr *cte = (CommonTableExpr *) lfirst(lc); + if (strcmp(cte->ctename, rte->ctename) == 0) + { + foundName = true; + break; + } + } + } + else + { + foundName = true; /* if we don't need a name check */ + } + + if (foundName) + { + int oldSup = rte->ctelevelsup; + rte->ctelevelsup += context->offset; /* e.g. -1 => (k+1) → k */ + elog(DEBUG2, "Shifting ctelevelsup for ctename=%s from %d to %d", + rte->ctename, oldSup, rte->ctelevelsup); + } + else + { + elog(DEBUG2, "ctename=%s not found in top-level list, skipping shift", + rte->ctename); + } + } + } + + return false; + } + + /* fallback for expressions, e.g. scanning function calls, sublinks, etc. */ + return expression_tree_walker(node, inline_cte_walker, (void *) context); } -/* - * ShiftAllCteLevelsInQuery - * A convenience wrapper to shift ctelevelsup by 1 for any references >= baseLevel - */ + static void -ShiftAllCteLevelsInQuery(Query *subquery, int baseLevel) +DecrementInsertLevelReferences(Query *subquery, + int offset, /* typically -1 */ + List *topLevelCteList) { - ShiftAllLevelsContext ctx; - ctx.baseLevel = baseLevel; + ShiftReferencesWalkerContext ctx; + ctx.levelsup = 0; /* so that top-level query => 0, subquery => 1, etc. */ + ctx.offset = offset; + ctx.topLevelCteList = topLevelCteList; - (void) query_tree_walker(subquery, - ShiftAllLevelsWalker, - (void *) &ctx, - QTW_EXAMINE_RTES_BEFORE); + query_tree_walker(subquery, + inline_cte_walker, + (void *) &ctx, + QTW_EXAMINE_RTES_AFTER /* or BEFORE, but snippet used AFTER */); } + /* * CreateCombineQueryForRouterPlan is used for creating a dummy combineQuery * for a router plan, since router plans normally don't have one.