From 1b78737485daf994949995655fa4898672276c4c Mon Sep 17 00:00:00 2001
From: Sergey Soloviev <sergey.soloviev@tantorlabs.ru>
Date: Thu, 11 Dec 2025 14:30:37 +0300
Subject: [PATCH v3 4/5] add support for Partial IndexAggregate

Now IndexAggregate support partial aggregates. The main problem was with
partial aggregates which creates SortGroupClause for same expression as
in target list, but different sortgroupclause, so make_pathkeys_for_sortclases
failed to find required target list entry and throws ERROR.

To fix this now we explicitly pass pathkeys to create_agg_path (but only
for AGG_INDEX for now), so caller is responsible for searching and
building pathkeys list.
---
 src/backend/optimizer/path/allpaths.c  | 76 ++++++++++++++++++++
 src/backend/optimizer/plan/planner.c   | 98 ++++++++++++++++++++++++--
 src/backend/optimizer/prep/prepunion.c |  2 +
 src/backend/optimizer/util/pathnode.c  | 16 +++--
 src/include/optimizer/pathnode.h       |  1 +
 5 files changed, 185 insertions(+), 8 deletions(-)

diff --git a/src/backend/optimizer/path/allpaths.c b/src/backend/optimizer/path/allpaths.c
index 4c43fd0b19b..5fcac30af84 100644
--- a/src/backend/optimizer/path/allpaths.c
+++ b/src/backend/optimizer/path/allpaths.c
@@ -3446,6 +3446,7 @@ generate_grouped_paths(PlannerInfo *root, RelOptInfo *grouped_rel,
 	AggClauseCosts agg_costs;
 	bool		can_hash;
 	bool		can_sort;
+	bool		can_index;
 	Path	   *cheapest_total_path = NULL;
 	Path	   *cheapest_partial_path = NULL;
 	double		dNumGroups = 0;
@@ -3498,6 +3499,12 @@ generate_grouped_paths(PlannerInfo *root, RelOptInfo *grouped_rel,
 	can_hash = (agg_info->group_clauses != NIL &&
 				grouping_is_hashable(agg_info->group_clauses));
 
+	/* 
+	 * Determine whether we should consider index-based implementations of
+	 * grouping.
+	 */
+	can_index = can_sort && can_hash;
+
 	/*
 	 * Consider whether we should generate partially aggregated non-partial
 	 * paths.  We can only do this if we have a non-partial path.
@@ -3615,6 +3622,7 @@ generate_grouped_paths(PlannerInfo *root, RelOptInfo *grouped_rel,
 											AGGSPLIT_INITIAL_SERIAL,
 											agg_info->group_clauses,
 											NIL,
+											NIL,
 											&agg_costs,
 											dNumGroups);
 
@@ -3691,6 +3699,7 @@ generate_grouped_paths(PlannerInfo *root, RelOptInfo *grouped_rel,
 											AGGSPLIT_INITIAL_SERIAL,
 											agg_info->group_clauses,
 											NIL,
+											NIL,
 											&agg_costs,
 											dNumPartialGroups);
 
@@ -3727,6 +3736,7 @@ generate_grouped_paths(PlannerInfo *root, RelOptInfo *grouped_rel,
 										AGGSPLIT_INITIAL_SERIAL,
 										agg_info->group_clauses,
 										NIL,
+										NIL,
 										&agg_costs,
 										dNumGroups);
 
@@ -3762,6 +3772,72 @@ generate_grouped_paths(PlannerInfo *root, RelOptInfo *grouped_rel,
 										AGGSPLIT_INITIAL_SERIAL,
 										agg_info->group_clauses,
 										NIL,
+										NIL,
+										&agg_costs,
+										dNumPartialGroups);
+
+		add_partial_path(grouped_rel, path);
+	}
+
+	if (can_index && cheapest_total_path != NULL)
+	{
+		Path	   *path;
+
+		/*
+		 * Since the path originates from a non-grouped relation that is
+		 * not aware of eager aggregation, we must ensure that it provides
+		 * the correct input for partial aggregation.
+		 */
+		path = (Path *) create_projection_path(root,
+											   grouped_rel,
+											   cheapest_total_path,
+											   agg_info->agg_input);
+		/*
+		 * qual is NIL because the HAVING clause cannot be evaluated until the
+		 * final value of the aggregate is known.
+		 */
+		path = (Path *) create_agg_path(root,
+										grouped_rel,
+										path,
+										agg_info->target,
+										AGG_INDEX,
+										AGGSPLIT_INITIAL_SERIAL,
+										agg_info->group_clauses,
+										NIL,
+										group_pathkeys,
+										&agg_costs,
+										dNumGroups);
+
+		add_path(grouped_rel, path);
+	}
+
+	if (can_index && cheapest_partial_path != NULL)
+	{
+		Path	   *path;
+
+		/*
+		 * Since the path originates from a non-grouped relation that is not
+		 * aware of eager aggregation, we must ensure that it provides the
+		 * correct input for partial aggregation.
+		 */
+		path = (Path *) create_projection_path(root,
+											   grouped_rel,
+											   cheapest_partial_path,
+											   agg_info->agg_input);
+
+		/*
+		 * qual is NIL because the HAVING clause cannot be evaluated until the
+		 * final value of the aggregate is known.
+		 */
+		path = (Path *) create_agg_path(root,
+										grouped_rel,
+										path,
+										agg_info->target,
+										AGG_INDEX,
+										AGGSPLIT_INITIAL_SERIAL,
+										agg_info->group_clauses,
+										NIL,
+										group_pathkeys,
 										&agg_costs,
 										dNumPartialGroups);
 
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index cfd2f3ff3a9..5598943d9f9 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -3888,6 +3888,7 @@ create_grouping_paths(PlannerInfo *root,
 		if (   gd == NULL
 			&& root->numOrderedAggs == 0
 			&& parse->groupClause != NIL
+			&& parse->groupingSets == NIL
 			&& grouping_is_sortable(root->processed_groupClause)
 			&& grouping_is_hashable(root->processed_groupClause))
 			flags |= GROUPING_CAN_USE_INDEX;
@@ -5030,6 +5031,7 @@ create_partial_distinct_paths(PlannerInfo *root, RelOptInfo *input_rel,
 										 AGGSPLIT_SIMPLE,
 										 root->processed_distinctClause,
 										 NIL,
+										 NIL,
 										 NULL,
 										 numDistinctRows));
 	}
@@ -5238,6 +5240,7 @@ create_final_distinct_paths(PlannerInfo *root, RelOptInfo *input_rel,
 								 AGGSPLIT_SIMPLE,
 								 root->processed_distinctClause,
 								 NIL,
+								 NIL,
 								 NULL,
 								 numDistinctRows));
 	}
@@ -7209,6 +7212,7 @@ add_paths_to_grouping_rel(PlannerInfo *root, RelOptInfo *input_rel,
 											 AGGSPLIT_SIMPLE,
 											 info->clauses,
 											 havingQual,
+											 NIL,
 											 agg_costs,
 											 dNumGroups));
 				}
@@ -7280,6 +7284,7 @@ add_paths_to_grouping_rel(PlannerInfo *root, RelOptInfo *input_rel,
 												 AGGSPLIT_FINAL_DESERIAL,
 												 info->clauses,
 												 havingQual,
+												 NIL,
 												 agg_final_costs,
 												 dNumFinalGroups));
 					else
@@ -7321,6 +7326,7 @@ add_paths_to_grouping_rel(PlannerInfo *root, RelOptInfo *input_rel,
 									 AGGSPLIT_SIMPLE,
 									 root->processed_groupClause,
 									 havingQual,
+									 NIL,
 									 agg_costs,
 									 dNumGroups));
 		}
@@ -7340,6 +7346,7 @@ add_paths_to_grouping_rel(PlannerInfo *root, RelOptInfo *input_rel,
 									 AGGSPLIT_FINAL_DESERIAL,
 									 root->processed_groupClause,
 									 havingQual,
+									 NIL,
 									 agg_final_costs,
 									 dNumFinalGroups));
 		}
@@ -7347,10 +7354,10 @@ add_paths_to_grouping_rel(PlannerInfo *root, RelOptInfo *input_rel,
 
 	if (can_index)
 	{
-		/* 
-		 * Generate IndexAgg path.
-		 */
-		Assert(!parse->groupingSets);
+		List *pathkeys = make_pathkeys_for_sortclauses(root,
+													   root->processed_groupClause,
+													   root->processed_tlist);
+
 		add_path(grouped_rel, (Path *)
 				 create_agg_path(root,
 								 grouped_rel,
@@ -7360,8 +7367,29 @@ add_paths_to_grouping_rel(PlannerInfo *root, RelOptInfo *input_rel,
 								 AGGSPLIT_SIMPLE,
 								 root->processed_groupClause,
 								 havingQual,
+								 pathkeys,
 								 agg_costs,
 								 dNumGroups));
+		
+		/*
+		 * Instead of operating directly on the input relation, we can
+		 * consider finalizing a partially aggregated path.
+		 */
+		if (partially_grouped_rel != NULL)
+		{
+			add_path(grouped_rel, (Path *)
+					 create_agg_path(root,
+									 grouped_rel,
+									 cheapest_partially_grouped_path,
+									 grouped_rel->reltarget,
+									 AGG_INDEX,
+									 AGGSPLIT_FINAL_DESERIAL,
+									 root->processed_groupClause,
+									 havingQual,
+									 pathkeys,
+									 agg_final_costs,
+									 dNumFinalGroups));
+		}
 	}
 
 	/*
@@ -7410,6 +7438,7 @@ create_partial_grouping_paths(PlannerInfo *root,
 	ListCell   *lc;
 	bool		can_hash = (extra->flags & GROUPING_CAN_USE_HASH) != 0;
 	bool		can_sort = (extra->flags & GROUPING_CAN_USE_SORT) != 0;
+	bool		can_index = (extra->flags & GROUPING_CAN_USE_INDEX) != 0;
 
 	/*
 	 * Check whether any partially aggregated paths have been generated
@@ -7561,6 +7590,7 @@ create_partial_grouping_paths(PlannerInfo *root,
 											 AGGSPLIT_INITIAL_SERIAL,
 											 info->clauses,
 											 NIL,
+											 NIL,
 											 agg_partial_costs,
 											 dNumPartialGroups));
 				else
@@ -7619,6 +7649,7 @@ create_partial_grouping_paths(PlannerInfo *root,
 													 AGGSPLIT_INITIAL_SERIAL,
 													 info->clauses,
 													 NIL,
+													 NIL,
 													 agg_partial_costs,
 													 dNumPartialPartialGroups));
 				else
@@ -7650,6 +7681,7 @@ create_partial_grouping_paths(PlannerInfo *root,
 								 AGGSPLIT_INITIAL_SERIAL,
 								 root->processed_groupClause,
 								 NIL,
+								 NIL,
 								 agg_partial_costs,
 								 dNumPartialGroups));
 	}
@@ -7668,6 +7700,62 @@ create_partial_grouping_paths(PlannerInfo *root,
 										 AGGSPLIT_INITIAL_SERIAL,
 										 root->processed_groupClause,
 										 NIL,
+										 NIL,
+										 agg_partial_costs,
+										 dNumPartialPartialGroups));
+	}
+	
+	/*
+	 * Add a partially-grouped IndexAgg Path where possible
+	 */
+	if (can_index && cheapest_total_path != NULL)
+	{
+		List *pathkeys;
+
+		/* This should have been checked previously */
+		Assert(parse->hasAggs || parse->groupClause);
+		
+		pathkeys = make_pathkeys_for_sortclauses(root,
+												 root->processed_groupClause,
+												 root->processed_tlist);
+
+		add_path(partially_grouped_rel, (Path *)
+				 create_agg_path(root,
+								 partially_grouped_rel,
+								 cheapest_total_path,
+								 partially_grouped_rel->reltarget,
+								 AGG_INDEX,
+								 AGGSPLIT_INITIAL_SERIAL,
+								 root->processed_groupClause,
+								 NIL,
+								 pathkeys,
+								 agg_partial_costs,
+								 dNumPartialGroups));
+	}
+
+	/*
+	 * Now add a partially-grouped IndexAgg partial Path where possible
+	 */
+	if (can_index && cheapest_partial_path != NULL)
+	{
+		List *pathkeys;
+
+		/* This should have been checked previously */
+		Assert(parse->hasAggs || parse->groupClause);
+		
+		pathkeys = make_pathkeys_for_sortclauses(root,
+												 root->processed_groupClause,
+												 root->processed_tlist);
+		add_partial_path(partially_grouped_rel, (Path *)
+						  create_agg_path(root,
+										 partially_grouped_rel,
+										 cheapest_partial_path,
+										 partially_grouped_rel->reltarget,
+										 AGG_INDEX,
+										 AGGSPLIT_INITIAL_SERIAL,
+										 root->processed_groupClause,
+										 NIL,
+										 pathkeys,
 										 agg_partial_costs,
 										 dNumPartialPartialGroups));
 	}
@@ -8829,6 +8917,7 @@ create_final_unique_paths(PlannerInfo *root, RelOptInfo *input_rel,
 										AGGSPLIT_SIMPLE,
 										groupClause,
 										NIL,
+										NIL,
 										NULL,
 										unique_rel->rows);
 
@@ -8971,6 +9060,7 @@ create_partial_unique_paths(PlannerInfo *root, RelOptInfo *input_rel,
 										AGGSPLIT_SIMPLE,
 										groupClause,
 										NIL,
+										NIL,
 										NULL,
 										partial_unique_rel->rows);
 
diff --git a/src/backend/optimizer/prep/prepunion.c b/src/backend/optimizer/prep/prepunion.c
index a01b02f3a7b..de6a1558044 100644
--- a/src/backend/optimizer/prep/prepunion.c
+++ b/src/backend/optimizer/prep/prepunion.c
@@ -949,6 +949,7 @@ generate_union_paths(SetOperationStmt *op, PlannerInfo *root,
 											AGGSPLIT_SIMPLE,
 											groupList,
 											NIL,
+											NIL,
 											NULL,
 											dNumGroups);
 			add_path(result_rel, path);
@@ -965,6 +966,7 @@ generate_union_paths(SetOperationStmt *op, PlannerInfo *root,
 												AGGSPLIT_SIMPLE,
 												groupList,
 												NIL,
+												NIL,
 												NULL,
 												dNumGroups);
 				add_path(result_rel, path);
diff --git a/src/backend/optimizer/util/pathnode.c b/src/backend/optimizer/util/pathnode.c
index 2bac26055a7..646762be43b 100644
--- a/src/backend/optimizer/util/pathnode.c
+++ b/src/backend/optimizer/util/pathnode.c
@@ -2988,6 +2988,7 @@ create_unique_path(PlannerInfo *root,
  * 'aggsplit' is the Agg node's aggregate-splitting mode
  * 'groupClause' is a list of SortGroupClause's representing the grouping
  * 'qual' is the HAVING quals if any
+ * 'pathkeys' for AGG_INDEX must be a list of PathKey used by this agg node
  * 'aggcosts' contains cost info about the aggregate functions to be computed
  * 'numGroups' is the estimated number of groups (1 if not grouping)
  */
@@ -3000,6 +3001,7 @@ create_agg_path(PlannerInfo *root,
 				AggSplit aggsplit,
 				List *groupClause,
 				List *qual,
+				List *pathkeys,
 				const AggClauseCosts *aggcosts,
 				double numGroups)
 {
@@ -3033,11 +3035,17 @@ create_agg_path(PlannerInfo *root,
 	else if (aggstrategy == AGG_INDEX)
 	{
 		/* 
-		 * When using index aggregation all grouping columns will be used as
-		 * comparator keys, so output is always sorted.
+		 * For IndexAgg we also must know used ordering just like for GroupAgg,
+		 * but for the latter this information is passed by child node, i.e.
+		 * Sort. But here we can not use make_pathkeys_for_sortclauses, because
+		 * in case of partial aggregates the node will contain different target
+		 * list and sortgroupref indexes, so this function will not find required
+		 * entries. So caller must build pathkeys for us.
+		 * 
+		 * NOTE: pathkeys CAN be NIL, i.e. if planner decided that all values
+		 * are same constant.
 		 */
-		pathnode->path.pathkeys = make_pathkeys_for_sortclauses(root, groupClause,
-																root->processed_tlist);
+		pathnode->path.pathkeys = pathkeys;
 	}
 	else
 		pathnode->path.pathkeys = NIL;	/* output is unordered */
diff --git a/src/include/optimizer/pathnode.h b/src/include/optimizer/pathnode.h
index 6b010f0b1a5..a2aad4ecba7 100644
--- a/src/include/optimizer/pathnode.h
+++ b/src/include/optimizer/pathnode.h
@@ -235,6 +235,7 @@ extern AggPath *create_agg_path(PlannerInfo *root,
 								AggSplit aggsplit,
 								List *groupClause,
 								List *qual,
+								List *pathkeys,
 								const AggClauseCosts *aggcosts,
 								double numGroups);
 extern GroupingSetsPath *create_groupingsets_path(PlannerInfo *root,
-- 
2.43.0

