From 429cab42ee94a88eef79dfb3575ded35b8056a1c Mon Sep 17 00:00:00 2001
From: Richard Guo <guofenglinux@gmail.com>
Date: Fri, 23 Feb 2024 13:33:09 +0800
Subject: [PATCH v1 7/9] Build grouped relations out of join relations

This commit builds grouped relations for each just-processed join
relation if possible, and generates aggregation paths for the grouped
join relations.

If we are joining rel1 and rel2, the aggregation paths for the grouped
join relation are generated by 1) joining the grouped paths of rel1 to
the plain paths of rel2, or joining the grouped paths of rel2 to the
plain paths of rel1, and 2) adding sorted and hashed partial aggregation
paths on top of paths of the plain join rel except for the topmost join
rel.

This commit also makes the grouped relation for the topmost join rel act
as the upper rel representing the result of partial aggregation, so that
we can add the final aggregation on top of that.

This commit also makes eager aggregation work for partitionwise join and
for geqo.

Starting from this commit, you should be able to see plans with eager
aggregation.
---
 src/backend/optimizer/geqo/geqo_eval.c  |  84 +++++++++++++----
 src/backend/optimizer/path/allpaths.c   |  48 ++++++++++
 src/backend/optimizer/path/joinrels.c   | 115 ++++++++++++++++++++++++
 src/backend/optimizer/plan/planner.c    |  35 ++++++--
 src/backend/optimizer/util/appendinfo.c |  64 +++++++++++++
 5 files changed, 320 insertions(+), 26 deletions(-)

diff --git a/src/backend/optimizer/geqo/geqo_eval.c b/src/backend/optimizer/geqo/geqo_eval.c
index 1141156899..278857d767 100644
--- a/src/backend/optimizer/geqo/geqo_eval.c
+++ b/src/backend/optimizer/geqo/geqo_eval.c
@@ -60,8 +60,12 @@ geqo_eval(PlannerInfo *root, Gene *tour, int num_gene)
 	MemoryContext oldcxt;
 	RelOptInfo *joinrel;
 	Cost		fitness;
-	int			savelength;
-	struct HTAB *savehash;
+	int			savelength_join_rel;
+	struct HTAB *savehash_join_rel;
+	int			savelength_grouped_rel;
+	struct HTAB *savehash_grouped_rel;
+	int			savelength_grouped_info;
+	struct HTAB *savehash_grouped_info;
 
 	/*
 	 * Create a private memory context that will hold all temp storage
@@ -78,25 +82,38 @@ geqo_eval(PlannerInfo *root, Gene *tour, int num_gene)
 	oldcxt = MemoryContextSwitchTo(mycontext);
 
 	/*
-	 * gimme_tree will add entries to root->join_rel_list, which may or may
-	 * not already contain some entries.  The newly added entries will be
-	 * recycled by the MemoryContextDelete below, so we must ensure that the
-	 * list is restored to its former state before exiting.  We can do this by
-	 * truncating the list to its original length.  NOTE this assumes that any
-	 * added entries are appended at the end!
+	 * gimme_tree will add entries to root->join_rel_list, root->agg_info_list
+	 * and root->upper_rels[UPPERREL_PARTIAL_GROUP_AGG], which may or may not
+	 * already contain some entries.  The newly added entries will be recycled
+	 * by the MemoryContextDelete below, so we must ensure that each list of
+	 * the RelInfoList structures is restored to its former state before
+	 * exiting.  We can do this by truncating each list to its original length.
+	 * NOTE this assumes that any added entries are appended at the end!
 	 *
-	 * We also must take care not to mess up the outer join_rel_list->hash, if
-	 * there is one.  We can do this by just temporarily setting the link to
-	 * NULL.  (If we are dealing with enough join rels, which we very likely
-	 * are, a new hash table will get built and used locally.)
+	 * We also must take care not to mess up the outer hash tables of the
+	 * RelInfoList structures, if any.  We can do this by just temporarily
+	 * setting each link to NULL.  (If we are dealing with enough join rels,
+	 * which we very likely are, new hash tables will get built and used
+	 * locally.)
 	 *
 	 * join_rel_level[] shouldn't be in use, so just Assert it isn't.
 	 */
-	savelength = list_length(root->join_rel_list->items);
-	savehash = root->join_rel_list->hash;
+	savelength_join_rel = list_length(root->join_rel_list->items);
+	savehash_join_rel = root->join_rel_list->hash;
+
+	savelength_grouped_rel =
+		list_length(root->upper_rels[UPPERREL_PARTIAL_GROUP_AGG].items);
+	savehash_grouped_rel =
+		root->upper_rels[UPPERREL_PARTIAL_GROUP_AGG].hash;
+
+	savelength_grouped_info = list_length(root->agg_info_list->items);
+	savehash_grouped_info = root->agg_info_list->hash;
+
 	Assert(root->join_rel_level == NULL);
 
 	root->join_rel_list->hash = NULL;
+	root->upper_rels[UPPERREL_PARTIAL_GROUP_AGG].hash = NULL;
+	root->agg_info_list->hash = NULL;
 
 	/* construct the best path for the given combination of relations */
 	joinrel = gimme_tree(root, tour, num_gene);
@@ -118,12 +135,22 @@ geqo_eval(PlannerInfo *root, Gene *tour, int num_gene)
 		fitness = DBL_MAX;
 
 	/*
-	 * Restore join_rel_list to its former state, and put back original
-	 * hashtable if any.
+	 * Restore each of the list in join_rel_list, agg_info_list and
+	 * upper_rels[UPPERREL_PARTIAL_GROUP_AGG] to its former state, and put back
+	 * original hashtable if any.
 	 */
 	root->join_rel_list->items = list_truncate(root->join_rel_list->items,
-											   savelength);
-	root->join_rel_list->hash = savehash;
+											   savelength_join_rel);
+	root->join_rel_list->hash = savehash_join_rel;
+
+	root->upper_rels[UPPERREL_PARTIAL_GROUP_AGG].items =
+		list_truncate(root->upper_rels[UPPERREL_PARTIAL_GROUP_AGG].items,
+					  savelength_grouped_rel);
+	root->upper_rels[UPPERREL_PARTIAL_GROUP_AGG].hash = savehash_grouped_rel;
+
+	root->agg_info_list->items = list_truncate(root->agg_info_list->items,
+											   savelength_grouped_info);
+	root->agg_info_list->hash = savehash_grouped_info;
 
 	/* release all the memory acquired within gimme_tree */
 	MemoryContextSwitchTo(oldcxt);
@@ -279,6 +306,27 @@ merge_clump(PlannerInfo *root, List *clumps, Clump *new_clump, int num_gene,
 				/* Find and save the cheapest paths for this joinrel */
 				set_cheapest(joinrel);
 
+				/*
+				 * Except for the topmost scan/join rel, consider generating
+				 * partial aggregation paths for the grouped relation on top of the
+				 * paths of this rel.  After that, we're done creating paths for
+				 * the grouped relation, so run set_cheapest().
+				 */
+				if (!bms_equal(joinrel->relids, root->all_query_rels))
+				{
+					RelOptInfo   *rel_grouped;
+					RelAggInfo   *agg_info;
+
+					rel_grouped = find_grouped_rel(root, joinrel->relids,
+												   &agg_info);
+					if (rel_grouped)
+					{
+						generate_grouped_paths(root, rel_grouped, joinrel,
+											   agg_info);
+						set_cheapest(rel_grouped);
+					}
+				}
+
 				/* Absorb new clump into old */
 				old_clump->joinrel = joinrel;
 				old_clump->size += new_clump->size;
diff --git a/src/backend/optimizer/path/allpaths.c b/src/backend/optimizer/path/allpaths.c
index ea2341d110..440a5daec7 100644
--- a/src/backend/optimizer/path/allpaths.c
+++ b/src/backend/optimizer/path/allpaths.c
@@ -3864,6 +3864,10 @@ standard_join_search(PlannerInfo *root, int levels_needed, List *initial_rels)
 		 *
 		 * After that, we're done creating paths for the joinrel, so run
 		 * set_cheapest().
+		 *
+		 * In addition, we also run generate_grouped_paths() for the grouped
+		 * relation of each just-processed joinrel, and run set_cheapest() for
+		 * the grouped relation afterwards.
 		 */
 		foreach(lc, root->join_rel_level[lev])
 		{
@@ -3884,6 +3888,27 @@ standard_join_search(PlannerInfo *root, int levels_needed, List *initial_rels)
 			/* Find and save the cheapest paths for this rel */
 			set_cheapest(rel);
 
+			/*
+			 * Except for the topmost scan/join rel, consider generating
+			 * partial aggregation paths for the grouped relation on top of the
+			 * paths of this rel.  After that, we're done creating paths for
+			 * the grouped relation, so run set_cheapest().
+			 */
+			if (!bms_equal(rel->relids, root->all_query_rels))
+			{
+				RelOptInfo   *rel_grouped;
+				RelAggInfo   *agg_info;
+
+				rel_grouped = find_grouped_rel(root, rel->relids,
+											   &agg_info);
+				if (rel_grouped)
+				{
+					generate_grouped_paths(root, rel_grouped, rel,
+										   agg_info);
+					set_cheapest(rel_grouped);
+				}
+			}
+
 #ifdef OPTIMIZER_DEBUG
 			pprint(rel);
 #endif
@@ -4752,6 +4777,29 @@ generate_partitionwise_join_paths(PlannerInfo *root, RelOptInfo *rel)
 		if (IS_DUMMY_REL(child_rel))
 			continue;
 
+		/*
+		 * Except for the topmost scan/join rel, consider generating partial
+		 * aggregation paths for the grouped relation on top of the paths of
+		 * this partitioned child-join.  After that, we're done creating paths
+		 * for the grouped relation, so run set_cheapest().
+		 */
+		if (!bms_equal(IS_OTHER_REL(rel) ?
+					   rel->top_parent_relids : rel->relids,
+					   root->all_query_rels))
+		{
+			RelOptInfo   *rel_grouped;
+			RelAggInfo   *agg_info;
+
+			rel_grouped = find_grouped_rel(root, child_rel->relids,
+										   &agg_info);
+			if (rel_grouped)
+			{
+				generate_grouped_paths(root, rel_grouped, child_rel,
+									   agg_info);
+				set_cheapest(rel_grouped);
+			}
+		}
+
 #ifdef OPTIMIZER_DEBUG
 		pprint(child_rel);
 #endif
diff --git a/src/backend/optimizer/path/joinrels.c b/src/backend/optimizer/path/joinrels.c
index 4750579b0a..a9ef081597 100644
--- a/src/backend/optimizer/path/joinrels.c
+++ b/src/backend/optimizer/path/joinrels.c
@@ -16,11 +16,13 @@
 
 #include "miscadmin.h"
 #include "optimizer/appendinfo.h"
+#include "optimizer/cost.h"
 #include "optimizer/joininfo.h"
 #include "optimizer/pathnode.h"
 #include "optimizer/paths.h"
 #include "partitioning/partbounds.h"
 #include "utils/memutils.h"
+#include "utils/selfuncs.h"
 
 
 static void make_rels_by_clause_joins(PlannerInfo *root,
@@ -35,6 +37,9 @@ static bool has_legal_joinclause(PlannerInfo *root, RelOptInfo *rel);
 static bool restriction_is_constant_false(List *restrictlist,
 										  RelOptInfo *joinrel,
 										  bool only_pushed_down);
+static void make_grouped_join_rel(PlannerInfo *root, RelOptInfo *rel1,
+								  RelOptInfo *rel2, RelOptInfo *joinrel,
+								  SpecialJoinInfo *sjinfo, List *restrictlist);
 static void populate_joinrel_with_paths(PlannerInfo *root, RelOptInfo *rel1,
 										RelOptInfo *rel2, RelOptInfo *joinrel,
 										SpecialJoinInfo *sjinfo, List *restrictlist);
@@ -753,6 +758,10 @@ make_join_rel(PlannerInfo *root, RelOptInfo *rel1, RelOptInfo *rel2)
 		return joinrel;
 	}
 
+	/* Build a grouped join relation for 'joinrel' if possible. */
+	make_grouped_join_rel(root, rel1, rel2, joinrel, sjinfo,
+						  restrictlist);
+
 	/* Add paths to the join relation. */
 	populate_joinrel_with_paths(root, rel1, rel2, joinrel, sjinfo,
 								restrictlist);
@@ -864,6 +873,107 @@ add_outer_joins_to_relids(PlannerInfo *root, Relids input_relids,
 	return input_relids;
 }
 
+/*
+ * make_grouped_join_rel
+ *	  Build a grouped join relation out of 'joinrel' if eager aggregation is
+ *	  possible and the 'joinrel' can produce grouped paths.
+ *
+ * We also generate partial aggregation paths for the grouped relation by
+ * joining the grouped paths of 'rel1' to the plain paths of 'rel2', or by
+ * joining the grouped paths of 'rel2' to the plain paths of 'rel1'.
+ */
+static void
+make_grouped_join_rel(PlannerInfo *root, RelOptInfo *rel1,
+					  RelOptInfo *rel2, RelOptInfo *joinrel,
+					  SpecialJoinInfo *sjinfo, List *restrictlist)
+{
+	Relids			joinrelids;
+	RelOptInfo	   *rel_grouped;
+	RelAggInfo	   *agg_info = NULL;
+	RelOptInfo	   *rel1_grouped;
+	RelOptInfo	   *rel2_grouped;
+
+	/*
+	 * If there are no aggregate expressions or grouping expressions, eager
+	 * aggregation is not possible.
+	 */
+	if (root->agg_clause_list == NIL ||
+		root->group_expr_list == NIL)
+		return;
+
+	joinrelids = bms_union(rel1->relids, rel2->relids);
+	rel_grouped = find_grouped_rel(root, joinrelids, &agg_info);
+
+	/*
+	 * Construct a new RelOptInfo for the grouped join relation if there is no
+	 * existing one.
+	 */
+	if (rel_grouped == NULL)
+	{
+		/*
+		 * Prepare the information we need to create grouped paths for this
+		 * join relation.
+		 */
+		agg_info = create_rel_agg_info(root, joinrel);
+		if (agg_info == NULL)
+			return;
+
+		/* build a grouped relation out of the plain relation */
+		rel_grouped = build_grouped_rel(root, joinrel);
+		rel_grouped->reltarget = agg_info->target;
+		rel_grouped->rows = agg_info->grouped_rows;
+
+		/*
+		 * Make the grouped relation available for further joining or for
+		 * acting as the upper rel representing the result of partial
+		 * aggregation.
+		 */
+		add_grouped_rel(root, rel_grouped, agg_info);
+	}
+
+	Assert(agg_info != NULL);
+
+	/* retrieve the grouped relations for the two input rels */
+	rel1_grouped = find_grouped_rel(root, rel1->relids, NULL);
+	rel2_grouped = find_grouped_rel(root, rel2->relids, NULL);
+
+	/* we should not see dummy grouped relation */
+	Assert(rel1_grouped == NULL || !IS_DUMMY_REL(rel1_grouped));
+	Assert(rel2_grouped == NULL || !IS_DUMMY_REL(rel2_grouped));
+
+	/* Nothing to do if there's no grouped relation. */
+	if (rel1_grouped == NULL &&
+		rel2_grouped == NULL)
+		return;
+
+	/*
+	 * Join of two grouped relations is currently not supported.  In such a
+	 * case, grouping of one side would change the occurrence of the other
+	 * side's aggregate transient states on the input of the final aggregation.
+	 * This can be handled by adjusting the transient states, but it's not
+	 * worth the effort for now.
+	 */
+	if (rel1_grouped != NULL &&
+		rel2_grouped != NULL)
+		return;
+
+	/* generate partial aggregation paths for the grouped relation */
+	if (rel1_grouped != NULL)
+	{
+		set_joinrel_size_estimates(root, rel_grouped, rel1_grouped, rel2,
+								   sjinfo, restrictlist);
+		populate_joinrel_with_paths(root, rel1_grouped, rel2, rel_grouped,
+									sjinfo, restrictlist);
+	}
+	else if (rel2_grouped != NULL)
+	{
+		set_joinrel_size_estimates(root, rel_grouped, rel1, rel2_grouped,
+								   sjinfo, restrictlist);
+		populate_joinrel_with_paths(root, rel1, rel2_grouped, rel_grouped,
+									sjinfo, restrictlist);
+	}
+}
+
 /*
  * populate_joinrel_with_paths
  *	  Add paths to the given joinrel for given pair of joining relations. The
@@ -1653,6 +1763,11 @@ try_partitionwise_join(PlannerInfo *root, RelOptInfo *rel1, RelOptInfo *rel2,
 						 adjust_child_relids(joinrel->relids,
 											 nappinfos, appinfos)));
 
+		/* Build a grouped join relation for 'child_joinrel' if possible */
+		make_grouped_join_rel(root, child_rel1, child_rel2,
+							  child_joinrel, child_sjinfo,
+							  child_restrictlist);
+
 		/* And make paths for the child join */
 		populate_joinrel_with_paths(root, child_rel1, child_rel2,
 									child_joinrel, child_sjinfo,
diff --git a/src/backend/optimizer/plan/planner.c b/src/backend/optimizer/plan/planner.c
index be4e182869..f8f2a09f1b 100644
--- a/src/backend/optimizer/plan/planner.c
+++ b/src/backend/optimizer/plan/planner.c
@@ -3946,10 +3946,16 @@ create_ordinary_grouping_paths(PlannerInfo *root, RelOptInfo *input_rel,
 
 	/* Gather any partially grouped partial paths. */
 	if (partially_grouped_rel && partially_grouped_rel->partial_pathlist)
-	{
 		gather_grouping_paths(root, partially_grouped_rel);
+
+	/*
+	 * Now choose the best path(s) for partially_grouped_rel.
+	 *
+	 * Note that the non-partial paths can come either from the Gather above or
+	 * from eager aggregation.
+	 */
+	if (partially_grouped_rel && partially_grouped_rel->pathlist)
 		set_cheapest(partially_grouped_rel);
-	}
 
 	/*
 	 * Estimate number of groups.
@@ -7043,6 +7049,13 @@ create_partial_grouping_paths(PlannerInfo *root,
 	bool		can_hash = (extra->flags & GROUPING_CAN_USE_HASH) != 0;
 	bool		can_sort = (extra->flags & GROUPING_CAN_USE_SORT) != 0;
 
+	/*
+	 * The partially_grouped_rel could have been already created due to eager
+	 * aggregation.
+	 */
+	partially_grouped_rel = find_grouped_rel(root, input_rel->relids, NULL);
+	Assert(enable_eager_aggregate || partially_grouped_rel == NULL);
+
 	/*
 	 * Consider whether we should generate partially aggregated non-partial
 	 * paths.  We can only do this if we have a non-partial path, and only if
@@ -7066,19 +7079,25 @@ create_partial_grouping_paths(PlannerInfo *root,
 	 * If we can't partially aggregate partial paths, and we can't partially
 	 * aggregate non-partial paths, then don't bother creating the new
 	 * RelOptInfo at all, unless the caller specified force_rel_creation.
+	 *
+	 * Note that the partially_grouped_rel could have been already created and
+	 * populated with appropriate paths by eager aggregation.
 	 */
 	if (cheapest_total_path == NULL &&
 		cheapest_partial_path == NULL &&
 		!force_rel_creation)
-		return NULL;
+		return partially_grouped_rel;
 
 	/*
 	 * Build a new upper relation to represent the result of partially
-	 * aggregating the rows from the input relation.
-	 */
-	partially_grouped_rel = fetch_upper_rel(root,
-											UPPERREL_PARTIAL_GROUP_AGG,
-											grouped_rel->relids);
+	 * aggregating the rows from the input relation.  The relation may already
+	 * exist due to eager aggregation, in which case we don't need to create
+	 * it.
+	 */
+	if (partially_grouped_rel == NULL)
+		partially_grouped_rel = fetch_upper_rel(root,
+												UPPERREL_PARTIAL_GROUP_AGG,
+												grouped_rel->relids);
 	partially_grouped_rel->consider_parallel =
 		grouped_rel->consider_parallel;
 	partially_grouped_rel->reloptkind = grouped_rel->reloptkind;
diff --git a/src/backend/optimizer/util/appendinfo.c b/src/backend/optimizer/util/appendinfo.c
index 51fdeace7d..7016473047 100644
--- a/src/backend/optimizer/util/appendinfo.c
+++ b/src/backend/optimizer/util/appendinfo.c
@@ -495,6 +495,70 @@ adjust_appendrel_attrs_mutator(Node *node,
 		return (Node *) newinfo;
 	}
 
+	/*
+	 * We have to process RelAggInfo nodes specially.
+	 */
+	if (IsA(node, RelAggInfo))
+	{
+		RelAggInfo *oldinfo = (RelAggInfo *) node;
+		RelAggInfo *newinfo = makeNode(RelAggInfo);
+
+		/* Copy all flat-copiable fields */
+		memcpy(newinfo, oldinfo, sizeof(RelAggInfo));
+
+		newinfo->relids = adjust_child_relids(oldinfo->relids,
+											  context->nappinfos,
+											  context->appinfos);
+
+		newinfo->target = (PathTarget *)
+			adjust_appendrel_attrs_mutator((Node *) oldinfo->target,
+										   context);
+
+		newinfo->agg_input = (PathTarget *)
+			adjust_appendrel_attrs_mutator((Node *) oldinfo->agg_input,
+										   context);
+
+		newinfo->group_clauses = (List *)
+			adjust_appendrel_attrs_mutator((Node *) oldinfo->group_clauses,
+										   context);
+
+		newinfo->group_exprs = (List *)
+			adjust_appendrel_attrs_mutator((Node *) oldinfo->group_exprs,
+										   context);
+
+		newinfo->agg_exprs = (List *)
+			adjust_appendrel_attrs_mutator((Node *) oldinfo->agg_exprs,
+										   context);
+
+		return (Node *) newinfo;
+	}
+
+	/*
+	 * We have to process PathTarget nodes specially.
+	 */
+	if (IsA(node, PathTarget))
+	{
+		PathTarget *oldtarget = (PathTarget *) node;
+		PathTarget *newtarget = makeNode(PathTarget);
+
+		/* Copy all flat-copiable fields */
+		memcpy(newtarget, oldtarget, sizeof(PathTarget));
+
+		if (oldtarget->sortgrouprefs)
+		{
+			Size	nbytes = list_length(oldtarget->exprs) * sizeof(Index);
+
+			newtarget->exprs = (List *)
+				adjust_appendrel_attrs_mutator((Node *) oldtarget->exprs,
+											   context);
+
+			newtarget->sortgrouprefs = (Index *) palloc(nbytes);
+			memcpy(newtarget->sortgrouprefs, oldtarget->sortgrouprefs, nbytes);
+		}
+
+		return (Node *) newtarget;
+	}
+
 	/*
 	 * NOTE: we do not need to recurse into sublinks, because they should
 	 * already have been converted to subplans before we see them.
-- 
2.31.0

