From 1b2e931dc4b883b5fd9167f139fdeade8d2b18d6 Mon Sep 17 00:00:00 2001 From: Hannu Krosing Date: Wed, 16 Sep 2026 13:41:19 +0000 Subject: [PATCH v3 4/5] Support dynamic pruning in recursive UNION DISTINCT ON Allows WITH RECURSIVE queries to perform state-space pruning using: UNION DISTINCT ON (keys ORDER BY sort_keys) During recursive iteration, tuples with worse sort keys are pruned, while tuples with better sort keys replace existing hashtable entries and are re-queued into the intermediate work table to re-expand exploration. Final results are buffered in a tuplestore and returned upon convergence, enabling shortest-path and graph-search algorithms in pure SQL. --- doc/src/sgml/ref/select.sgml | 4 +- src/backend/executor/nodeRecursiveunion.c | 162 +++++++++++++++++++++- src/backend/optimizer/plan/createplan.c | 42 ++++++ src/backend/optimizer/prep/prepunion.c | 24 ++++ src/backend/optimizer/util/pathnode.c | 2 + src/include/nodes/execnodes.h | 5 + src/include/nodes/pathnodes.h | 1 + src/include/nodes/plannodes.h | 7 + src/include/optimizer/pathnode.h | 1 + src/test/regress/expected/union.out | 19 +++ src/test/regress/sql/union.sql | 13 ++ 11 files changed, 275 insertions(+), 5 deletions(-) diff --git a/doc/src/sgml/ref/select.sgml b/doc/src/sgml/ref/select.sgml index ca09f75c027..b467bdaa49c 100644 --- a/doc/src/sgml/ref/select.sgml +++ b/doc/src/sgml/ref/select.sgml @@ -1287,7 +1287,9 @@ SELECT DISTINCT ON (location ORDER BY time DESC) location, time, report If the optional ORDER BY clause is specified inside DISTINCT ON, it determines which row is kept from each set of duplicates (the first row according to this sort order). Otherwise, the kept - row is unpredictable. + row is unpredictable. When used in recursive CTEs (WITH RECURSIVE), + UNION DISTINCT ON with ORDER BY can be used + to prune the search space and find optimal solutions (e.g., shortest paths). diff --git a/src/backend/executor/nodeRecursiveunion.c b/src/backend/executor/nodeRecursiveunion.c index 7166397e59b..d150aed56c7 100644 --- a/src/backend/executor/nodeRecursiveunion.c +++ b/src/backend/executor/nodeRecursiveunion.c @@ -23,6 +23,7 @@ #include "miscadmin.h" #include "utils/memutils.h" #include "utils/tuplestore.h" +#include "utils/sortsupport.h" @@ -77,6 +78,23 @@ build_hash_table(RecursiveUnionState *rustate) * 2.6 go back to 2.2 * ---------------------------------------------------------------- */ +static void +populate_result_table_from_hash(RecursiveUnionState *rustate) +{ + TupleHashTable hashtable = rustate->hashtable; + tuplehash_iterator iter; + TupleHashEntry entry; + TupleTableSlot *slot = rustate->ps.ps_ResultTupleSlot; + + tuplehash_start_iterate(hashtable->hashtab, &iter); + while ((entry = tuplehash_iterate(hashtable->hashtab, &iter)) != NULL) + { + ExecStoreMinimalTuple(entry->firstTuple, slot, false); + tuplestore_puttupleslot(rustate->result_table, slot); + ExecClearTuple(slot); + } +} + static TupleTableSlot * ExecRecursiveUnion(PlanState *pstate) { @@ -89,6 +107,100 @@ ExecRecursiveUnion(PlanState *pstate) CHECK_FOR_INTERRUPTS(); + /* If we need sorting, we must buffer and return from result_table */ + if (plan->numSortCols > 0) + { + if (node->result_table == NULL) + { + /* Run the entire recursion loop and buffer results */ + node->result_table = tuplestore_begin_heap(false, false, work_mem); + + /* 1. Process non-recursive term */ + for (;;) + { + slot = ExecProcNode(outerPlan); + if (TupIsNull(slot)) + break; + + if (plan->numCols > 0) + { + TupleHashEntry entry; + entry = LookupTupleHashEntry(node->hashtable, slot, &isnew, NULL); + if (!isnew) + { + ReplaceTupleHashEntryIfBetter(node->hashtable, + entry, + slot, + node->sort_firstTupleSlot, + node->sortKeys, + plan->numSortCols); + continue; + } + } + tuplestore_puttupleslot(node->working_table, slot); + } + + /* 2. Process recursive term */ + node->recursing = true; + for (;;) + { + slot = ExecProcNode(innerPlan); + if (TupIsNull(slot)) + { + Tuplestorestate *swaptemp; + + if (node->intermediate_empty) + break; /* End of recursion */ + + tuplestore_clear(node->working_table); + swaptemp = node->working_table; + node->working_table = node->intermediate_table; + node->intermediate_table = swaptemp; + node->intermediate_empty = true; + innerPlan->chgParam = bms_add_member(innerPlan->chgParam, + plan->wtParam); + continue; + } + + if (plan->numCols > 0) + { + TupleHashEntry entry; + entry = LookupTupleHashEntry(node->hashtable, slot, &isnew, NULL); + if (!isnew) + { + bool replaced = ReplaceTupleHashEntryIfBetter(node->hashtable, + entry, + slot, + node->sort_firstTupleSlot, + node->sortKeys, + plan->numSortCols); + if (replaced) + { + /* Replaced! Explore this better path */ + node->intermediate_empty = false; + tuplestore_puttupleslot(node->intermediate_table, slot); + } + continue; + } + } + + node->intermediate_empty = false; + tuplestore_puttupleslot(node->intermediate_table, slot); + } + + /* Populate result_table from hashtable */ + populate_result_table_from_hash(node); + } + + /* Read from result_table */ + slot = node->ps.ps_ResultTupleSlot; + if (tuplestore_gettupleslot(node->result_table, true, false, slot)) + return slot; + + return NULL; + } + + /* Original pipelined behavior (numSortCols == 0) */ /* 1. Evaluate non-recursive term */ if (!node->recursing) { @@ -198,6 +310,9 @@ ExecInitRecursiveUnion(RecursiveUnion *node, EState *estate, int eflags) rustate->hashtable = NULL; rustate->tempContext = NULL; rustate->tuplesContext = NULL; + rustate->result_table = NULL; + rustate->sortKeys = NULL; + rustate->sort_firstTupleSlot = NULL; /* initialize processing state */ rustate->recursing = false; @@ -218,10 +333,20 @@ ExecInitRecursiveUnion(RecursiveUnion *node, EState *estate, int eflags) AllocSetContextCreate(CurrentMemoryContext, "RecursiveUnion", ALLOCSET_DEFAULT_SIZES); - rustate->tuplesContext = - BumpContextCreate(CurrentMemoryContext, - "RecursiveUnion hashed tuples", - ALLOCSET_DEFAULT_SIZES); + if (node->numSortCols > 0) + { + rustate->tuplesContext = + AllocSetContextCreate(CurrentMemoryContext, + "RecursiveUnion hashed tuples", + ALLOCSET_DEFAULT_SIZES); + } + else + { + rustate->tuplesContext = + BumpContextCreate(CurrentMemoryContext, + "RecursiveUnion hashed tuples", + ALLOCSET_DEFAULT_SIZES); + } } /* @@ -246,6 +371,8 @@ ExecInitRecursiveUnion(RecursiveUnion *node, EState *estate, int eflags) * tuples, so we have to initialize them. */ ExecInitResultTypeTL(&rustate->ps); + if (node->numSortCols > 0) + ExecInitResultSlot(&rustate->ps, &TTSOpsMinimalTuple); /* * Initialize result tuple type. (Note: we have to set up the result type @@ -273,6 +400,26 @@ ExecInitRecursiveUnion(RecursiveUnion *node, EState *estate, int eflags) build_hash_table(rustate); } + if (node->numSortCols > 0) + { + TupleDesc desc = ExecGetResultType(outerPlanState(rustate)); + int i; + + rustate->sortKeys = (SortSupportData *) palloc0(node->numSortCols * sizeof(SortSupportData)); + rustate->sort_firstTupleSlot = ExecInitExtraTupleSlot(estate, desc, &TTSOpsMinimalTuple); + + for (i = 0; i < node->numSortCols; i++) + { + SortSupport skey = &rustate->sortKeys[i]; + + skey->ssup_cxt = CurrentMemoryContext; + skey->ssup_collation = node->sortCollations[i]; + skey->ssup_nulls_first = node->sortNullsFirst[i]; + skey->ssup_attno = node->sortColIdx[i]; + PrepareSortSupportFromOrderingOp(node->sortOperators[i], skey); + } + } + return rustate; } @@ -294,6 +441,8 @@ ExecEndRecursiveUnion(RecursiveUnionState *node) MemoryContextDelete(node->tempContext); if (node->tuplesContext) MemoryContextDelete(node->tuplesContext); + if (node->result_table) + tuplestore_end(node->result_table); /* * close down subplans @@ -338,4 +487,9 @@ ExecReScanRecursiveUnion(RecursiveUnionState *node) node->intermediate_empty = true; tuplestore_clear(node->working_table); tuplestore_clear(node->intermediate_table); + if (node->result_table) + { + tuplestore_end(node->result_table); + node->result_table = NULL; + } } diff --git a/src/backend/optimizer/plan/createplan.c b/src/backend/optimizer/plan/createplan.c index 78a6d6a3c20..59fdcc4a82b 100644 --- a/src/backend/optimizer/plan/createplan.c +++ b/src/backend/optimizer/plan/createplan.c @@ -225,6 +225,7 @@ static RecursiveUnion *make_recursive_union(List *tlist, Plan *righttree, int wtParam, List *distinctList, + List *distinctSortClause, Cardinality numGroups); static BitmapAnd *make_bitmap_and(List *bitmapplans); static BitmapOr *make_bitmap_or(List *bitmapplans); @@ -2667,6 +2668,7 @@ create_recursiveunion_plan(PlannerInfo *root, RecursiveUnionPath *best_path) rightplan, best_path->wtParam, best_path->distinctList, + best_path->distinctSortClause, best_path->numGroups); copy_generic_path_info(&plan->plan, (Path *) best_path); @@ -5949,6 +5951,7 @@ make_recursive_union(List *tlist, Plan *righttree, int wtParam, List *distinctList, + List *distinctSortClause, Cardinality numGroups) { RecursiveUnion *node = makeNode(RecursiveUnion); @@ -5994,6 +5997,45 @@ make_recursive_union(List *tlist, node->dupOperators = dupOperators; node->dupCollations = dupCollations; } + + /* Extract sort keys for inline DISTINCT ON ORDER BY */ + if (distinctSortClause) + { + int numsortkeys = list_length(distinctSortClause); + AttrNumber *sortColIdx = (AttrNumber *) palloc(numsortkeys * sizeof(AttrNumber)); + Oid *sortOperators = (Oid *) palloc(numsortkeys * sizeof(Oid)); + Oid *collations = (Oid *) palloc(numsortkeys * sizeof(Oid)); + bool *nullsFirst = (bool *) palloc(numsortkeys * sizeof(bool)); + int keyno = 0; + ListCell *l; + + foreach(l, distinctSortClause) + { + SortGroupClause *sortcl = (SortGroupClause *) lfirst(l); + TargetEntry *tle = get_sortgroupclause_tle(sortcl, plan->targetlist); + + sortColIdx[keyno] = tle->resno; + sortOperators[keyno] = sortcl->sortop; + collations[keyno] = exprCollation((Node *) tle->expr); + nullsFirst[keyno] = sortcl->nulls_first; + keyno++; + } + + node->numSortCols = numsortkeys; + node->sortColIdx = sortColIdx; + node->sortOperators = sortOperators; + node->sortCollations = collations; + node->sortNullsFirst = nullsFirst; + } + else + { + node->numSortCols = 0; + node->sortColIdx = NULL; + node->sortOperators = NULL; + node->sortCollations = NULL; + node->sortNullsFirst = NULL; + } + node->numGroups = numGroups; return node; diff --git a/src/backend/optimizer/prep/prepunion.c b/src/backend/optimizer/prep/prepunion.c index aa65bb6c6fa..4f20ffefb06 100644 --- a/src/backend/optimizer/prep/prepunion.c +++ b/src/backend/optimizer/prep/prepunion.c @@ -36,6 +36,7 @@ #include "optimizer/planner.h" #include "optimizer/prep.h" #include "optimizer/tlist.h" +#include "optimizer/optimizer.h" #include "parser/parse_coerce.h" #include "port/pg_bitutils.h" #include "utils/selfuncs.h" @@ -360,6 +361,28 @@ recurse_set_operations(Node *setOp, PlannerInfo *root, /* * Generate paths for a recursive UNION node */ +static List * +adjust_setop_sortclauses(List *sortClauses, List *query_tlist) +{ + List *result = NIL; + ListCell *lc; + + if (sortClauses == NIL) + return NIL; + + foreach(lc, sortClauses) + { + SortGroupClause *sortcl = (SortGroupClause *) lfirst(lc); + SortGroupClause *newcl = copyObject(sortcl); + TargetEntry *tle; + + tle = get_sortgroupref_tle(sortcl->tleSortGroupRef, query_tlist); + newcl->tleSortGroupRef = tle->resno; + result = lappend(result, newcl); + } + return result; +} + static RelOptInfo * generate_recursion_path(SetOperationStmt *setOp, PlannerInfo *root, List *refnames_tlist, @@ -463,6 +486,7 @@ generate_recursion_path(SetOperationStmt *setOp, PlannerInfo *root, rpath, result_rel->reltarget, groupList, + adjust_setop_sortclauses(setOp->sortClauses, root->parse->targetList), root->wt_param_id, dNumGroups); diff --git a/src/backend/optimizer/util/pathnode.c b/src/backend/optimizer/util/pathnode.c index fd2dd6808b0..341629b87ad 100644 --- a/src/backend/optimizer/util/pathnode.c +++ b/src/backend/optimizer/util/pathnode.c @@ -3630,6 +3630,7 @@ create_recursiveunion_path(PlannerInfo *root, Path *rightpath, PathTarget *target, List *distinctList, + List *distinctSortClause, int wtParam, double numGroups) { @@ -3651,6 +3652,7 @@ create_recursiveunion_path(PlannerInfo *root, pathnode->leftpath = leftpath; pathnode->rightpath = rightpath; pathnode->distinctList = distinctList; + pathnode->distinctSortClause = distinctSortClause; pathnode->wtParam = wtParam; pathnode->numGroups = numGroups; diff --git a/src/include/nodes/execnodes.h b/src/include/nodes/execnodes.h index e3cea52e7df..af2208b5de4 100644 --- a/src/include/nodes/execnodes.h +++ b/src/include/nodes/execnodes.h @@ -1591,6 +1591,11 @@ typedef struct RecursiveUnionState MemoryContext tempContext; /* short-term context for comparisons */ TupleHashTable hashtable; /* hash table for tuples already seen */ MemoryContext tuplesContext; /* context containing hash table's tuples */ + + /* Sort keys for replacement (for subset DISTINCT ON with ORDER BY) */ + SortSupport sortKeys; + TupleTableSlot *sort_firstTupleSlot; + Tuplestorestate *result_table; /* buffered results for UNION DISTINCT ON */ } RecursiveUnionState; /* ---------------- diff --git a/src/include/nodes/pathnodes.h b/src/include/nodes/pathnodes.h index d798a631a29..dad02886fc3 100644 --- a/src/include/nodes/pathnodes.h +++ b/src/include/nodes/pathnodes.h @@ -2687,6 +2687,7 @@ typedef struct RecursiveUnionPath Path *leftpath; /* paths representing input sources */ Path *rightpath; List *distinctList; /* SortGroupClauses identifying target cols */ + List *distinctSortClause; /* SortGroupClauses for ordering (if DISTINCT ON) */ int wtParam; /* ID of Param representing work table */ Cardinality numGroups; /* estimated number of groups in input */ } RecursiveUnionPath; diff --git a/src/include/nodes/plannodes.h b/src/include/nodes/plannodes.h index b687edb3a9f..855357f512d 100644 --- a/src/include/nodes/plannodes.h +++ b/src/include/nodes/plannodes.h @@ -495,6 +495,13 @@ typedef struct RecursiveUnion /* estimated number of groups in input */ Cardinality numGroups; + + /* sort keys for replacement (for subset DISTINCT ON with ORDER BY) */ + int numSortCols; + AttrNumber *sortColIdx pg_node_attr(array_size(numSortCols)); + Oid *sortOperators pg_node_attr(array_size(numSortCols)); + Oid *sortCollations pg_node_attr(array_size(numSortCols)); + bool *sortNullsFirst pg_node_attr(array_size(numSortCols)); } RecursiveUnion; /* ---------------- diff --git a/src/include/optimizer/pathnode.h b/src/include/optimizer/pathnode.h index d514eb0c3ba..0c34acad269 100644 --- a/src/include/optimizer/pathnode.h +++ b/src/include/optimizer/pathnode.h @@ -316,6 +316,7 @@ extern RecursiveUnionPath *create_recursiveunion_path(PlannerInfo *root, Path *rightpath, PathTarget *target, List *distinctList, + List *distinctSortClause, int wtParam, double numGroups); extern LockRowsPath *create_lockrows_path(PlannerInfo *root, RelOptInfo *rel, diff --git a/src/test/regress/expected/union.out b/src/test/regress/expected/union.out index 91fd248fcad..0c6f4962eef 100644 --- a/src/test/regress/expected/union.out +++ b/src/test/regress/expected/union.out @@ -1852,3 +1852,22 @@ SELECT * FROM ( (3 rows) DROP TABLE union_distinct_u1, union_distinct_u2; +-- Recursive CTE Shortest Path (Pruning) +CREATE TABLE union_distinct_edges (src int, dst int, cost int); +INSERT INTO union_distinct_edges VALUES (1, 2, 10), (1, 3, 2), (3, 2, 3), (2, 4, 1); +WITH RECURSIVE search(node, cost, path) AS ( + SELECT 1 AS node, 0 AS cost, ARRAY[1] AS path + UNION DISTINCT ON (node ORDER BY cost ASC) + SELECT e.dst, s.cost + e.cost, s.path || e.dst + FROM search s JOIN union_distinct_edges e ON s.node = e.src +) +SELECT * FROM search ORDER BY node; + node | cost | path +------+------+----------- + 1 | 0 | {1} + 2 | 5 | {1,3,2} + 3 | 2 | {1,3} + 4 | 6 | {1,3,2,4} +(4 rows) + +DROP TABLE union_distinct_edges; diff --git a/src/test/regress/sql/union.sql b/src/test/regress/sql/union.sql index e6813edc244..1569d3c3c98 100644 --- a/src/test/regress/sql/union.sql +++ b/src/test/regress/sql/union.sql @@ -738,3 +738,16 @@ SELECT * FROM ( DROP TABLE union_distinct_u1, union_distinct_u2; +-- Recursive CTE Shortest Path (Pruning) +CREATE TABLE union_distinct_edges (src int, dst int, cost int); +INSERT INTO union_distinct_edges VALUES (1, 2, 10), (1, 3, 2), (3, 2, 3), (2, 4, 1); + +WITH RECURSIVE search(node, cost, path) AS ( + SELECT 1 AS node, 0 AS cost, ARRAY[1] AS path + UNION DISTINCT ON (node ORDER BY cost ASC) + SELECT e.dst, s.cost + e.cost, s.path || e.dst + FROM search s JOIN union_distinct_edges e ON s.node = e.src +) +SELECT * FROM search ORDER BY node; + +DROP TABLE union_distinct_edges; -- 2.56.0.176.g3575e65682-goog