v2-0004-Implement-functions-that-create-RelAggInfos-if-applicable.patch

From 6c2afe0414f7c8a8b3df18d241f2912965ffda1c Mon Sep 17 00:00:00 2001
From: Richard Guo <guofenglinux@gmail.com>
Date: Fri, 23 Feb 2024 11:27:49 +0800
Subject: [PATCH v2 4/9] Implement functions that create RelAggInfos if
 applicable

This commit implements the functions that check if eager aggregation is
applicable for a given relation, and if so, create RelAggInfo structure
for the relation, using the infos about aggregate expressions and
grouping expressions we collected earlier.
---
 src/backend/optimizer/path/equivclass.c |  26 +-
 src/backend/optimizer/plan/planmain.c   |   3 +
 src/backend/optimizer/util/relnode.c    | 624 ++++++++++++++++++++++++
 src/backend/utils/adt/selfuncs.c        |   5 +-
 src/include/nodes/pathnodes.h           |   6 +
 src/include/optimizer/pathnode.h        |   5 +
 src/include/optimizer/paths.h           |   3 +-
 7 files changed, 662 insertions(+), 10 deletions(-)

diff --git a/src/backend/optimizer/path/equivclass.c b/src/backend/optimizer/path/equivclass.c
index 4bd60a09c6..1890dbb852 100644
--- a/src/backend/optimizer/path/equivclass.c
+++ b/src/backend/optimizer/path/equivclass.c
@@ -2439,15 +2439,17 @@ find_join_domain(PlannerInfo *root, Relids relids)
  *	  Detect whether two expressions are known equal due to equivalence
  *	  relationships.
  *
- * Actually, this only shows that the expressions are equal according
- * to some opfamily's notion of equality --- but we only use it for
- * selectivity estimation, so a fuzzy idea of equality is OK.
+ * If opfamily is given, the expressions must be known equal per the semantics
+ * of that opfamily (note it has to be a btree opfamily, since those are the
+ * only opfamilies equivclass.c deals with).  If opfamily is InvalidOid, we'll
+ * return true if they're equal according to any opfamily, which is fuzzy but
+ * OK for estimation purposes.
  *
  * Note: does not bother to check for "equal(item1, item2)"; caller must
  * check that case if it's possible to pass identical items.
  */
 bool
-exprs_known_equal(PlannerInfo *root, Node *item1, Node *item2)
+exprs_known_equal(PlannerInfo *root, Node *item1, Node *item2, Oid opfamily)
 {
 	ListCell   *lc1;
 
@@ -2462,6 +2464,17 @@ exprs_known_equal(PlannerInfo *root, Node *item1, Node *item2)
 		if (ec->ec_has_volatile)
 			continue;
 
+		/*
+		 * It's okay to consider ec_broken ECs here.  Brokenness just means we
+		 * couldn't derive all the implied clauses we'd have liked to; it does
+		 * not invalidate our knowledge that the members are equal.
+		 */
+
+		/* Ignore if this EC doesn't use specified opfamily */
+		if (OidIsValid(opfamily) &&
+			!list_member_oid(ec->ec_opfamilies, opfamily))
+			continue;
+
 		foreach(lc2, ec->ec_members)
 		{
 			EquivalenceMember *em = (EquivalenceMember *) lfirst(lc2);
@@ -2490,8 +2503,7 @@ exprs_known_equal(PlannerInfo *root, Node *item1, Node *item2)
  * (In principle there might be more than one matching eclass if multiple
  * collations are involved, but since collation doesn't matter for equality,
  * we ignore that fine point here.)  This is much like exprs_known_equal,
- * except that we insist on the comparison operator matching the eclass, so
- * that the result is definite not approximate.
+ * except for the format of the input.
  *
  * On success, we also set fkinfo->eclass[colno] to the matching eclass,
  * and set fkinfo->fk_eclass_member[colno] to the eclass member for the
@@ -2532,7 +2544,7 @@ match_eclasses_to_foreign_key_col(PlannerInfo *root,
 		/* Never match to a volatile EC */
 		if (ec->ec_has_volatile)
 			continue;
-		/* Note: it seems okay to match to "broken" eclasses here */
+		/* It's okay to consider "broken" ECs here, see exprs_known_equal */
 
 		foreach(lc2, ec->ec_members)
 		{
diff --git a/src/backend/optimizer/plan/planmain.c b/src/backend/optimizer/plan/planmain.c
index 197a3f905e..0ff0ca99cb 100644
--- a/src/backend/optimizer/plan/planmain.c
+++ b/src/backend/optimizer/plan/planmain.c
@@ -67,6 +67,9 @@ query_planner(PlannerInfo *root,
 	root->join_rel_list = makeNode(RelInfoList);
 	root->join_rel_list->items = NIL;
 	root->join_rel_list->hash = NULL;
+	root->agg_info_list = makeNode(RelInfoList);
+	root->agg_info_list->items = NIL;
+	root->agg_info_list->hash = NULL;
 	root->join_rel_level = NULL;
 	root->join_cur_level = 0;
 	root->canon_pathkeys = NIL;
diff --git a/src/backend/optimizer/util/relnode.c b/src/backend/optimizer/util/relnode.c
index c88da963db..e7f465ef7b 100644
--- a/src/backend/optimizer/util/relnode.c
+++ b/src/backend/optimizer/util/relnode.c
@@ -87,6 +87,14 @@ static void build_child_join_reltarget(PlannerInfo *root,
 									   RelOptInfo *childrel,
 									   int nappinfos,
 									   AppendRelInfo **appinfos);
+static bool eager_aggregation_possible_for_relation(PlannerInfo *root,
+													RelOptInfo *rel);
+static bool init_grouping_targets(PlannerInfo *root, RelOptInfo *rel,
+								  PathTarget *target, PathTarget *agg_input,
+								  List **group_exprs_extra_p);
+static bool is_var_in_aggref_only(PlannerInfo *root, Var *var);
+static bool is_var_needed_by_join(PlannerInfo *root, Var *var, RelOptInfo *rel);
+static Index get_expression_sortgroupref(PlannerInfo *root, Expr *expr);
 
 
 /*
@@ -640,6 +648,58 @@ add_join_rel(PlannerInfo *root, RelOptInfo *joinrel)
 	add_rel_info(root->join_rel_list, joinrel);
 }
 
+/*
+ * add_grouped_rel
+ *		Add grouped base or join relation to the list of grouped relations in
+ *		the given PlannerInfo.  Also add the corresponding RelAggInfo to
+ *		root->agg_info_list.
+ */
+void
+add_grouped_rel(PlannerInfo *root, RelOptInfo *rel, RelAggInfo *agg_info)
+{
+	add_rel_info(&root->upper_rels[UPPERREL_PARTIAL_GROUP_AGG], rel);
+	add_rel_info(root->agg_info_list, agg_info);
+}
+
+/*
+ * find_grouped_rel
+ *	  Returns grouped relation entry (base or join relation) corresponding to
+ *	  'relids' or NULL if none exists.
+ *
+ * If agg_info_p is not NULL, then also the corresponding RelAggInfo (if one
+ * exists) will be returned in *agg_info_p.
+ */
+RelOptInfo *
+find_grouped_rel(PlannerInfo *root, Relids relids, RelAggInfo **agg_info_p)
+{
+	RelOptInfo *rel;
+
+	rel = (RelOptInfo *) find_rel_info(&root->upper_rels[UPPERREL_PARTIAL_GROUP_AGG],
+									   relids);
+	if (rel == NULL)
+	{
+		if (agg_info_p)
+			*agg_info_p = NULL;
+
+		return NULL;
+	}
+
+	/* also return the corresponding RelAggInfo, if asked */
+	if (agg_info_p)
+	{
+		RelAggInfo *agg_info;
+
+		agg_info = (RelAggInfo *) find_rel_info(root->agg_info_list, relids);
+
+		/* The relation exists, so the agg_info should be there too. */
+		Assert(agg_info != NULL);
+
+		*agg_info_p = agg_info;
+	}
+
+	return rel;
+}
+
 /*
  * set_foreign_rel_properties
  *		Set up foreign-join fields if outer and inner relation are foreign
@@ -2464,3 +2524,567 @@ build_child_join_reltarget(PlannerInfo *root,
 	childrel->reltarget->cost.per_tuple = parentrel->reltarget->cost.per_tuple;
 	childrel->reltarget->width = parentrel->reltarget->width;
 }
+
+/*
+ * create_rel_agg_info
+ *	  Check if the given relation can produce grouped paths and return the
+ *	  information it'll need for it.  The given relation is the non-grouped one
+ *	  which has the reltarget already constructed.
+ */
+RelAggInfo *
+create_rel_agg_info(PlannerInfo *root, RelOptInfo *rel)
+{
+	ListCell   *lc;
+	RelAggInfo *result;
+	PathTarget *agg_input;
+	PathTarget *target;
+	List	   *grp_exprs_extra = NIL;
+	List	   *group_clauses_final;
+	int			i;
+
+	/*
+	 * The lists of aggregate expressions and grouping expressions should have
+	 * been constructed.
+	 */
+	Assert(root->agg_clause_list != NIL);
+	Assert(root->group_expr_list != NIL);
+
+	/*
+	 * If this is a child rel, the grouped rel for its parent rel must have
+	 * been created if it can.  So we can just use parent's RelAggInfo if there
+	 * is one, with appropriate variable substitutions.
+	 */
+	if (IS_OTHER_REL(rel))
+	{
+		RelOptInfo	   *rel_grouped;
+		RelAggInfo	   *agg_info;
+		Relids			parent_relids = NULL;
+		AppendRelInfo **appinfos;
+		int				nappinfos;
+		int				cnt;
+
+		appinfos = find_appinfos_by_relids(root, rel->relids, &nappinfos);
+
+		for (cnt = 0; cnt < nappinfos; cnt++)
+			parent_relids = bms_add_member(parent_relids,
+										   appinfos[cnt]->parent_relid);
+
+		Assert(!bms_is_empty(parent_relids));
+		rel_grouped = find_grouped_rel(root, parent_relids, &agg_info);
+
+		if (rel_grouped == NULL)
+			return NULL;
+
+		Assert(agg_info != NULL);
+
+		agg_info = (RelAggInfo *) adjust_appendrel_attrs(root,
+														 (Node *) agg_info,
+														 nappinfos,
+														 appinfos);
+
+		pfree(appinfos);
+
+		agg_info->input_rows = rel->rows;
+		agg_info->grouped_rows =
+			estimate_num_groups(root, agg_info->group_exprs,
+								agg_info->input_rows, NULL, NULL);
+
+		return agg_info;
+	}
+
+	/* Check if it's possible to produce grouped paths for this relation. */
+	if (!eager_aggregation_possible_for_relation(root, rel))
+		return NULL;
+
+	/*
+	 * Create targets for the grouped paths and for the input paths of the
+	 * grouped paths.
+	 */
+	target = create_empty_pathtarget();
+	agg_input = create_empty_pathtarget();
+
+	/* initialize 'target' and 'agg_input' */
+	if (!init_grouping_targets(root, rel, target, agg_input, &grp_exprs_extra))
+		return NULL;
+
+	/* Eager aggregation makes no sense w/o grouping expressions */
+	if ((list_length(target->exprs) + list_length(grp_exprs_extra)) == 0)
+		return NULL;
+
+	group_clauses_final = root->parse->groupClause;
+
+	/*
+	 * If the aggregation target should have extra grouping expressions (in
+	 * order to emit input vars for join conditions), add them now. This step
+	 * includes assignment of tleSortGroupRef's which we can generate now.
+	 */
+	if (list_length(grp_exprs_extra) > 0)
+	{
+		Index		sortgroupref;
+
+		/*
+		 * Make a copy of the group clauses as we'll need to add some more
+		 * clauses.
+		 */
+		group_clauses_final = list_copy(group_clauses_final);
+
+		/* find out the current max sortgroupref */
+		sortgroupref = 0;
+		foreach(lc, root->processed_tlist)
+		{
+			Index		ref = ((TargetEntry *) lfirst(lc))->ressortgroupref;
+
+			if (ref > sortgroupref)
+				sortgroupref = ref;
+		}
+
+		/*
+		 * Generate the SortGroupClause's and add the expressions to the
+		 * target.
+		 */
+		foreach(lc, grp_exprs_extra)
+		{
+			Var		   *var = lfirst_node(Var, lc);
+			SortGroupClause *cl = makeNode(SortGroupClause);
+
+			/*
+			 * Initialize the SortGroupClause.
+			 *
+			 * As the final aggregation will not use this grouping expression,
+			 * we don't care whether sortop is < or >. The value of nulls_first
+			 * should not matter for the same reason.
+			 */
+			cl->tleSortGroupRef = ++sortgroupref;
+			get_sort_group_operators(var->vartype,
+									 false, true, false,
+									 &cl->sortop, &cl->eqop, NULL,
+									 &cl->hashable);
+			group_clauses_final = lappend(group_clauses_final, cl);
+			add_column_to_pathtarget(target, (Expr *) var,
+									 cl->tleSortGroupRef);
+
+			/*
+			 * The aggregation input target must emit this var too.
+			 */
+			add_column_to_pathtarget(agg_input, (Expr *) var,
+									 cl->tleSortGroupRef);
+		}
+	}
+
+	/*
+	 * Build a list of grouping expressions and a list of the corresponding
+	 * SortGroupClauses.
+	 */
+	i = 0;
+	result = makeNode(RelAggInfo);
+	foreach(lc, target->exprs)
+	{
+		Index		sortgroupref = 0;
+		SortGroupClause *cl;
+		Expr	   *texpr;
+
+		texpr = (Expr *) lfirst(lc);
+
+		Assert(IsA(texpr, Var));
+
+		sortgroupref = target->sortgrouprefs[i++];
+		if (sortgroupref == 0)
+			continue;
+
+		/* find the SortGroupClause in group_clauses_final */
+		cl = get_sortgroupref_clause(sortgroupref, group_clauses_final);
+
+		/* do not add this SortGroupClause if it has already been added */
+		if (list_member(result->group_clauses, cl))
+			continue;
+
+		result->group_clauses = lappend(result->group_clauses, cl);
+		result->group_exprs = list_append_unique(result->group_exprs,
+												 texpr);
+	}
+
+	/*
+	 * Calculate pathkeys that represent this grouping requirements.
+	 */
+	result->group_pathkeys =
+		make_pathkeys_for_sortclauses(root, result->group_clauses,
+									  make_tlist_from_pathtarget(target));
+
+	/*
+	 * Add aggregates to the grouping target.
+	 */
+	foreach(lc, root->agg_clause_list)
+	{
+		AggClauseInfo *ac_info = lfirst_node(AggClauseInfo, lc);
+		Aggref   *aggref;
+
+		Assert(IsA(ac_info->aggref, Aggref));
+
+		aggref = (Aggref *) copyObject(ac_info->aggref);
+		mark_partial_aggref(aggref, AGGSPLIT_INITIAL_SERIAL);
+
+		add_column_to_pathtarget(target, (Expr *) aggref, 0);
+
+		result->agg_exprs = lappend(result->agg_exprs, aggref);
+	}
+
+	/*
+	 * Since neither target nor agg_input is supposed to be identical to the
+	 * source reltarget, compute the width and cost again.
+	 */
+	set_pathtarget_cost_width(root, target);
+	set_pathtarget_cost_width(root, agg_input);
+
+	result->relids = bms_copy(rel->relids);
+	result->target = target;
+	result->agg_input = agg_input;
+
+	/*
+	 * The number of aggregation input rows is simply the number of rows of the
+	 * non-grouped relation, which should have been estimated by now.
+	 */
+	result->input_rows = rel->rows;
+
+	/* Estimate the number of groups with equal grouped exprs. */
+	result->grouped_rows = estimate_num_groups(root, result->group_exprs,
+											   result->input_rows, NULL, NULL);
+
+	return result;
+}
+
+/*
+ * eager_aggregation_possible_for_relation
+ * 	  Check if it's possible to produce grouped paths for the given relation.
+ */
+static bool
+eager_aggregation_possible_for_relation(PlannerInfo *root, RelOptInfo *rel)
+{
+	ListCell   *lc;
+
+	/*
+	 * The current implementation of eager aggregation cannot handle
+	 * PlaceHolderVar (PHV).
+	 *
+	 * If we knew that the PHV should be evaluated in this target (and of
+	 * course, if its expression matched some Aggref argument), we'd just let
+	 * init_grouping_targets add that Aggref. On the other hand, if we knew
+	 * that the PHV is evaluated below the current rel, we could ignore it
+	 * because the referencing Aggref would take care of propagation of the
+	 * value to upper joins.
+	 *
+	 * The problem is that the same PHV can be evaluated in the target of the
+	 * current rel or in that of lower rel --- depending on the input paths.
+	 * For example, consider rel->relids = {A, B, C} and if ph_eval_at = {B,
+	 * C}. Path "A JOIN (B JOIN C)" implies that the PHV is evaluated by the
+	 * "(B JOIN C)", while path "(A JOIN B) JOIN C" evaluates the PHV itself.
+	 */
+	foreach(lc, rel->reltarget->exprs)
+	{
+		Expr	   *expr = lfirst(lc);
+
+		if (IsA(expr, PlaceHolderVar))
+			return false;
+	}
+
+	if (IS_SIMPLE_REL(rel))
+	{
+		RangeTblEntry *rte = root->simple_rte_array[rel->relid];;
+
+		/*
+		 * rtekind != RTE_RELATION case is not supported yet.
+		 */
+		if (rte->rtekind != RTE_RELATION)
+			return false;
+	}
+
+	/* Caller should only pass base relations or joins. */
+	Assert(rel->reloptkind == RELOPT_BASEREL ||
+		   rel->reloptkind == RELOPT_JOINREL);
+
+	/*
+	 * Check if all aggregate expressions can be evaluated on this relation
+	 * level.
+	 */
+	foreach(lc, root->agg_clause_list)
+	{
+		AggClauseInfo *ac_info = lfirst_node(AggClauseInfo, lc);
+
+		Assert(IsA(ac_info->aggref, Aggref));
+
+		/*
+		 * Give up if any aggregate needs relations other than the current one.
+		 *
+		 * If the aggregate needs the current rel plus anything else, then the
+		 * problem is that grouping of the current relation could make some
+		 * input variables unavailable for the "higher aggregate", and it'd
+		 * also decrease the number of input rows the "higher aggregate"
+		 * receives.
+		 *
+		 * If the aggregate does not even need the current rel, then the
+		 * current rel should be grouped because we do not support join of two
+		 * grouped relations.
+		 */
+		if (!bms_is_subset(ac_info->agg_eval_at, rel->relids))
+			return false;
+	}
+
+	return true;
+}
+
+/*
+ * init_grouping_targets
+ *	  Initialize target for grouped paths (target) as well as a target for
+ *	  paths that generate input for the grouped paths (agg_input).
+ *
+ * group_exprs_extra_p receives a list of Var nodes for which we need to
+ * construct SortGroupClause. Those vars will then be used as additional
+ * grouping expressions, for the sake of join clauses.
+ *
+ * Return true iff the targets could be initialized.
+ */
+static bool
+init_grouping_targets(PlannerInfo *root, RelOptInfo *rel,
+					  PathTarget *target, PathTarget *agg_input,
+					  List **group_exprs_extra_p)
+{
+	ListCell   *lc;
+	List	   *possibly_dependent = NIL;
+
+	foreach(lc, rel->reltarget->exprs)
+	{
+		Expr	   *expr = (Expr *) lfirst(lc);
+		Index		sortgroupref;
+
+		/*
+		 * Given that PlaceHolderVar currently prevents us from doing eager
+		 * aggregation, the source target cannot contain anything more complex
+		 * than a Var.
+		 */
+		Assert(IsA(expr, Var));
+
+		/* Get the sortgroupref if the expr can act as grouping expression. */
+		sortgroupref = get_expression_sortgroupref(root, expr);
+		if (sortgroupref > 0)
+		{
+			/*
+			 * If the target expression can be used as the grouping key, it
+			 * should be emitted by the grouped paths that have been pushed
+			 * down to this relation level.
+			 */
+			add_column_to_pathtarget(target, expr, sortgroupref);
+
+			/*
+			 * ... and it also should be emitted by the input paths
+			 */
+			add_column_to_pathtarget(agg_input, expr, sortgroupref);
+		}
+		else
+		{
+			if (is_var_needed_by_join(root, (Var *) expr, rel))
+			{
+				/*
+				 * The variable is needed for a join, however it's neither in
+				 * the GROUP BY clause nor can it be derived from it using EC.
+				 * (Otherwise it would have to be added to the targets above.)
+				 * We need to construct special SortGroupClause for this
+				 * variable.
+				 *
+				 * Note that its tleSortGroupRef needs to be unique within
+				 * agg_input, so we need to postpone creation of the
+				 * SortGroupClause's until we're done with the iteration of
+				 * rel->reltarget->exprs. Also it makes sense for the caller to
+				 * do some more check before it starts to create those
+				 * SortGroupClause's.
+				 */
+				*group_exprs_extra_p = lappend(*group_exprs_extra_p, expr);
+			}
+			else if (is_var_in_aggref_only(root, (Var *) expr))
+			{
+				/*
+				 * Another reason we might need this variable is that some
+				 * aggregate pushed down to this relation references it. In
+				 * such a case, add it to "agg_input", but not to "target".
+				 * However, if the aggregate is not the only reason for the var
+				 * to be in the target, some more checks need to be performed
+				 * below.
+				 */
+				add_new_column_to_pathtarget(agg_input, expr);
+			}
+			else
+			{
+				/*
+				 * The Var can be functionally dependent on another expression
+				 * of the target, but we cannot check that until we've built
+				 * all the expressions for the target.
+				 */
+				possibly_dependent = lappend(possibly_dependent, expr);
+			}
+		}
+	}
+
+	/*
+	 * Now we can check whether the expression is functionally dependent on
+	 * another one.
+	 */
+	foreach(lc, possibly_dependent)
+	{
+		Var		   *tvar;
+		List	   *deps = NIL;
+		RangeTblEntry *rte;
+
+		tvar = lfirst_node(Var, lc);
+		rte = root->simple_rte_array[tvar->varno];
+
+		/*
+		 * Check if the Var can be in the grouping key even though it's not
+		 * mentioned by the GROUP BY clause (and could not be derived using
+		 * ECs).
+		 */
+		if (check_functional_grouping(rte->relid, tvar->varno,
+									  tvar->varlevelsup,
+									  target->exprs, &deps))
+		{
+			/*
+			 * The var shouldn't be actually used for grouping key evaluation
+			 * (instead, the one this depends on will be), so sortgroupref
+			 * should not be important.
+			 */
+			add_new_column_to_pathtarget(target, (Expr *) tvar);
+			add_new_column_to_pathtarget(agg_input, (Expr *) tvar);
+		}
+		else
+		{
+			/*
+			 * As long as the query is semantically correct, arriving here
+			 * means that the var is referenced by a generic grouping
+			 * expression but not referenced by any join.
+			 *
+			 * If the eager aggregation will support generic grouping
+			 * expression in the future, create_rel_agg_info() will have to add
+			 * this variable to "agg_input" target and also add the whole
+			 * generic expression to "target".
+			 */
+			return false;
+		}
+	}
+
+	return true;
+}
+
+/*
+ * is_var_in_aggref_only
+ *	  Check whether given Var appears in Aggref(s) which we consider usable at
+ *	  relation / join level, and only in the Aggref(s).
+ */
+static bool
+is_var_in_aggref_only(PlannerInfo *root, Var *var)
+{
+	ListCell   *lc;
+
+	foreach(lc, root->agg_clause_list)
+	{
+		AggClauseInfo *ac_info = lfirst_node(AggClauseInfo, lc);
+
+		Assert(IsA(ac_info->aggref, Aggref));
+
+		if (bms_is_member(var->varno, ac_info->agg_eval_at))
+			break;;
+	}
+
+	/* No aggregate references the Var? */
+	if (lc == NULL)
+		return false;
+
+	/* Does the Var appear in the target outside aggregates? */
+	foreach(lc, root->processed_tlist)
+	{
+		TargetEntry *tle = lfirst_node(TargetEntry, lc);
+		List	   *vars;
+
+		if (IsA(tle->expr, Aggref))
+			continue;
+
+		vars = pull_var_clause((Node *) tle->expr,
+							   PVC_RECURSE_WINDOWFUNCS |
+							   PVC_RECURSE_PLACEHOLDERS);
+		if (list_member(vars, var))
+		{
+			list_free(vars);
+			return false;
+		}
+
+		list_free(vars);
+	}
+
+	/* The Var is in aggregate(s) and only there. */
+	return true;
+}
+
+/*
+ * is_var_needed_by_join
+ *	  Check if the given Var is needed by joins above the current rel.
+ *
+ * Consider pushing the aggregate avg(b.y) down to relation b for the following
+ * query:
+ *
+ *    SELECT a.i, avg(b.y)
+ *    FROM a JOIN b ON a.j = b.j
+ *    GROUP BY a.i;
+ *
+ * Column b.j needs to be used as the grouping key because otherwise it cannot
+ * find its way to the input of the join expression.
+ */
+static bool
+is_var_needed_by_join(PlannerInfo *root, Var *var, RelOptInfo *rel)
+{
+	Relids		relids;
+	int			attno;
+	RelOptInfo *baserel;
+
+	/*
+	 * Note that when we are checking if the Var is needed by joins above, we
+	 * want to exclude the situation where the Var is only needed in final
+	 * output.  So include "relation 0" here.
+	 */
+	relids = bms_copy(rel->relids);
+	relids = bms_add_member(relids, 0);
+
+	baserel = find_base_rel(root, var->varno);
+	attno = var->varattno - baserel->min_attr;
+
+
+	return bms_nonempty_difference(baserel->attr_needed[attno], relids);
+}
+
+/*
+ * get_expression_sortgroupref
+ *	  Return sortgroupref if the given 'expr' can be used as a grouping
+ *	  expression in grouped paths for base or join relations, or 0 otherwise.
+ *
+ * Note that we also need to check if the 'expr' is known equal to other exprs
+ * due to equivalence relationships that can act as grouping expressions.
+ */
+static Index
+get_expression_sortgroupref(PlannerInfo *root, Expr *expr)
+{
+	ListCell   *lc;
+
+	foreach(lc, root->group_expr_list)
+	{
+		GroupExprInfo *ge_info = lfirst_node(GroupExprInfo, lc);
+
+		Assert(IsA(ge_info->expr, Var));
+
+		if (equal(ge_info->expr, expr) ||
+			exprs_known_equal(root, (Node *) expr, (Node *) ge_info->expr,
+							  ge_info->btree_opfamily))
+		{
+			Assert(ge_info->sortgroupref > 0);
+
+			return ge_info->sortgroupref;
+		}
+	}
+
+	/* The expression cannot be used as grouping key. */
+	return 0;
+}
diff --git a/src/backend/utils/adt/selfuncs.c b/src/backend/utils/adt/selfuncs.c
index cea777e9d4..d1365229f7 100644
--- a/src/backend/utils/adt/selfuncs.c
+++ b/src/backend/utils/adt/selfuncs.c
@@ -3313,10 +3313,11 @@ add_unique_group_var(PlannerInfo *root, List *varinfos,
 
 		/*
 		 * Drop known-equal vars, but only if they belong to different
-		 * relations (see comments for estimate_num_groups)
+		 * relations (see comments for estimate_num_groups).  We aren't too
+		 * fussy about the semantics of "equal" here.
 		 */
 		if (vardata->rel != varinfo->rel &&
-			exprs_known_equal(root, var, varinfo->var))
+			exprs_known_equal(root, var, varinfo->var, InvalidOid))
 		{
 			if (varinfo->ndistinct <= ndistinct)
 			{
diff --git a/src/include/nodes/pathnodes.h b/src/include/nodes/pathnodes.h
index 69ed9eb1f6..3ef5195323 100644
--- a/src/include/nodes/pathnodes.h
+++ b/src/include/nodes/pathnodes.h
@@ -429,6 +429,12 @@ struct PlannerInfo
 	 */
 	RelInfoList		upper_rels[UPPERREL_FINAL + 1] pg_node_attr(read_write_ignore);;
 
+	/*
+	 * list of grouped relation RelAggInfos. One instance of RelAggInfo per
+	 * item of the upper_rels[UPPERREL_PARTIAL_GROUP_AGG] list.
+	 */
+	RelInfoList *agg_info_list;
+
 	/* Result tlists chosen by grouping_planner for upper-stage processing */
 	struct PathTarget *upper_targets[UPPERREL_FINAL + 1] pg_node_attr(read_write_ignore);
 
diff --git a/src/include/optimizer/pathnode.h b/src/include/optimizer/pathnode.h
index c43d97b48a..8d03ce2c57 100644
--- a/src/include/optimizer/pathnode.h
+++ b/src/include/optimizer/pathnode.h
@@ -310,6 +310,10 @@ extern RelOptInfo *find_base_rel(PlannerInfo *root, int relid);
 extern RelOptInfo *find_base_rel_noerr(PlannerInfo *root, int relid);
 extern RelOptInfo *find_base_rel_ignore_join(PlannerInfo *root, int relid);
 extern RelOptInfo *find_join_rel(PlannerInfo *root, Relids relids);
+extern void add_grouped_rel(PlannerInfo *root, RelOptInfo *rel,
+							RelAggInfo *agg_info);
+extern RelOptInfo *find_grouped_rel(PlannerInfo *root, Relids relids,
+									RelAggInfo **agg_info_p);
 extern RelOptInfo *build_join_rel(PlannerInfo *root,
 								  Relids joinrelids,
 								  RelOptInfo *outer_rel,
@@ -344,4 +348,5 @@ extern RelOptInfo *build_child_join_rel(PlannerInfo *root,
 										RelOptInfo *parent_joinrel, List *restrictlist,
 										SpecialJoinInfo *sjinfo);
 
+extern RelAggInfo *create_rel_agg_info(PlannerInfo *root, RelOptInfo *rel);
 #endif							/* PATHNODE_H */
diff --git a/src/include/optimizer/paths.h b/src/include/optimizer/paths.h
index 040a047b81..dcea10888b 100644
--- a/src/include/optimizer/paths.h
+++ b/src/include/optimizer/paths.h
@@ -160,7 +160,8 @@ extern List *generate_join_implied_equalities_for_ecs(PlannerInfo *root,
 													  Relids join_relids,
 													  Relids outer_relids,
 													  RelOptInfo *inner_rel);
-extern bool exprs_known_equal(PlannerInfo *root, Node *item1, Node *item2);
+extern bool exprs_known_equal(PlannerInfo *root, Node *item1, Node *item2,
+							  Oid opfamily);
 extern EquivalenceClass *match_eclasses_to_foreign_key_col(PlannerInfo *root,
 														   ForeignKeyOptInfo *fkinfo,
 														   int colno);
-- 
2.31.0