From c0e8a3d74a9be9577b4390ac14152671d3f9a5e6 Mon Sep 17 00:00:00 2001 From: Hannu Krosing Date: Wed, 16 Sep 2026 13:41:19 +0000 Subject: [PATCH 4/4] 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 788ed99935f..7e53c353870 100644 --- a/doc/src/sgml/ref/select.sgml +++ b/doc/src/sgml/ref/select.sgml @@ -1330,7 +1330,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 4524a6e11f0..208c043534a 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); @@ -5921,6 +5923,7 @@ make_recursive_union(List *tlist, Plan *righttree, int wtParam, List *distinctList, + List *distinctSortClause, Cardinality numGroups) { RecursiveUnion *node = makeNode(RecursiveUnion); @@ -5966,6 +5969,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 a424865f61e..cf47adbd7ad 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 d1d129c2f23..397b1367766 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 f0cb21444b2..dd4db8e6798 100644 --- a/src/include/nodes/execnodes.h +++ b/src/include/nodes/execnodes.h @@ -1613,6 +1613,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 aebaaa60304..512514e9c5f 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 accdd1f5098..10410d87f47 100644 --- a/src/include/nodes/plannodes.h +++ b/src/include/nodes/plannodes.h @@ -497,6 +497,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 85c4d4fe9d3..3f06e4c6d73 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 30efade7968..5f925e76c65 100644 --- a/src/test/regress/expected/union.out +++ b/src/test/regress/expected/union.out @@ -1754,3 +1754,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 9c4844bfdb1..31f13f350de 100644 --- a/src/test/regress/sql/union.sql +++ b/src/test/regress/sql/union.sql @@ -702,3 +702,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.55.0.1082.g2b9226bbc0-goog