Rebased
From 96f5a3d8f90aeab2b6d562a3fefb799240bafe45 Mon Sep 17 00:00:00 2001
From: Hannu Krosing <[email protected]>
Date: Wed, 16 Sep 2026 13:12:52 +0000
Subject: [PATCH 1/4] Refactor set operation group clauses to decouple from
 targetlist positions

Previously, generate_setop_child_grouplist, generate_setop_grouplist,
and add_setop_child_rel_equivalences assumed that groupClauses had a
strict 1:1 positional correspondence with all non-junk targetlist columns.

This patch decouples set operation group clauses from targetlist positions
by matching them through SortGroupClause.tleSortGroupRef. This prepares
the planner and equivalence class machinery for set operations that group
on a subset of columns.
---
 src/backend/optimizer/path/equivclass.c  | 35 +++++++-------
 src/backend/optimizer/plan/planner.c     | 58 ++++++++++++++----------
 src/backend/optimizer/prep/prepunion.c   | 24 +++-------
 src/backend/parser/analyze.c             |  9 +++-
 src/backend/rewrite/rewriteSearchCycle.c | 21 ++++++---
 5 files changed, 80 insertions(+), 67 deletions(-)

diff --git a/src/backend/optimizer/path/equivclass.c b/src/backend/optimizer/path/equivclass.c
index 7995c48e0e2..d48b7567ac0 100644
--- a/src/backend/optimizer/path/equivclass.c
+++ b/src/backend/optimizer/path/equivclass.c
@@ -3053,30 +3053,31 @@ add_setop_child_rel_equivalences(PlannerInfo *root, RelOptInfo *child_rel,
 								 List *child_tlist, List *setop_pathkeys)
 {
 	ListCell   *lc;
-	ListCell   *lc2 = list_head(setop_pathkeys);
 
-	foreach(lc, child_tlist)
+	foreach(lc, setop_pathkeys)
 	{
-		TargetEntry *tle = lfirst_node(TargetEntry, lc);
+		PathKey    *pk = lfirst_node(PathKey, lc);
+		Index		ref = pk->pk_eclass->ec_sortref;
+		TargetEntry *tle = NULL;
+		ListCell   *lc2;
 		EquivalenceMember *parent_em;
-		PathKey    *pk;
 
-		if (tle->resjunk)
-			continue;
+		/* Find matching non-resjunk tle in child_tlist */
+		foreach(lc2, child_tlist)
+		{
+			TargetEntry *cur_tle = lfirst_node(TargetEntry, lc2);
+			if (cur_tle->ressortgroupref == ref && !cur_tle->resjunk)
+			{
+				tle = cur_tle;
+				break;
+			}
+		}
 
-		if (lc2 == NULL)
-			elog(ERROR, "too few pathkeys for set operation");
+		if (tle == NULL)
+			elog(ERROR, "could not find target entry for setop pathkey ref %d", ref);
 
-		pk = lfirst_node(PathKey, lc2);
 		parent_em = linitial(pk->pk_eclass->ec_members);
 
-		/*
-		 * We can safely pass the parent member as the first member in the
-		 * ec_members list as this is added first in generate_union_paths,
-		 * likewise, the JoinDomain can be that of the initial member of the
-		 * Pathkey's EquivalenceClass.  We pass -1 for ec_index since we
-		 * maintain the eclass_indexes for the child_rel after the loop.
-		 */
 		add_child_eq_member(root,
 							pk->pk_eclass,
 							-1,
@@ -3086,8 +3087,6 @@ add_setop_child_rel_equivalences(PlannerInfo *root, RelOptInfo *child_rel,
 							parent_em,
 							exprType((Node *) tle->expr),
 							child_rel->relid);
-
-		lc2 = lnext(setop_pathkeys, lc2);
 	}
 
 	/*
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index 8d30131855a..524019a6efc 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -8622,49 +8622,59 @@ group_by_has_partkey(RelOptInfo *input_rel,
  * then we return an empty list.  This may leave some TLEs with unreferenced
  * ressortgroupref markings, but that's harmless.
  */
+static TargetEntry *
+get_nth_nonjunk_tle(List *tlist, int n)
+{
+	ListCell   *lc;
+	int			count = 0;
+
+	foreach(lc, tlist)
+	{
+		TargetEntry *tle = lfirst_node(TargetEntry, lc);
+
+		if (!tle->resjunk)
+		{
+			count++;
+			if (count == n)
+				return tle;
+		}
+	}
+	return NULL;
+}
+
 static List *
 generate_setop_child_grouplist(SetOperationStmt *op, List *targetlist)
 {
 	List	   *grouplist = copyObject(op->groupClauses);
 	ListCell   *lg;
-	ListCell   *lt;
-	ListCell   *ct;
 
-	lg = list_head(grouplist);
-	ct = list_head(op->colTypes);
-	foreach(lt, targetlist)
+	foreach(lg, grouplist)
 	{
-		TargetEntry *tle = (TargetEntry *) lfirst(lt);
-		SortGroupClause *sgc;
+		SortGroupClause *sgc = (SortGroupClause *) lfirst(lg);
+		Index		ref = sgc->tleSortGroupRef;
+		TargetEntry *tle;
 		Oid			coltype;
 
-		/* resjunk columns could have sortgrouprefs.  Leave these alone */
-		if (tle->resjunk)
-			continue;
+		/* If tleSortGroupRef is not set, we can't map it. */
+		if (ref == 0)
+			elog(ERROR, "missing tleSortGroupRef in setop groupClause");
 
-		/*
-		 * We expect every non-resjunk target to have a SortGroupClause and
-		 * colTypes.
-		 */
-		Assert(lg != NULL);
-		Assert(ct != NULL);
-		sgc = (SortGroupClause *) lfirst(lg);
-		coltype = lfirst_oid(ct);
+		tle = get_nth_nonjunk_tle(targetlist, ref);
+		if (tle == NULL)
+			elog(ERROR, "missing target entry for setop groupClause ref %d", ref);
+
+		/* We also need to get the type from op->colTypes */
+		Assert(ref <= list_length(op->colTypes));
+		coltype = list_nth_oid(op->colTypes, ref - 1);
 
 		/* reject if target type isn't the same as the setop target type */
 		if (coltype != exprType((Node *) tle->expr))
 			return NIL;
 
-		lg = lnext(grouplist, lg);
-		ct = lnext(op->colTypes, ct);
-
 		/* assign a tleSortGroupRef, or reuse the existing one */
 		sgc->tleSortGroupRef = assignSortGroupRef(tle, targetlist);
 	}
 
-	Assert(lg == NULL);
-	Assert(ct == NULL);
-
 	return grouplist;
 }
 
diff --git a/src/backend/optimizer/prep/prepunion.c b/src/backend/optimizer/prep/prepunion.c
index b136f12ff3b..1efd0d915b7 100644
--- a/src/backend/optimizer/prep/prepunion.c
+++ b/src/backend/optimizer/prep/prepunion.c
@@ -1719,28 +1719,18 @@ generate_setop_grouplist(SetOperationStmt *op, List *targetlist)
 {
 	List	   *grouplist = copyObject(op->groupClauses);
 	ListCell   *lg;
-	ListCell   *lt;
 
-	lg = list_head(grouplist);
-	foreach(lt, targetlist)
+	foreach(lg, grouplist)
 	{
-		TargetEntry *tle = (TargetEntry *) lfirst(lt);
-		SortGroupClause *sgc;
+		SortGroupClause *sgc = (SortGroupClause *) lfirst(lg);
+		Index		ref = sgc->tleSortGroupRef;
+		TargetEntry *tle;
 
+		Assert(ref > 0 && ref <= list_length(targetlist));
+		tle = list_nth(targetlist, ref - 1);
 		Assert(!tle->resjunk);
-
-		/* non-resjunk columns should have sortgroupref = resno */
-		Assert(tle->ressortgroupref == tle->resno);
-
-		/* non-resjunk columns should have grouping clauses */
-		Assert(lg != NULL);
-		sgc = (SortGroupClause *) lfirst(lg);
-		lg = lnext(grouplist, lg);
-		Assert(sgc->tleSortGroupRef == 0);
-
-		sgc->tleSortGroupRef = tle->ressortgroupref;
+		Assert(tle->ressortgroupref == ref);
 	}
-	Assert(lg == NULL);
 	return grouplist;
 }
 
diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c
index 08f99dff711..fb7e106a0c9 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -2283,6 +2283,7 @@ constructSetOpTargetlist(ParseState *pstate, SetOperationStmt *op,
 {
 	ListCell   *ltl;
 	ListCell   *rtl;
+	int			resno = 1;
 
 	/*
 	 * Verify that the two children have the same number of non-junk columns,
@@ -2402,13 +2403,15 @@ constructSetOpTargetlist(ParseState *pstate, SetOperationStmt *op,
 		if (op->op != SETOP_UNION || !op->all)
 		{
 			ParseCallbackState pcbstate;
+			SortGroupClause *grpcl;
 
 			setup_parser_errposition_callback(&pcbstate, pstate,
 											  bestlocation);
 
 			/* If it's a recursive union, we need to require hashing support. */
-			op->groupClauses = lappend(op->groupClauses,
-									   makeSortGroupClauseForSetOp(rescoltype, recursive));
+			grpcl = makeSortGroupClauseForSetOp(rescoltype, recursive);
+			grpcl->tleSortGroupRef = resno;
+			op->groupClauses = lappend(op->groupClauses, grpcl);
 
 			cancel_parser_errposition_callback(&pcbstate);
 		}
@@ -2433,6 +2436,8 @@ constructSetOpTargetlist(ParseState *pstate, SetOperationStmt *op,
 									 false);
 			*targetlist = lappend(*targetlist, restle);
 		}
+
+		resno++;
 	}
 }
 
diff --git a/src/backend/rewrite/rewriteSearchCycle.c b/src/backend/rewrite/rewriteSearchCycle.c
index 75943072817..524a265a4c9 100644
--- a/src/backend/rewrite/rewriteSearchCycle.c
+++ b/src/backend/rewrite/rewriteSearchCycle.c
@@ -606,8 +606,11 @@ rewriteSearchAndCycle(CommonTableExpr *cte)
 		sos->colTypmods = lappend_int(sos->colTypmods, -1);
 		sos->colCollations = lappend_oid(sos->colCollations, InvalidOid);
 		if (!sos->all)
-			sos->groupClauses = lappend(sos->groupClauses,
-										makeSortGroupClauseForSetOp(search_seq_type, true));
+		{
+			SortGroupClause *sgc = makeSortGroupClauseForSetOp(search_seq_type, true);
+			sgc->tleSortGroupRef = list_length(sos->colTypes);
+			sos->groupClauses = lappend(sos->groupClauses, sgc);
+		}
 	}
 	if (cte->cycle_clause)
 	{
@@ -615,15 +618,21 @@ rewriteSearchAndCycle(CommonTableExpr *cte)
 		sos->colTypmods = lappend_int(sos->colTypmods, cte->cycle_clause->cycle_mark_typmod);
 		sos->colCollations = lappend_oid(sos->colCollations, cte->cycle_clause->cycle_mark_collation);
 		if (!sos->all)
-			sos->groupClauses = lappend(sos->groupClauses,
-										makeSortGroupClauseForSetOp(cte->cycle_clause->cycle_mark_type, true));
+		{
+			SortGroupClause *sgc = makeSortGroupClauseForSetOp(cte->cycle_clause->cycle_mark_type, true);
+			sgc->tleSortGroupRef = list_length(sos->colTypes);
+			sos->groupClauses = lappend(sos->groupClauses, sgc);
+		}
 
 		sos->colTypes = lappend_oid(sos->colTypes, RECORDARRAYOID);
 		sos->colTypmods = lappend_int(sos->colTypmods, -1);
 		sos->colCollations = lappend_oid(sos->colCollations, InvalidOid);
 		if (!sos->all)
-			sos->groupClauses = lappend(sos->groupClauses,
-										makeSortGroupClauseForSetOp(RECORDARRAYOID, true));
+		{
+			SortGroupClause *sgc = makeSortGroupClauseForSetOp(RECORDARRAYOID, true);
+			sgc->tleSortGroupRef = list_length(sos->colTypes);
+			sos->groupClauses = lappend(sos->groupClauses, sgc);
+		}
 	}
 
 	/*
-- 
2.55.0.1082.g2b9226bbc0-goog

From 0445719915a897a98fd5888d9bd15f88dcc2ff46 Mon Sep 17 00:00:00 2001
From: Hannu Krosing <[email protected]>
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 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 <literal>ORDER BY</literal> clause is specified inside
     <literal>DISTINCT ON</literal>, 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 (<literal>WITH RECURSIVE</literal>),
+    <literal>UNION DISTINCT ON</literal> with <literal>ORDER BY</literal> can be used
+    to prune the search space and find optimal solutions (e.g., shortest paths).
    </para>
 
    <para>
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 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 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 91bb0bd2e13..1d0b52e8580 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 73e518ce501..25d324ac6e6 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 2c66ba830a9..5b07dc6e7b5 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 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

From 91f030f58542918c630a70772e41c2f25f3f4b89 Mon Sep 17 00:00:00 2001
From: Hannu Krosing <[email protected]>
Date: Wed, 16 Sep 2026 13:32:07 +0000
Subject: [PATCH 3/4] Support UNION DISTINCT ON for non-recursive set
 operations

Extends set operations to support duplicate elimination on a subset of
columns using the syntax:
    select_stmt UNION DISTINCT ON (keys [ORDER BY sort_keys]) select_stmt

Builds upon the setop group clause refactoring to plan both HashAgg and
Sort->Unique paths for set operations with partial group lists.
Also prevents invalid subquery qual pushdown onto non-distinct attributes
via the UNSAFE_NOTIN_DISTINCTON_CLAUSE safety check.
---
 doc/src/sgml/ref/select.sgml           |   8 +-
 src/backend/optimizer/path/allpaths.c  |  19 ++++
 src/backend/optimizer/plan/planner.c   |   6 +-
 src/backend/optimizer/prep/prepunion.c |  57 +++++------
 src/backend/parser/analyze.c           | 126 +++++++++++++++++++++----
 src/backend/parser/gram.y              |  26 ++++-
 src/include/nodes/parsenodes.h         |   1 +
 src/test/regress/expected/union.out    |  48 ++++++++++
 src/test/regress/sql/union.sql         |  28 ++++++
 9 files changed, 268 insertions(+), 51 deletions(-)

diff --git a/doc/src/sgml/ref/select.sgml b/doc/src/sgml/ref/select.sgml
index 1ed53b2b88e..ca09f75c027 100644
--- a/doc/src/sgml/ref/select.sgml
+++ b/doc/src/sgml/ref/select.sgml
@@ -1251,7 +1251,7 @@ SELECT DISTINCT ON (location ORDER BY time DESC) location, time, report
    <para>
     The <literal>UNION</literal> clause has this general form:
 <synopsis>
-<replaceable class="parameter">select_statement</replaceable> UNION [ ALL | DISTINCT ] <replaceable class="parameter">select_statement</replaceable>
+<replaceable class="parameter">select_statement</replaceable> UNION [ ALL | DISTINCT | DISTINCT ON ( <replaceable class="parameter">expression</replaceable> [, ...] [ ORDER BY <replaceable class="parameter">sort_expression</replaceable> [ ASC | DESC ] [, ...] ] ) ] <replaceable class="parameter">select_statement</replaceable>
 </synopsis><replaceable class="parameter">select_statement</replaceable> is
     any <command>SELECT</command> statement without an <literal>ORDER
     BY</literal>, <literal>LIMIT</literal>, <literal>FOR NO KEY UPDATE</literal>, <literal>FOR UPDATE</literal>,
@@ -1282,6 +1282,12 @@ SELECT DISTINCT ON (location ORDER BY time DESC) location, time, report
     <literal>UNION</literal>; use <literal>ALL</literal> when you can.)
     <literal>DISTINCT</literal> can be written to explicitly specify the
     default behavior of eliminating duplicate rows.
+    <literal>DISTINCT ON</literal> keeps only the first row of each set of rows
+    that are duplicates according to the specified expressions.
+    If the optional <literal>ORDER BY</literal> clause is specified inside
+    <literal>DISTINCT ON</literal>, 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.
    </para>
 
    <para>
diff --git a/src/backend/optimizer/path/allpaths.c b/src/backend/optimizer/path/allpaths.c
index de8f29c2b57..9cb6bfe0b60 100644
--- a/src/backend/optimizer/path/allpaths.c
+++ b/src/backend/optimizer/path/allpaths.c
@@ -4161,6 +4161,25 @@ subquery_is_pushdown_safe(Query *subquery, Query *topquery,
 	 */
 	if (subquery->setOperations == NULL)
 		check_output_expressions(subquery, safetyInfo);
+	else if (subquery->hasDistinctOn)
+	{
+		ListCell   *lc;
+
+		foreach(lc, subquery->targetList)
+		{
+			TargetEntry *tle = (TargetEntry *) lfirst(lc);
+
+			if (tle->resjunk)
+				continue;
+
+			if ((safetyInfo->unsafeFlags[tle->resno] &
+				 UNSAFE_NOTIN_DISTINCTON_CLAUSE) == 0 &&
+				!targetIsInSortList(tle, InvalidOid, subquery->distinctClause))
+			{
+				safetyInfo->unsafeFlags[tle->resno] |= UNSAFE_NOTIN_DISTINCTON_CLAUSE;
+			}
+		}
+	}
 
 	/* Are we at top level, or looking at a setop component? */
 	if (subquery == topquery)
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index 5ed0de6224d..f540f930302 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -1806,7 +1806,6 @@ grouping_planner(PlannerInfo *root, double tuple_fraction,
 		/*
 		 * Calculate pathkeys that represent result ordering requirements
 		 */
-		Assert(parse->distinctClause == NIL);
 		root->sort_pathkeys = make_pathkeys_for_sortclauses(root,
 															parse->sortClause,
 															root->processed_tlist);
@@ -3920,7 +3919,7 @@ standard_qp_callback(PlannerInfo *root, void *extra)
 												   false,
 												   false,
 												   &sortable,
-												   false);
+												   true);
 		if (!sortable)
 			root->setop_pathkeys = NIL;
 	}
@@ -8675,7 +8674,8 @@ get_nth_nonjunk_tle(List *tlist, int n)
 static List *
 generate_setop_child_grouplist(SetOperationStmt *op, List *targetlist)
 {
-	List	   *grouplist = copyObject(op->groupClauses);
+	List	   *clauses = op->sortClauses ? op->sortClauses : op->groupClauses;
+	List	   *grouplist = copyObject(clauses);
 	ListCell   *lg;
 
 	foreach(lg, grouplist)
diff --git a/src/backend/optimizer/prep/prepunion.c b/src/backend/optimizer/prep/prepunion.c
index 1efd0d915b7..a424865f61e 100644
--- a/src/backend/optimizer/prep/prepunion.c
+++ b/src/backend/optimizer/prep/prepunion.c
@@ -113,7 +113,6 @@ plan_set_operations(PlannerInfo *root)
 	Assert(parse->groupClause == NIL);
 	Assert(parse->havingQual == NULL);
 	Assert(parse->windowClause == NIL);
-	Assert(parse->distinctClause == NIL);
 
 	/*
 	 * In the outer query level, equivalence classes are limited to classes
@@ -934,32 +933,34 @@ generate_union_paths(SetOperationStmt *op, PlannerInfo *root,
 			 * Try a hash aggregate plan on 'apath'.  This is the cheapest
 			 * available path containing each append child.
 			 */
-			path = (Path *) create_agg_path(root,
-											result_rel,
-											apath,
-											result_rel->reltarget,
-											AGG_HASHED,
-											AGGSPLIT_SIMPLE,
-											groupList,
-											NIL,
-											NULL,
-											dNumChildGroups);
-			add_path(result_rel, path);
-
-			/* Try hash aggregate on the Gather path, if valid */
-			if (gpath != NULL)
-			{
-				/* Hashed aggregate plan --- no sort needed */
-				path = (Path *) create_agg_path(root,
+			path = (Path *) create_agg_path_ext(root,
 												result_rel,
-												gpath,
+												apath,
 												result_rel->reltarget,
 												AGG_HASHED,
 												AGGSPLIT_SIMPLE,
 												groupList,
 												NIL,
 												NULL,
-												dNumChildGroups);
+												dNumChildGroups,
+												op->sortClauses);
+			add_path(result_rel, path);
+
+			/* Try hash aggregate on the Gather path, if valid */
+			if (gpath != NULL)
+			{
+				/* Hashed aggregate plan --- no sort needed */
+				path = (Path *) create_agg_path_ext(root,
+													result_rel,
+													gpath,
+													result_rel->reltarget,
+													AGG_HASHED,
+													AGGSPLIT_SIMPLE,
+													groupList,
+													NIL,
+													NULL,
+													dNumChildGroups,
+													op->sortClauses);
 				add_path(result_rel, path);
 			}
 		}
@@ -967,17 +968,18 @@ generate_union_paths(SetOperationStmt *op, PlannerInfo *root,
 		if (can_sort)
 		{
 			Path	   *path = apath;
+			List	   *sort_clauses = op->sortClauses ? op->sortClauses : groupList;
 
 			/* Try Sort -> Unique on the Append path */
-			if (groupList != NIL)
+			if (sort_clauses != NIL)
 				path = (Path *) create_sort_path(root, result_rel, path,
-												 make_pathkeys_for_sortclauses(root, groupList, tlist),
+												 make_pathkeys_for_sortclauses(root, sort_clauses, tlist),
 												 -1.0);
 
 			path = (Path *) create_unique_path(root,
 											   result_rel,
 											   path,
-											   list_length(path->pathkeys),
+											   list_length(groupList),
 											   dNumChildGroups);
 
 			add_path(result_rel, path);
@@ -987,14 +989,15 @@ generate_union_paths(SetOperationStmt *op, PlannerInfo *root,
 			{
 				path = gpath;
 
-				path = (Path *) create_sort_path(root, result_rel, path,
-												 make_pathkeys_for_sortclauses(root, groupList, tlist),
-												 -1.0);
+				if (sort_clauses != NIL)
+					path = (Path *) create_sort_path(root, result_rel, path,
+													 make_pathkeys_for_sortclauses(root, sort_clauses, tlist),
+													 -1.0);
 
 				path = (Path *) create_unique_path(root,
 												   result_rel,
 												   path,
-												   list_length(path->pathkeys),
+												   list_length(groupList),
 												   dNumChildGroups);
 				add_path(result_rel, path);
 			}
diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c
index b63813360f1..32fa02c73ae 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -78,10 +78,13 @@ static Query *transformSelectStmt(ParseState *pstate, SelectStmt *stmt,
 static Query *transformValuesClause(ParseState *pstate, SelectStmt *stmt);
 static Query *transformSetOperationStmt(ParseState *pstate, SelectStmt *stmt);
 static Node *transformSetOperationTree(ParseState *pstate, SelectStmt *stmt,
-									   bool isTopLevel, List **targetlist);
+									   bool isTopLevel, List **targetlist,
+									   List *distinctClause);
 static void constructSetOpTargetlist(ParseState *pstate, SetOperationStmt *op,
 									 const List *ltargetlist, const List *rtargetlist,
-									 List **targetlist, const char *context, bool recursive);
+									 List **targetlist, const char *context, bool recursive,
+									 List *distinctClause);
+static bool col_in_distinct_on(const char *colname, int resno, List *distinctClause);
 static void determineRecursiveColTypes(ParseState *pstate,
 									   Node *larg, List *nrtargetlist);
 static Query *transformReturnStmt(ParseState *pstate, ReturnStmt *stmt);
@@ -1854,6 +1857,10 @@ transformSetOperationStmt(ParseState *pstate, SelectStmt *stmt)
 	ParseNamespaceColumn *sortnscolumns;
 	int			sortcolindex;
 	int			tllen;
+	List	   *distinctClause = stmt->distinctClause;
+	List	   *distinctSortClause = stmt->distinctSortClause;
+	List	   *transformed_distinctSortClause = NIL;
+	int			distinct_tllen;
 
 	qry->commandType = CMD_SELECT;
 
@@ -1889,6 +1896,8 @@ transformSetOperationStmt(ParseState *pstate, SelectStmt *stmt)
 	withClause = stmt->withClause;
 
 	stmt->sortClause = NIL;
+	stmt->distinctClause = NIL;
+	stmt->distinctSortClause = NIL;
 	stmt->limitOffset = NULL;
 	stmt->limitCount = NULL;
 	stmt->lockingClause = NIL;
@@ -1916,7 +1925,7 @@ transformSetOperationStmt(ParseState *pstate, SelectStmt *stmt)
 	 * Recursively transform the components of the tree.
 	 */
 	sostmt = castNode(SetOperationStmt,
-					  transformSetOperationTree(pstate, stmt, true, NULL));
+					  transformSetOperationTree(pstate, stmt, true, NULL, distinctClause));
 	Assert(sostmt);
 	qry->setOperations = (Node *) sostmt;
 
@@ -2032,11 +2041,22 @@ transformSetOperationStmt(ParseState *pstate, SelectStmt *stmt)
 										  EXPR_KIND_ORDER_BY,
 										  false /* allow SQL92 rules */ );
 
+	distinct_tllen = list_length(qry->targetList);
+	if (distinctSortClause)
+	{
+		List *full_sortby = prepend_distinct_to_sortby(distinctClause, distinctSortClause);
+		transformed_distinctSortClause = transformSortClause(pstate,
+															 full_sortby,
+															 &qry->targetList,
+															 EXPR_KIND_ORDER_BY,
+															 false);
+	}
+
 	/* restore namespace, remove join RTE from rtable */
 	pstate->p_namespace = sv_namespace;
 	pstate->p_rtable = list_truncate(pstate->p_rtable, sv_rtable_length);
 
-	if (tllen != list_length(qry->targetList))
+	if (tllen != distinct_tllen)
 		ereport(ERROR,
 				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 				 errmsg("invalid UNION/INTERSECT/EXCEPT ORDER BY clause"),
@@ -2045,6 +2065,26 @@ transformSetOperationStmt(ParseState *pstate, SelectStmt *stmt)
 				 parser_errposition(pstate,
 									exprLocation(list_nth(qry->targetList, tllen)))));
 
+	if (distinct_tllen != list_length(qry->targetList))
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("invalid UNION/INTERSECT/EXCEPT DISTINCT ON ORDER BY clause"),
+				 errdetail("Only result column names can be used, not expressions or functions."),
+				 parser_errposition(pstate,
+									exprLocation(list_nth(qry->targetList, distinct_tllen)))));
+
+	qry->distinctSortClause = transformed_distinctSortClause;
+	sostmt->sortClauses = transformed_distinctSortClause;
+
+	if (distinctClause)
+	{
+		qry->distinctClause = transformDistinctOnClause(pstate,
+														distinctClause,
+														&qry->targetList,
+														transformed_distinctSortClause ? transformed_distinctSortClause : qry->sortClause);
+		qry->hasDistinctOn = true;
+	}
+
 	qry->limitOffset = transformLimitClause(pstate, limitOffset,
 											EXPR_KIND_OFFSET, "OFFSET",
 											stmt->limitOption);
@@ -2133,7 +2173,8 @@ makeSortGroupClauseForSetOp(Oid rescoltype, bool require_hash)
  */
 static Node *
 transformSetOperationTree(ParseState *pstate, SelectStmt *stmt,
-						  bool isTopLevel, List **targetlist)
+						  bool isTopLevel, List **targetlist,
+						  List *distinctClause)
 {
 	bool		isLeaf;
 
@@ -2278,7 +2319,8 @@ transformSetOperationTree(ParseState *pstate, SelectStmt *stmt,
 		 */
 		op->larg = transformSetOperationTree(pstate, stmt->larg,
 											 false,
-											 &ltargetlist);
+											 &ltargetlist,
+											 NIL);
 
 		/*
 		 * If we are processing a recursive union query, now is the time to
@@ -2294,10 +2336,11 @@ transformSetOperationTree(ParseState *pstate, SelectStmt *stmt,
 		 */
 		op->rarg = transformSetOperationTree(pstate, stmt->rarg,
 											 false,
-											 &rtargetlist);
+											 &rtargetlist,
+											 NIL);
 
 		constructSetOpTargetlist(pstate, op, ltargetlist, rtargetlist, targetlist,
-								 context, recursive);
+								 context, recursive, distinctClause);
 
 		return (Node *) op;
 	}
@@ -2318,7 +2361,8 @@ transformSetOperationTree(ParseState *pstate, SelectStmt *stmt,
 static void
 constructSetOpTargetlist(ParseState *pstate, SetOperationStmt *op,
 						 const List *ltargetlist, const List *rtargetlist,
-						 List **targetlist, const char *context, bool recursive)
+						 List **targetlist, const char *context, bool recursive,
+						 List *distinctClause)
 {
 	ListCell   *ltl;
 	ListCell   *rtl;
@@ -2441,18 +2485,27 @@ constructSetOpTargetlist(ParseState *pstate, SetOperationStmt *op,
 		 */
 		if (op->op != SETOP_UNION || !op->all)
 		{
-			ParseCallbackState pcbstate;
-			SortGroupClause *grpcl;
+			bool active = false;
+			if (distinctClause == NIL)
+				active = true;
+			else
+				active = col_in_distinct_on(ltle->resname, resno, distinctClause);
 
-			setup_parser_errposition_callback(&pcbstate, pstate,
-											  bestlocation);
+			if (active)
+			{
+				ParseCallbackState pcbstate;
+				SortGroupClause *grpcl;
+
+				setup_parser_errposition_callback(&pcbstate, pstate,
+												  bestlocation);
 
-			/* If it's a recursive union, we need to require hashing support. */
-			grpcl = makeSortGroupClauseForSetOp(rescoltype, recursive);
-			grpcl->tleSortGroupRef = resno;
-			op->groupClauses = lappend(op->groupClauses, grpcl);
+				/* If it's a recursive union, we need to require hashing support. */
+				grpcl = makeSortGroupClauseForSetOp(rescoltype, recursive);
+				grpcl->tleSortGroupRef = resno;
+				op->groupClauses = lappend(op->groupClauses, grpcl);
 
-			cancel_parser_errposition_callback(&pcbstate);
+				cancel_parser_errposition_callback(&pcbstate);
+			}
 		}
 
 		/*
@@ -3795,3 +3848,40 @@ test_raw_expression_coverage(Node *node, void *context)
 									  context);
 }
 #endif							/* DEBUG_NODE_TESTS_ENABLED */
+
+/*
+ * col_in_distinct_on -
+ *	  Check if a column name or its 1-based position matches any expression in distinctClause
+ */
+static bool
+col_in_distinct_on(const char *colname, int resno, List *distinctClause)
+{
+	ListCell   *lc;
+
+	foreach(lc, distinctClause)
+	{
+		Node	   *n = (Node *) lfirst(lc);
+
+		if (IsA(n, ColumnRef))
+		{
+			ColumnRef  *cr = (ColumnRef *) n;
+
+			if (list_length(cr->fields) == 1 && IsA(linitial(cr->fields), String))
+			{
+				if (strcmp(strVal(linitial(cr->fields)), colname) == 0)
+					return true;
+			}
+		}
+		else if (IsA(n, A_Const))
+		{
+			A_Const    *aconst = (A_Const *) n;
+
+			if (IsA(&aconst->val, Integer))
+			{
+				if (intVal(&aconst->val) == resno)
+					return true;
+			}
+		}
+	}
+	return false;
+}
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index a8bb0a4d8d5..6247ba94fd4 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -13263,9 +13263,31 @@ simple_select:
 					n->fromClause = list_make1($2);
 					$$ = (Node *) n;
 				}
-			| select_clause UNION set_quantifier select_clause
+			| select_clause UNION select_clause				%prec UNION
 				{
-					$$ = makeSetOp(SETOP_UNION, $3 == SET_QUANTIFIER_ALL, $1, $4);
+					$$ = makeSetOp(SETOP_UNION, false, $1, $3);
+				}
+			| select_clause UNION ALL select_clause			%prec UNION
+				{
+					$$ = makeSetOp(SETOP_UNION, true, $1, $4);
+				}
+			| select_clause UNION distinct_clause select_clause	%prec UNION
+				{
+					List *distinctClause = linitial($3);
+					List *distinctSortClause = lsecond($3);
+					Node *n = makeSetOp(SETOP_UNION, false, $1, $4);
+					SelectStmt *s = (SelectStmt *) n;
+					if (linitial(distinctClause) == NULL && distinctSortClause == NIL)
+					{
+						s->distinctClause = NIL;
+						s->distinctSortClause = NIL;
+					}
+					else
+					{
+						s->distinctClause = distinctClause;
+						s->distinctSortClause = distinctSortClause;
+					}
+					$$ = (Node *) s;
 				}
 			| select_clause INTERSECT set_quantifier select_clause
 				{
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index 135445f78b0..4fe1174a1d0 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -2313,6 +2313,7 @@ typedef struct SetOperationStmt
 	/* a list of SortGroupClause's */
 	List	   *groupClauses pg_node_attr(query_jumble_ignore);
 	/* groupClauses is NIL if UNION ALL, but must be set otherwise */
+	List	   *sortClauses pg_node_attr(query_jumble_ignore); /* inline DISTINCT ON ORDER BY clauses */
 } SetOperationStmt;
 
 
diff --git a/src/test/regress/expected/union.out b/src/test/regress/expected/union.out
index 84abcd6b14f..30efade7968 100644
--- a/src/test/regress/expected/union.out
+++ b/src/test/regress/expected/union.out
@@ -1706,3 +1706,51 @@ join (select ten from tenk1 union select ten from onek) s on s.ten = t.unique1;
          Index Cond: (unique1 = tenk1.ten)
 (8 rows)
 
+--
+-- UNION DISTINCT ON
+--
+CREATE TABLE union_distinct_u1 (a int, b int);
+CREATE TABLE union_distinct_u2 (a int, b int);
+INSERT INTO union_distinct_u1 VALUES (1, 10), (2, 20);
+INSERT INTO union_distinct_u2 VALUES (1, 100), (3, 30);
+-- UNION DISTINCT ON without ORDER BY
+SELECT a, b FROM union_distinct_u1 UNION DISTINCT ON (a) SELECT a, b FROM union_distinct_u2 ORDER BY a;
+ a | b  
+---+----
+ 1 | 10
+ 2 | 20
+ 3 | 30
+(3 rows)
+
+-- UNION DISTINCT ON with ORDER BY (DESC)
+SELECT a, b FROM union_distinct_u1 UNION DISTINCT ON (a ORDER BY b DESC) SELECT a, b FROM union_distinct_u2 ORDER BY a;
+ a |  b  
+---+-----
+ 1 | 100
+ 2 |  20
+ 3 |  30
+(3 rows)
+
+-- UNION DISTINCT ON with ORDER BY (ASC)
+SELECT a, b FROM union_distinct_u1 UNION DISTINCT ON (a ORDER BY b ASC) SELECT a, b FROM union_distinct_u2 ORDER BY a;
+ a | b  
+---+----
+ 1 | 10
+ 2 | 20
+ 3 | 30
+(3 rows)
+
+-- Test subquery pushdown safety with UNION DISTINCT ON
+SELECT * FROM (
+    SELECT a, b FROM union_distinct_u1
+    UNION DISTINCT ON (a ORDER BY b DESC)
+    SELECT a, b FROM union_distinct_u2
+) s WHERE b > 15 ORDER BY a;
+ a |  b  
+---+-----
+ 1 | 100
+ 2 |  20
+ 3 |  30
+(3 rows)
+
+DROP TABLE union_distinct_u1, union_distinct_u2;
diff --git a/src/test/regress/sql/union.sql b/src/test/regress/sql/union.sql
index c8de276c2b5..9c4844bfdb1 100644
--- a/src/test/regress/sql/union.sql
+++ b/src/test/regress/sql/union.sql
@@ -674,3 +674,31 @@ select null::int[] union all select null::int[] union all select null::bigint[];
 explain (costs off)
 select * from tenk1 t
 join (select ten from tenk1 union select ten from onek) s on s.ten = t.unique1;
+
+--
+-- UNION DISTINCT ON
+--
+
+CREATE TABLE union_distinct_u1 (a int, b int);
+CREATE TABLE union_distinct_u2 (a int, b int);
+INSERT INTO union_distinct_u1 VALUES (1, 10), (2, 20);
+INSERT INTO union_distinct_u2 VALUES (1, 100), (3, 30);
+
+-- UNION DISTINCT ON without ORDER BY
+SELECT a, b FROM union_distinct_u1 UNION DISTINCT ON (a) SELECT a, b FROM union_distinct_u2 ORDER BY a;
+
+-- UNION DISTINCT ON with ORDER BY (DESC)
+SELECT a, b FROM union_distinct_u1 UNION DISTINCT ON (a ORDER BY b DESC) SELECT a, b FROM union_distinct_u2 ORDER BY a;
+
+-- UNION DISTINCT ON with ORDER BY (ASC)
+SELECT a, b FROM union_distinct_u1 UNION DISTINCT ON (a ORDER BY b ASC) SELECT a, b FROM union_distinct_u2 ORDER BY a;
+
+-- Test subquery pushdown safety with UNION DISTINCT ON
+SELECT * FROM (
+    SELECT a, b FROM union_distinct_u1
+    UNION DISTINCT ON (a ORDER BY b DESC)
+    SELECT a, b FROM union_distinct_u2
+) s WHERE b > 15 ORDER BY a;
+
+DROP TABLE union_distinct_u1, union_distinct_u2;
+
-- 
2.55.0.1082.g2b9226bbc0-goog

From e11421380502405a642e196dda6766ea2d5566aa Mon Sep 17 00:00:00 2001
From: Hannu Krosing <[email protected]>
Date: Wed, 16 Sep 2026 13:22:47 +0000
Subject: [PATCH 2/4] Support inline ORDER BY in DISTINCT ON for SELECT and
 enable HashAgg

This patch introduces the syntax:
    SELECT DISTINCT ON (keys ORDER BY sort_keys) ...

This decouples the tie-breaker sort order used for duplicate elimination
from the final output sort order of the query, avoiding the need for
subqueries.

In addition, this enables Hash Aggregate to execute DISTINCT ON queries
with tie-breaking by introducing ReplaceTupleHashEntryIfBetter() and
using AllocSetContext for hashed tuples when sort columns are present.
---
 doc/src/sgml/ref/select.sgml                  | 37 +++++----
 src/backend/executor/execGrouping.c           | 80 +++++++++++++++++++
 src/backend/executor/nodeAgg.c                | 78 ++++++++++++++++--
 src/backend/optimizer/plan/createplan.c       | 46 +++++++++++
 src/backend/optimizer/plan/planner.c          | 42 ++++++++--
 src/backend/optimizer/util/pathnode.c         | 19 +++++
 src/backend/parser/analyze.c                  | 41 +++++++++-
 src/backend/parser/gram.y                     | 19 ++++-
 src/include/executor/executor.h               |  6 ++
 src/include/executor/nodeAgg.h                |  2 +
 src/include/nodes/parsenodes.h                |  3 +
 src/include/nodes/pathnodes.h                 |  1 +
 src/include/nodes/plannodes.h                 |  7 ++
 src/include/optimizer/pathnode.h              | 11 +++
 .../regress/expected/select_distinct_on.out   | 75 +++++++++++++++++
 src/test/regress/sql/select_distinct_on.sql   | 22 +++++
 16 files changed, 459 insertions(+), 30 deletions(-)

diff --git a/doc/src/sgml/ref/select.sgml b/doc/src/sgml/ref/select.sgml
index 18392f7caec..1ed53b2b88e 100644
--- a/doc/src/sgml/ref/select.sgml
+++ b/doc/src/sgml/ref/select.sgml
@@ -1203,30 +1203,39 @@ EXCLUDE NO OTHERS
 
    <para>
     <literal>SELECT DISTINCT ON ( <replaceable
-    class="parameter">expression</replaceable> [, ...] )</literal>
+    class="parameter">expression</replaceable> [, ...] [ ORDER BY <replaceable
+    class="parameter">sort_expression</replaceable> [ ASC | DESC ] [ NULLS { FIRST | LAST } ] [, ...] ] )</literal>
     keeps only the first row of each set of rows where the given
     expressions evaluate to equal.  The <literal>DISTINCT ON</literal>
     expressions are interpreted using the same rules as for
-    <literal>ORDER BY</literal> (see above).  Note that the <quote>first
-    row</quote> of each set is unpredictable unless <literal>ORDER
-    BY</literal> is used to ensure that the desired row appears first.  For
-    example:
+    <literal>ORDER BY</literal> (see above).
+    If the optional inline <literal>ORDER BY</literal> clause is specified,
+    it determines which row is kept from each set of duplicates (the first row
+    according to this sort order). This inline sorting applies only for duplicate
+    resolution and does not dictate the final output order of the query.
+    Note that if no inline <literal>ORDER BY</literal> is used, the <quote>first
+    row</quote> of each set is unpredictable unless a global <literal>ORDER
+    BY</literal> is used at the end of the query to ensure that the desired row appears first.
+    For example:
 <programlisting>
-SELECT DISTINCT ON (location) location, time, report
-    FROM weather_reports
-    ORDER BY location, time DESC;
+SELECT DISTINCT ON (location ORDER BY time DESC) location, time, report
+    FROM weather_reports;
 </programlisting>
     retrieves the most recent weather report for each location.  But
-    if we had not used <literal>ORDER BY</literal> to force descending order
-    of time values for each location, we'd have gotten a report from
-    an unpredictable time for each location.
+    if we had not used inline <literal>ORDER BY time DESC</literal>, we'd have gotten a report from
+    an unpredictable time for each location (unless a global <literal>ORDER BY</literal> was used).
    </para>
 
    <para>
-    The <literal>DISTINCT ON</literal> expression(s) must match the leftmost
-    <literal>ORDER BY</literal> expression(s).  The <literal>ORDER BY</literal> clause
-    will normally contain additional expression(s) that determine the
+    If a global <literal>ORDER BY</literal> clause is used at the end of the query
+    without an inline <literal>ORDER BY</literal> inside <literal>DISTINCT ON</literal>,
+    the <literal>DISTINCT ON</literal> expression(s) must match the leftmost
+    <literal>ORDER BY</literal> expression(s). The global <literal>ORDER BY</literal> clause
+    will then normally contain additional expression(s) that determine the
     desired precedence of rows within each <literal>DISTINCT ON</literal> group.
+    If an inline <literal>ORDER BY</literal> is specified, this matching requirement
+    is relaxed, and the global <literal>ORDER BY</literal> can sort the final results
+    by any columns, independent of the distinct keys.
    </para>
 
    <para>
diff --git a/src/backend/executor/execGrouping.c b/src/backend/executor/execGrouping.c
index 948aa2f177a..7a46b29ee2f 100644
--- a/src/backend/executor/execGrouping.c
+++ b/src/backend/executor/execGrouping.c
@@ -22,6 +22,7 @@
 #include "executor/executor.h"
 #include "miscadmin.h"
 #include "utils/lsyscache.h"
+#include "utils/sortsupport.h"
 
 static int	TupleHashTableMatch(struct tuplehash_hash *tb, MinimalTuple tuple1, MinimalTuple tuple2);
 static inline uint32 TupleHashTableHash_internal(struct tuplehash_hash *tb,
@@ -621,3 +622,82 @@ TupleHashTableMatch(struct tuplehash_hash *tb, MinimalTuple tuple1, MinimalTuple
 	econtext->ecxt_outertuple = slot1;
 	return !ExecQualAndReset(hashtable->cur_eq_func, econtext);
 }
+
+/*
+ * ReplaceTupleHashEntryIfBetter
+ *
+ * Compare the new slot with the stored tuple in the entry using the sort keys.
+ * If the new slot is "better" (comes before in sort order), replace the stored
+ * tuple in the entry.
+ *
+ * Returns true if replaced, false otherwise.
+ */
+bool
+ReplaceTupleHashEntryIfBetter(TupleHashTable hashtable,
+							  TupleHashEntry entry,
+							  TupleTableSlot *newslot,
+							  TupleTableSlot *firstslot,
+							  SortSupport sortKeys,
+							  int numSortCols)
+{
+	int			i;
+	bool		replace = false;
+
+	/* If no sort keys, we shouldn't be here */
+	if (numSortCols == 0)
+		return false;
+
+	/* Retrieve stored tuple and store it in firstslot */
+	ExecStoreMinimalTuple(entry->firstTuple, firstslot, false);
+
+	/* Compare sort keys one by one */
+	for (i = 0; i < numSortCols; i++)
+	{
+		SortSupport skey = &sortKeys[i];
+		AttrNumber	attno = skey->ssup_attno;
+		Datum		datum1, datum2;
+		bool		isnull1, isnull2;
+		int			compare;
+
+		datum1 = slot_getattr(firstslot, attno, &isnull1);
+		datum2 = slot_getattr(newslot, attno, &isnull2);
+
+		compare = ApplySortComparator(datum1, isnull1,
+									  datum2, isnull2,
+									  skey);
+
+		if (compare != 0)
+		{
+			/*
+			 * ApplySortComparator returns < 0 if datum1 comes BEFORE datum2.
+			 * So if compare > 0, datum2 comes BEFORE datum1, so it is better.
+			 */
+			if (compare > 0)
+				replace = true;
+			break; /* Found a difference, no need to compare further */
+		}
+	}
+
+	if (replace)
+	{
+		MinimalTuple oldtuple = entry->firstTuple;
+		MinimalTuple newtuple;
+		MemoryContext oldcxt;
+
+		/* Copy new tuple into the long-lived context */
+		oldcxt = MemoryContextSwitchTo(hashtable->tuplescxt);
+		newtuple = ExecCopySlotMinimalTuple(newslot);
+		MemoryContextSwitchTo(oldcxt);
+
+		/* Replace in entry */
+		entry->firstTuple = newtuple;
+
+		/* Free old tuple */
+		pfree(oldtuple);
+	}
+
+	/* Clear the comparison slot to avoid holding references */
+	ExecClearTuple(firstslot);
+
+	return replace;
+}
diff --git a/src/backend/executor/nodeAgg.c b/src/backend/executor/nodeAgg.c
index 29037cf3122..aba6bc05d87 100644
--- a/src/backend/executor/nodeAgg.c
+++ b/src/backend/executor/nodeAgg.c
@@ -277,6 +277,7 @@
 #include "utils/memutils_memorychunk.h"
 #include "utils/syscache.h"
 #include "utils/tuplesort.h"
+#include "utils/sortsupport.h"
 
 /*
  * Control how many partitions are created when spilling HashAgg to
@@ -1684,6 +1685,30 @@ find_hash_columns(AggState *aggstate)
 			ExecAllocTableSlot(&estate->es_tupleTable, hashDesc,
 							   &TTSOpsMinimalTuple, 0);
 
+		if (perhash->aggnode->numSortCols > 0)
+		{
+			perhash->hash_firstTupleSlot =
+				ExecAllocTableSlot(&estate->es_tupleTable, hashDesc,
+								   &TTSOpsMinimalTuple, 0);
+
+			/* Initialize sort support */
+			perhash->sortKeys = (SortSupport) palloc0(sizeof(SortSupportData) * perhash->aggnode->numSortCols);
+			for (i = 0; i < perhash->aggnode->numSortCols; i++)
+			{
+				SortSupport	skey = &perhash->sortKeys[i];
+
+				skey->ssup_collation = perhash->aggnode->sortCollations[i];
+				skey->ssup_nulls_first = perhash->aggnode->sortNullsFirst[i];
+				skey->ssup_attno = perhash->aggnode->sortColIdx[i];
+				PrepareSortSupportFromOrderingOp(perhash->aggnode->sortOperators[i], skey);
+			}
+		}
+		else
+		{
+			perhash->hash_firstTupleSlot = NULL;
+			perhash->sortKeys = NULL;
+		}
+
 		list_free(hashTlist);
 		bms_free(colnos);
 	}
@@ -1996,7 +2021,25 @@ hash_agg_update_metrics(AggState *aggstate, bool from_tape, int npartitions)
 static void
 hash_create_memory(AggState *aggstate)
 {
+	Agg		   *node = (Agg *) aggstate->ss.ps.plan;
 	Size		maxBlockSize = ALLOCSET_DEFAULT_MAXSIZE;
+	bool		use_allocset = false;
+	ListCell   *lc;
+
+	if (node->numSortCols > 0)
+		use_allocset = true;
+	else
+	{
+		foreach(lc, node->chain)
+		{
+			Agg *chained_node = lfirst_node(Agg, lc);
+			if (chained_node->numSortCols > 0)
+			{
+				use_allocset = true;
+				break;
+			}
+		}
+	}
 
 	/*
 	 * The hashcontext's per-tuple memory will be used for byref transition
@@ -2040,11 +2083,20 @@ hash_create_memory(AggState *aggstate)
 	/* and no smaller than ALLOCSET_DEFAULT_INITSIZE */
 	maxBlockSize = Max(maxBlockSize, ALLOCSET_DEFAULT_INITSIZE);
 
-	aggstate->hash_tuplescxt = BumpContextCreate(aggstate->ss.ps.state->es_query_cxt,
-												 "HashAgg hashed tuples",
-												 ALLOCSET_DEFAULT_MINSIZE,
-												 ALLOCSET_DEFAULT_INITSIZE,
-												 maxBlockSize);
+	if (use_allocset)
+	{
+		aggstate->hash_tuplescxt = AllocSetContextCreate(aggstate->ss.ps.state->es_query_cxt,
+														 "HashAgg hashed tuples",
+														 ALLOCSET_DEFAULT_SIZES);
+	}
+	else
+	{
+		aggstate->hash_tuplescxt = BumpContextCreate(aggstate->ss.ps.state->es_query_cxt,
+													 "HashAgg hashed tuples",
+													 ALLOCSET_DEFAULT_MINSIZE,
+													 ALLOCSET_DEFAULT_INITSIZE,
+													 maxBlockSize);
+	}
 
 }
 
@@ -2178,6 +2230,18 @@ initialize_hash_entry(AggState *aggstate, TupleHashTable hashtable,
  * for each grouping set, making the refilling of the hash table very
  * efficient.
  */
+static void
+replace_hash_entry_if_better(AggState *aggstate, AggStatePerHash perhash,
+							TupleHashEntry entry, TupleTableSlot *newslot)
+{
+	ReplaceTupleHashEntryIfBetter(perhash->hashtable,
+								  entry,
+								  newslot,
+								  perhash->hash_firstTupleSlot,
+								  perhash->sortKeys,
+								  perhash->aggnode->numSortCols);
+}
+
 static void
 lookup_hash_entries(AggState *aggstate)
 {
@@ -2210,6 +2274,8 @@ lookup_hash_entries(AggState *aggstate)
 		{
 			if (isnew)
 				initialize_hash_entry(aggstate, hashtable, entry);
+			else if (perhash->aggnode->numSortCols > 0)
+				replace_hash_entry_if_better(aggstate, perhash, entry, hashslot);
 			pergroup[setno] = TupleHashEntryGetAdditional(hashtable, entry);
 		}
 		else
@@ -2770,6 +2836,8 @@ agg_refill_hash_table(AggState *aggstate)
 		{
 			if (isnew)
 				initialize_hash_entry(aggstate, hashtable, entry);
+			else if (perhash->aggnode->numSortCols > 0)
+				replace_hash_entry_if_better(aggstate, perhash, entry, hashslot);
 			aggstate->hash_pergroup[batch->setno] = TupleHashEntryGetAdditional(hashtable, entry);
 			advance_aggregates(aggstate);
 		}
diff --git a/src/backend/optimizer/plan/createplan.c b/src/backend/optimizer/plan/createplan.c
index 2f8a08ad78e..78a6d6a3c20 100644
--- a/src/backend/optimizer/plan/createplan.c
+++ b/src/backend/optimizer/plan/createplan.c
@@ -2193,6 +2193,52 @@ create_agg_plan(PlannerInfo *root, AggPath *best_path)
 					best_path->transitionSpace,
 					subplan);
 
+	/* Extract sort keys for inline DISTINCT ON ORDER BY */
+	if (best_path->distinctSortClause)
+	{
+		List	   *sortcls = best_path->distinctSortClause;
+		List	   *sub_tlist = subplan->targetlist;
+		ListCell   *l;
+		int			numsortkeys;
+		AttrNumber *sortColIdx;
+		Oid		   *sortOperators;
+		Oid		   *collations;
+		bool	   *nullsFirst;
+
+		numsortkeys = list_length(sortcls);
+		sortColIdx = (AttrNumber *) palloc(numsortkeys * sizeof(AttrNumber));
+		sortOperators = (Oid *) palloc(numsortkeys * sizeof(Oid));
+		collations = (Oid *) palloc(numsortkeys * sizeof(Oid));
+		nullsFirst = (bool *) palloc(numsortkeys * sizeof(bool));
+
+		numsortkeys = 0;
+		foreach(l, sortcls)
+		{
+			SortGroupClause *sortcl = (SortGroupClause *) lfirst(l);
+			TargetEntry *tle = get_sortgroupclause_tle(sortcl, sub_tlist);
+
+			sortColIdx[numsortkeys] = tle->resno;
+			sortOperators[numsortkeys] = sortcl->sortop;
+			collations[numsortkeys] = exprCollation((Node *) tle->expr);
+			nullsFirst[numsortkeys] = sortcl->nulls_first;
+			numsortkeys++;
+		}
+
+		plan->numSortCols = numsortkeys;
+		plan->sortColIdx = sortColIdx;
+		plan->sortOperators = sortOperators;
+		plan->sortCollations = collations;
+		plan->sortNullsFirst = nullsFirst;
+	}
+	else
+	{
+		plan->numSortCols = 0;
+		plan->sortColIdx = NULL;
+		plan->sortOperators = NULL;
+		plan->sortCollations = NULL;
+		plan->sortNullsFirst = NULL;
+	}
+
 	copy_generic_path_info(&plan->plan, (Path *) best_path);
 
 	return plan;
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index 524019a6efc..5ed0de6224d 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -3856,16 +3856,46 @@ standard_qp_callback(PlannerInfo *root, void *extra)
 
 		/* Make a copy since pathkey processing can modify the list */
 		root->processed_distinctClause = list_copy(parse->distinctClause);
-		root->distinct_pathkeys =
+		
+		if (parse->distinctSortClause)
+		{
+			/* We have DISTINCT ON with ORDER BY */
+			List *temp_distinct_clause = list_copy(parse->distinctSortClause);
+			bool temp_sortable;
+			root->distinct_pathkeys =
+				make_pathkeys_for_sortclauses_extended(root,
+													   &temp_distinct_clause,
+													   tlist,
+													   true,
+													   false,
+													   &sortable,
+													   false);
+			if (!sortable)
+				root->distinct_pathkeys = NIL;
+
+			/* We ALSO need to remove redundant keys from processed_distinctClause */
 			make_pathkeys_for_sortclauses_extended(root,
 												   &root->processed_distinctClause,
 												   tlist,
 												   true,
 												   false,
-												   &sortable,
+												   &temp_sortable,
 												   false);
-		if (!sortable)
-			root->distinct_pathkeys = NIL;
+		}
+		else
+		{
+			/* Standard DISTINCT or DISTINCT ON without ORDER BY */
+			root->distinct_pathkeys =
+				make_pathkeys_for_sortclauses_extended(root,
+													   &root->processed_distinctClause,
+													   tlist,
+													   true,
+													   false,
+													   &sortable,
+													   false);
+			if (!sortable)
+				root->distinct_pathkeys = NIL;
+		}
 	}
 	else
 		root->distinct_pathkeys = NIL;
@@ -5254,7 +5284,7 @@ create_partial_distinct_paths(PlannerInfo *root, RelOptInfo *input_rel,
 					add_partial_path(partial_distinct_rel, (Path *)
 									 create_unique_path(root, partial_distinct_rel,
 														sorted_path,
-														list_length(root->distinct_pathkeys),
+														list_length(root->processed_distinctClause),
 														numDistinctRows));
 				}
 			}
@@ -5448,7 +5478,7 @@ create_final_distinct_paths(PlannerInfo *root, RelOptInfo *input_rel,
 					add_path(distinct_rel, (Path *)
 							 create_unique_path(root, distinct_rel,
 												sorted_path,
-												list_length(root->distinct_pathkeys),
+												list_length(root->processed_distinctClause),
 												numDistinctRows));
 				}
 			}
diff --git a/src/backend/optimizer/util/pathnode.c b/src/backend/optimizer/util/pathnode.c
index 67c47fd7c6c..fd2dd6808b0 100644
--- a/src/backend/optimizer/util/pathnode.c
+++ b/src/backend/optimizer/util/pathnode.c
@@ -3077,6 +3077,24 @@ create_agg_path(PlannerInfo *root,
 				List *qual,
 				const AggClauseCosts *aggcosts,
 				double numGroups)
+{
+	return create_agg_path_ext(root, rel, subpath, target,
+							   aggstrategy, aggsplit, groupClause,
+							   qual, aggcosts, numGroups, NIL);
+}
+
+AggPath *
+create_agg_path_ext(PlannerInfo *root,
+					RelOptInfo *rel,
+					Path *subpath,
+					PathTarget *target,
+					AggStrategy aggstrategy,
+					AggSplit aggsplit,
+					List *groupClause,
+					List *qual,
+					const AggClauseCosts *aggcosts,
+					double numGroups,
+					List *distinctSortClause)
 {
 	AggPath    *pathnode = makeNode(AggPath);
 
@@ -3115,6 +3133,7 @@ create_agg_path(PlannerInfo *root,
 	pathnode->numGroups = numGroups;
 	pathnode->transitionSpace = aggcosts ? aggcosts->transitionSpace : 0;
 	pathnode->groupClause = groupClause;
+	pathnode->distinctSortClause = distinctSortClause;
 	pathnode->qual = qual;
 
 	cost_agg(&pathnode->path, root,
diff --git a/src/backend/parser/analyze.c b/src/backend/parser/analyze.c
index fb7e106a0c9..b63813360f1 100644
--- a/src/backend/parser/analyze.c
+++ b/src/backend/parser/analyze.c
@@ -1417,6 +1417,33 @@ count_rowexpr_columns(ParseState *pstate, Node *expr)
  * Note: this covers only cases with no set operations and no VALUES lists;
  * see below for the other cases.
  */
+static List *
+prepend_distinct_to_sortby(List *distinctClause, List *distinctSortClause)
+{
+	List	   *result = list_copy(distinctSortClause);
+	ListCell   *lc;
+	List	   *prepended = NIL;
+
+	/* If distinctClause is empty or has NULL (SELECT DISTINCT), do nothing */
+	if (distinctClause == NIL || linitial(distinctClause) == NULL)
+		return distinctSortClause;
+
+	foreach(lc, distinctClause)
+	{
+		Node	   *key = (Node *) lfirst(lc);
+		SortBy	   *sb = makeNode(SortBy);
+
+		sb->node = key;
+		sb->sortby_dir = SORTBY_DEFAULT;
+		sb->sortby_nulls = SORTBY_NULLS_DEFAULT;
+		sb->useOp = NIL;
+		sb->location = -1;
+		prepended = lappend(prepended, sb);
+	}
+
+	return list_concat(prepended, result);
+}
+
 static Query *
 transformSelectStmt(ParseState *pstate, SelectStmt *stmt,
 					SelectStmtPassthrough *passthru)
@@ -1424,6 +1451,7 @@ transformSelectStmt(ParseState *pstate, SelectStmt *stmt,
 	Query	   *qry = makeNode(Query);
 	Node	   *qual;
 	ListCell   *l;
+	List	   *distinctSortClause = NIL;
 
 	qry->commandType = CMD_SELECT;
 
@@ -1496,6 +1524,17 @@ transformSelectStmt(ParseState *pstate, SelectStmt *stmt,
 											false /* allow SQL92 rules */ );
 	qry->groupDistinct = stmt->groupDistinct;
 
+	if (stmt->distinctSortClause)
+	{
+		List *full_sortby = prepend_distinct_to_sortby(stmt->distinctClause, stmt->distinctSortClause);
+		distinctSortClause = transformSortClause(pstate,
+												 full_sortby,
+												 &qry->targetList,
+												 EXPR_KIND_ORDER_BY,
+												 false);
+	}
+	qry->distinctSortClause = distinctSortClause;
+
 	if (stmt->distinctClause == NIL)
 	{
 		qry->distinctClause = NIL;
@@ -1516,7 +1555,7 @@ transformSelectStmt(ParseState *pstate, SelectStmt *stmt,
 		qry->distinctClause = transformDistinctOnClause(pstate,
 														stmt->distinctClause,
 														&qry->targetList,
-														qry->sortClause);
+														distinctSortClause ? distinctSortClause : qry->sortClause);
 		qry->hasDistinctOn = true;
 	}
 
diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 0563453fe24..a8bb0a4d8d5 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -13231,7 +13231,8 @@ simple_select:
 				{
 					SelectStmt *n = makeNode(SelectStmt);
 
-					n->distinctClause = $2;
+					n->distinctClause = linitial($2);
+					n->distinctSortClause = lsecond($2);
 					n->targetList = $3;
 					n->intoClause = $4;
 					n->fromClause = $5;
@@ -13485,8 +13486,9 @@ set_quantifier:
  * should be placed in the DISTINCT list during parsetree analysis.
  */
 distinct_clause:
-			DISTINCT								{ $$ = list_make1(NIL); }
-			| DISTINCT ON '(' expr_list ')'			{ $$ = $4; }
+			DISTINCT								{ $$ = list_make2(list_make1(NIL), NIL); }
+			| DISTINCT ON '(' expr_list ')'			{ $$ = list_make2($4, NIL); }
+			| DISTINCT ON '(' expr_list sort_clause ')' { $$ = list_make2($4, $5); }
 		;
 
 opt_all_clause:
@@ -17951,7 +17953,16 @@ PLpgSQL_Expr: opt_distinct_clause opt_target_list
 				{
 					SelectStmt *n = makeNode(SelectStmt);
 
-					n->distinctClause = $1;
+					if ($1)
+					{
+						n->distinctClause = linitial($1);
+						n->distinctSortClause = lsecond($1);
+					}
+					else
+					{
+						n->distinctClause = NIL;
+						n->distinctSortClause = NIL;
+					}
 					n->targetList = $2;
 					n->fromClause = $3;
 					n->whereClause = $4;
diff --git a/src/include/executor/executor.h b/src/include/executor/executor.h
index 8bb6c7bda2f..71cd51d0d46 100644
--- a/src/include/executor/executor.h
+++ b/src/include/executor/executor.h
@@ -153,6 +153,12 @@ extern uint32 TupleHashTableHash(TupleHashTable hashtable,
 extern TupleHashEntry LookupTupleHashEntryHash(TupleHashTable hashtable,
 											   TupleTableSlot *slot,
 											   bool *isnew, uint32 hash);
+extern bool ReplaceTupleHashEntryIfBetter(TupleHashTable hashtable,
+										  TupleHashEntry entry,
+										  TupleTableSlot *newslot,
+										  TupleTableSlot *firstslot,
+										  SortSupport sortKeys,
+										  int numSortCols);
 extern TupleHashEntry FindTupleHashEntry(TupleHashTable hashtable,
 										 TupleTableSlot *slot,
 										 ExprState *eqcomp,
diff --git a/src/include/executor/nodeAgg.h b/src/include/executor/nodeAgg.h
index 1e1be9666ae..e9f0cddec32 100644
--- a/src/include/executor/nodeAgg.h
+++ b/src/include/executor/nodeAgg.h
@@ -319,6 +319,8 @@ typedef struct AggStatePerHashData
 	AttrNumber *hashGrpColIdxInput; /* hash col indices in input slot */
 	AttrNumber *hashGrpColIdxHash;	/* indices in hash table tuples */
 	Agg		   *aggnode;		/* original Agg node, for numGroups etc. */
+	TupleTableSlot *hash_firstTupleSlot; /* slot for comparing stored tuples */
+	SortSupport  sortKeys;		/* sort support for distinct ON comparisons */
 } AggStatePerHashData;
 
 
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index aea7311faad..135445f78b0 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -226,6 +226,8 @@ typedef struct Query
 
 	List	   *distinctClause; /* a list of SortGroupClause's */
 
+	List	   *distinctSortClause; /* a list of SortGroupClause's for inline DISTINCT ON ORDER BY */
+
 	List	   *sortClause;		/* a list of SortGroupClause's */
 
 	Node	   *limitOffset;	/* # of result tuples to skip (int8 expr) */
@@ -2269,6 +2271,7 @@ typedef struct SelectStmt
 	struct SelectStmt *larg;	/* left child */
 	struct SelectStmt *rarg;	/* right child */
 	/* Eventually add fields for CORRESPONDING spec here */
+	List	   *distinctSortClause; /* inline DISTINCT ON ORDER BY clause (list of SortBy's) */
 } SelectStmt;
 
 
diff --git a/src/include/nodes/pathnodes.h b/src/include/nodes/pathnodes.h
index 1c6d1fe3d04..73e518ce501 100644
--- a/src/include/nodes/pathnodes.h
+++ b/src/include/nodes/pathnodes.h
@@ -2596,6 +2596,7 @@ typedef struct AggPath
 	Cardinality numGroups;		/* estimated number of groups in input */
 	uint64		transitionSpace;	/* for pass-by-ref transition data */
 	List	   *groupClause;	/* a list of SortGroupClause's */
+	List	   *distinctSortClause; /* a list of SortGroupClause's for inline DISTINCT ON ORDER BY */
 	List	   *qual;			/* quals (HAVING quals), if any */
 } AggPath;
 
diff --git a/src/include/nodes/plannodes.h b/src/include/nodes/plannodes.h
index 09a1ec73180..2c66ba830a9 100644
--- a/src/include/nodes/plannodes.h
+++ b/src/include/nodes/plannodes.h
@@ -1242,6 +1242,13 @@ typedef struct Agg
 	/* grouping sets to use */
 	List	   *groupingSets;
 
+	/* sort keys for inline DISTINCT ON ORDER BY (if any) */
+	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));
+
 	/* chained Agg/Sort nodes */
 	List	   *chain;
 } Agg;
diff --git a/src/include/optimizer/pathnode.h b/src/include/optimizer/pathnode.h
index da2d9b384b5..d514eb0c3ba 100644
--- a/src/include/optimizer/pathnode.h
+++ b/src/include/optimizer/pathnode.h
@@ -269,6 +269,17 @@ extern AggPath *create_agg_path(PlannerInfo *root,
 								List *qual,
 								const AggClauseCosts *aggcosts,
 								double numGroups);
+extern AggPath *create_agg_path_ext(PlannerInfo *root,
+									RelOptInfo *rel,
+									Path *subpath,
+									PathTarget *target,
+									AggStrategy aggstrategy,
+									AggSplit aggsplit,
+									List *groupClause,
+									List *qual,
+									const AggClauseCosts *aggcosts,
+									double numGroups,
+									List *distinctSortClause);
 extern GroupingSetsPath *create_groupingsets_path(PlannerInfo *root,
 												  RelOptInfo *rel,
 												  Path *subpath,
diff --git a/src/test/regress/expected/select_distinct_on.out b/src/test/regress/expected/select_distinct_on.out
index 75b1e7d300f..4c24a8b7fdb 100644
--- a/src/test/regress/expected/select_distinct_on.out
+++ b/src/test/regress/expected/select_distinct_on.out
@@ -246,3 +246,78 @@ SELECT DISTINCT ON (y, x) x, y FROM (select * from distinct_on_tbl order by x, z
 
 RESET enable_hashagg;
 DROP TABLE distinct_on_tbl;
+--
+-- Test SELECT DISTINCT ON with inline ORDER BY
+--
+CREATE TABLE distinct_inline_tbl (a int, b int);
+INSERT INTO distinct_inline_tbl VALUES (1, 10), (1, 20), (2, 5), (2, 15);
+EXPLAIN (COSTS OFF) SELECT DISTINCT ON (a ORDER BY b DESC) a, b FROM distinct_inline_tbl ORDER BY a;
+                 QUERY PLAN                  
+---------------------------------------------
+ Unique
+   ->  Sort
+         Sort Key: a, b DESC
+         ->  Seq Scan on distinct_inline_tbl
+(4 rows)
+
+SELECT DISTINCT ON (a ORDER BY b DESC) a, b FROM distinct_inline_tbl ORDER BY a;
+ a | b  
+---+----
+ 1 | 20
+ 2 | 15
+(2 rows)
+
+EXPLAIN (COSTS OFF) SELECT DISTINCT ON (a ORDER BY b ASC) a, b FROM distinct_inline_tbl ORDER BY a;
+                 QUERY PLAN                  
+---------------------------------------------
+ Unique
+   ->  Sort
+         Sort Key: a, b
+         ->  Seq Scan on distinct_inline_tbl
+(4 rows)
+
+SELECT DISTINCT ON (a ORDER BY b ASC) a, b FROM distinct_inline_tbl ORDER BY a;
+ a | b  
+---+----
+ 1 | 10
+ 2 |  5
+(2 rows)
+
+-- Test with HashAgg
+SET enable_sort TO OFF;
+EXPLAIN (COSTS OFF) SELECT DISTINCT ON (a ORDER BY b DESC) a, b FROM distinct_inline_tbl;
+                 QUERY PLAN                  
+---------------------------------------------
+ Unique
+   ->  Sort
+         Disabled: true
+         Sort Key: a, b DESC
+         ->  Seq Scan on distinct_inline_tbl
+(5 rows)
+
+SELECT DISTINCT ON (a ORDER BY b DESC) a, b FROM distinct_inline_tbl ORDER BY a;
+ a | b  
+---+----
+ 1 | 20
+ 2 | 15
+(2 rows)
+
+EXPLAIN (COSTS OFF) SELECT DISTINCT ON (a ORDER BY b ASC) a, b FROM distinct_inline_tbl;
+                 QUERY PLAN                  
+---------------------------------------------
+ Unique
+   ->  Sort
+         Disabled: true
+         Sort Key: a, b
+         ->  Seq Scan on distinct_inline_tbl
+(5 rows)
+
+SELECT DISTINCT ON (a ORDER BY b ASC) a, b FROM distinct_inline_tbl ORDER BY a;
+ a | b  
+---+----
+ 1 | 10
+ 2 |  5
+(2 rows)
+
+RESET enable_sort;
+DROP TABLE distinct_inline_tbl;
diff --git a/src/test/regress/sql/select_distinct_on.sql b/src/test/regress/sql/select_distinct_on.sql
index 8680749e49a..df844119938 100644
--- a/src/test/regress/sql/select_distinct_on.sql
+++ b/src/test/regress/sql/select_distinct_on.sql
@@ -82,3 +82,25 @@ SELECT DISTINCT ON (y, x) x, y FROM (select * from distinct_on_tbl order by x, z
 RESET enable_hashagg;
 
 DROP TABLE distinct_on_tbl;
+
+--
+-- Test SELECT DISTINCT ON with inline ORDER BY
+--
+CREATE TABLE distinct_inline_tbl (a int, b int);
+INSERT INTO distinct_inline_tbl VALUES (1, 10), (1, 20), (2, 5), (2, 15);
+
+EXPLAIN (COSTS OFF) SELECT DISTINCT ON (a ORDER BY b DESC) a, b FROM distinct_inline_tbl ORDER BY a;
+SELECT DISTINCT ON (a ORDER BY b DESC) a, b FROM distinct_inline_tbl ORDER BY a;
+EXPLAIN (COSTS OFF) SELECT DISTINCT ON (a ORDER BY b ASC) a, b FROM distinct_inline_tbl ORDER BY a;
+SELECT DISTINCT ON (a ORDER BY b ASC) a, b FROM distinct_inline_tbl ORDER BY a;
+
+-- Test with HashAgg
+SET enable_sort TO OFF;
+EXPLAIN (COSTS OFF) SELECT DISTINCT ON (a ORDER BY b DESC) a, b FROM distinct_inline_tbl;
+SELECT DISTINCT ON (a ORDER BY b DESC) a, b FROM distinct_inline_tbl ORDER BY a;
+EXPLAIN (COSTS OFF) SELECT DISTINCT ON (a ORDER BY b ASC) a, b FROM distinct_inline_tbl;
+SELECT DISTINCT ON (a ORDER BY b ASC) a, b FROM distinct_inline_tbl ORDER BY a;
+RESET enable_sort;
+
+DROP TABLE distinct_inline_tbl;
+
-- 
2.55.0.1082.g2b9226bbc0-goog

Reply via email to