Improve performance of ORDER BY / DISTINCT aggregates

ORDER BY / DISTINCT aggreagtes have, since implemented in Postgres, been
executed by always performing a sort in nodeAgg.c to sort the tuples in
the current group into the correct order before calling the transition
function on the sorted tuples.  This was not great as often there might be
an index that could have provided pre-sorted input and allowed the
transition functions to be called as the rows come in, rather than having
to store them in a tuplestore in order to sort them once all the tuples
for the group have arrived.

Here we change the planner so it requests a path with a sort order which
supports the most amount of ORDER BY / DISTINCT aggregate functions and
add new code to the executor to allow it to support the processing of
ORDER BY / DISTINCT aggregates where the tuples are already sorted in the
correct order.

Since there can be many ORDER BY / DISTINCT aggregates in any given query
level, it's very possible that we can't find an order that suits all of
these aggregates.  The sort order that the planner chooses is simply the
one that suits the most aggregate functions.  We take the most strictly
sorted variation of each order and see how many aggregate functions can
use that, then we try again with the order of the remaining aggregates to
see if another order would suit more aggregate functions.  For example:

SELECT agg(a ORDER BY a),agg2(a ORDER BY a,b) ...

would request the sort order to be {a, b} because {a} is a subset of the
sort order of {a,b}, but;

SELECT agg(a ORDER BY a),agg2(a ORDER BY c) ...

would just pick a plan ordered by {a} (we give precedence to aggregates
which are earlier in the targetlist).

SELECT agg(a ORDER BY a),agg2(a ORDER BY b),agg3(a ORDER BY b) ...

would choose to order by {b} since two aggregates suit that vs just one
that requires input ordered by {a}.

Author: David Rowley
Reviewed-by: Ronan Dunklau, James Coleman, Ranier Vilela, Richard Guo, Tom Lane
Discussion: https://postgr.es/m/CAApHDvpHzfo92%3DR4W0%2BxVua3BUYCKMckWAmo-2t_KiXN-wYH%3Dw%40mail.gmail.com

src/backend/optimizer/path/pathkeys.c

@@ -103,6 +103,28 @@ make_canonical_pathkey(PlannerInfo *root,
 	return pk;
 }
 
+/*
+ * append_pathkeys
+ *		Append all non-redundant PathKeys in 'source' onto 'target' and
+ *		returns the updated 'target' list.
+ */
+List *
+append_pathkeys(List *target, List *source)
+{
+	ListCell   *lc;
+
+	Assert(target != NIL);
+
+	foreach(lc, source)
+	{
+		PathKey    *pk = lfirst_node(PathKey, lc);
+
+		if (!pathkey_is_redundant(pk, target))
+			target = lappend(target, pk);
+	}
+	return target;
+}
+
 /*
  * pathkey_is_redundant
  *	   Is a pathkey redundant with one already in the given list?
@@ -746,7 +768,8 @@ get_cheapest_group_keys_order(PlannerInfo *root, double nrows,
 List *
 get_useful_group_keys_orderings(PlannerInfo *root, double nrows,
 								List *path_pathkeys,
-								List *group_pathkeys, List *group_clauses)
+								List *group_pathkeys, List *group_clauses,
+								List *aggregate_pathkeys)
 {
 	Query	   *parse = root->parse;
 	List	   *infos = NIL;
@@ -758,7 +781,10 @@ get_useful_group_keys_orderings(PlannerInfo *root, double nrows,
 
 	/* always return at least the original pathkeys/clauses */
 	info = makeNode(PathKeyInfo);
-	info->pathkeys = pathkeys;
+	if (aggregate_pathkeys != NIL)
+		info->pathkeys = list_concat_copy(pathkeys, aggregate_pathkeys);
+	else
+		info->pathkeys = pathkeys;
 	info->clauses = clauses;
 
 	infos = lappend(infos, info);
@@ -783,7 +809,10 @@ get_useful_group_keys_orderings(PlannerInfo *root, double nrows,
 									  n_preordered))
 	{
 		info = makeNode(PathKeyInfo);
-		info->pathkeys = pathkeys;
+		if (aggregate_pathkeys != NIL)
+			info->pathkeys = list_concat_copy(pathkeys, aggregate_pathkeys);
+		else
+			info->pathkeys = pathkeys;
 		info->clauses = clauses;
 
 		infos = lappend(infos, info);
@@ -822,7 +851,10 @@ get_useful_group_keys_orderings(PlannerInfo *root, double nrows,
 		 * still want to keep the keys reordered to path_pathkeys.
 		 */
 		info = makeNode(PathKeyInfo);
-		info->pathkeys = pathkeys;
+		if (aggregate_pathkeys != NIL)
+			info->pathkeys = list_concat_copy(pathkeys, aggregate_pathkeys);
+		else
+			info->pathkeys = pathkeys;
 		info->clauses = clauses;
 
 		infos = lappend(infos, info);
@@ -850,7 +882,10 @@ get_useful_group_keys_orderings(PlannerInfo *root, double nrows,
 
 		/* keep the group keys reordered to match ordering of input path */
 		info = makeNode(PathKeyInfo);
-		info->pathkeys = pathkeys;
+		if (aggregate_pathkeys != NIL)
+			info->pathkeys = list_concat_copy(pathkeys, aggregate_pathkeys);
+		else
+			info->pathkeys = pathkeys;
 		info->clauses = clauses;
 
 		infos = lappend(infos, info);

src/backend/optimizer/plan/planagg.c

@@ -245,7 +245,7 @@ can_minmax_aggs(PlannerInfo *root, List **context)
 	foreach(lc, root->agginfos)
 	{
 		AggInfo    *agginfo = lfirst_node(AggInfo, lc);
-		Aggref	   *aggref = agginfo->representative_aggref;
+		Aggref	   *aggref = linitial_node(Aggref, agginfo->aggrefs);
 		Oid			aggsortop;
 		TargetEntry *curTarget;
 		MinMaxAggInfo *mminfo;

src/backend/optimizer/plan/planner.c

@@ -24,6 +24,7 @@
 #include "access/sysattr.h"
 #include "access/table.h"
 #include "access/xact.h"
+#include "catalog/pg_aggregate.h"
 #include "catalog/pg_constraint.h"
 #include "catalog/pg_inherits.h"
 #include "catalog/pg_proc.h"
@@ -3125,6 +3126,217 @@ reorder_grouping_sets(List *groupingsets, List *sortclause)
 	return result;
 }
 
+/*
+ * make_pathkeys_for_groupagg
+ *		Determine the pathkeys for the GROUP BY clause and/or any ordered
+ *		aggregates.  We expect at least one of these here.
+ *
+ * Building the pathkeys for the GROUP BY is simple.  Most of the complexity
+ * involved here comes from calculating the best pathkeys for ordered
+ * aggregates.  We define "best" as the pathkeys that suit the most number of
+ * aggregate functions.  We find these by looking at the first ORDER BY /
+ * DISTINCT aggregate and take the pathkeys for that before searching for
+ * other aggregates that require the same or a more strict variation of the
+ * same pathkeys.  We then repeat that process for any remaining aggregates
+ * with different pathkeys and if we find another set of pathkeys that suits a
+ * larger number of aggregates then we return those pathkeys instead.
+ *
+ * *number_groupby_pathkeys gets set to the number of elements in the returned
+ * list that belong to the GROUP BY clause.  Any elements above this number
+ * must belong to ORDER BY / DISTINCT aggregates.
+ *
+ * When the best pathkeys are found we also mark each Aggref that can use
+ * those pathkeys as aggpresorted = true.
+ */
+static List *
+make_pathkeys_for_groupagg(PlannerInfo *root, List *groupClause, List *tlist,
+						   int *number_groupby_pathkeys)
+{
+	List	   *grouppathkeys = NIL;
+	List	   *bestpathkeys;
+	Bitmapset  *bestaggs;
+	Bitmapset  *unprocessed_aggs;
+	ListCell   *lc;
+	int			i;
+
+	Assert(groupClause != NIL || root->numOrderedAggs > 0);
+
+	if (groupClause != NIL)
+	{
+		/* no pathkeys possible if there's an unsortable GROUP BY */
+		if (!grouping_is_sortable(groupClause))
+		{
+			*number_groupby_pathkeys = 0;
+			return NIL;
+		}
+
+		grouppathkeys = make_pathkeys_for_sortclauses(root, groupClause,
+													  tlist);
+		*number_groupby_pathkeys = list_length(grouppathkeys);
+	}
+	else
+		*number_groupby_pathkeys = 0;
+
+	/*
+	 * We can't add pathkeys for ordered aggregates if there are any grouping
+	 * sets.  All handling specific to ordered aggregates must be done by the
+	 * executor in that case.
+	 */
+	if (root->numOrderedAggs == 0 || root->parse->groupingSets != NIL)
+		return grouppathkeys;
+
+	/*
+	 * Make a first pass over all AggInfos to collect a Bitmapset containing
+	 * the indexes of all AggInfos to be processed below.
+	 */
+	unprocessed_aggs = NULL;
+	foreach(lc, root->agginfos)
+	{
+		AggInfo    *agginfo = lfirst_node(AggInfo, lc);
+		Aggref	   *aggref = linitial_node(Aggref, agginfo->aggrefs);
+
+		if (AGGKIND_IS_ORDERED_SET(aggref->aggkind))
+			continue;
+
+		/* only add aggregates with a DISTINCT or ORDER BY */
+		if (aggref->aggdistinct != NIL || aggref->aggorder != NIL)
+			unprocessed_aggs = bms_add_member(unprocessed_aggs,
+											  foreach_current_index(lc));
+	}
+
+	/*
+	 * Now process all the unprocessed_aggs to find the best set of pathkeys
+	 * for the given set of aggregates.
+	 *
+	 * On the first outer loop here 'bestaggs' will be empty.   We'll populate
+	 * this during the first loop using the pathkeys for the very first
+	 * AggInfo then taking any stronger pathkeys from any other AggInfos with
+	 * a more strict set of compatible pathkeys.  Once the outer loop is
+	 * complete, we mark off all the aggregates with compatible pathkeys then
+	 * remove those from the unprocessed_aggs and repeat the process to try to
+	 * find another set of pathkeys that are suitable for a larger number of
+	 * aggregates.  The outer loop will stop when there are not enough
+	 * unprocessed aggregates for it to be possible to find a set of pathkeys
+	 * to suit a larger number of aggregates.
+	 */
+	bestpathkeys = NIL;
+	bestaggs = NULL;
+	while (bms_num_members(unprocessed_aggs) > bms_num_members(bestaggs))
+	{
+		Bitmapset  *aggindexes = NULL;
+		List	   *currpathkeys = NIL;
+
+		i = -1;
+		while ((i = bms_next_member(unprocessed_aggs, i)) >= 0)
+		{
+			AggInfo    *agginfo = list_nth_node(AggInfo, root->agginfos, i);
+			Aggref	   *aggref = linitial_node(Aggref, agginfo->aggrefs);
+			List	   *sortlist;
+
+			if (aggref->aggdistinct != NIL)
+				sortlist = aggref->aggdistinct;
+			else
+				sortlist = aggref->aggorder;
+
+			/*
+			 * When not set yet, take the pathkeys from the first unprocessed
+			 * aggregate.
+			 */
+			if (currpathkeys == NIL)
+			{
+				currpathkeys = make_pathkeys_for_sortclauses(root, sortlist,
+															 aggref->args);
+
+				/* include the GROUP BY pathkeys, if they exist */
+				if (grouppathkeys != NIL)
+					currpathkeys = append_pathkeys(list_copy(grouppathkeys),
+												   currpathkeys);
+
+				/* record that we found pathkeys for this aggregate */
+				aggindexes = bms_add_member(aggindexes, i);
+			}
+			else
+			{
+				List	   *pathkeys;
+
+				/* now look for a stronger set of matching pathkeys */
+				pathkeys = make_pathkeys_for_sortclauses(root, sortlist,
+														 aggref->args);
+
+				/* include the GROUP BY pathkeys, if they exist */
+				if (grouppathkeys != NIL)
+					pathkeys = append_pathkeys(list_copy(grouppathkeys),
+											   pathkeys);
+
+				/* are 'pathkeys' compatible or better than 'currpathkeys'? */
+				switch (compare_pathkeys(currpathkeys, pathkeys))
+				{
+					case PATHKEYS_BETTER2:
+						/* 'pathkeys' are stronger, use these ones instead */
+						currpathkeys = pathkeys;
+						/* FALLTHROUGH */
+
+					case PATHKEYS_BETTER1:
+						/* 'pathkeys' are less strict */
+						/* FALLTHROUGH */
+
+					case PATHKEYS_EQUAL:
+						/* mark this aggregate as covered by 'currpathkeys' */
+						aggindexes = bms_add_member(aggindexes, i);
+						break;
+
+					case PATHKEYS_DIFFERENT:
+						break;
+				}
+			}
+		}
+
+		/* remove the aggregates that we've just processed */
+		unprocessed_aggs = bms_del_members(unprocessed_aggs, aggindexes);
+
+		/*
+		 * If this pass included more aggregates than the previous best then
+		 * use these ones as the best set.
+		 */
+		if (bms_num_members(aggindexes) > bms_num_members(bestaggs))
+		{
+			bestaggs = aggindexes;
+			bestpathkeys = currpathkeys;
+		}
+	}
+
+	/*
+	 * Now that we've found the best set of aggregates we can set the
+	 * presorted flag to indicate to the executor that it needn't bother
+	 * performing a sort for these Aggrefs.  We're able to do this now as
+	 * there's no chance of a Hash Aggregate plan as create_grouping_paths
+	 * will not mark the GROUP BY as GROUPING_CAN_USE_HASH due to the presence
+	 * of ordered aggregates.
+	 */
+	i = -1;
+	while ((i = bms_next_member(bestaggs, i)) >= 0)
+	{
+		AggInfo    *agginfo = list_nth_node(AggInfo, root->agginfos, i);
+
+		foreach(lc, agginfo->aggrefs)
+		{
+			Aggref	   *aggref = lfirst_node(Aggref, lc);
+
+			aggref->aggpresorted = true;
+		}
+	}
+
+	/*
+	 * bestpathkeys includes the GROUP BY pathkeys, so if we found any ordered
+	 * aggregates, then return bestpathkeys, otherwise return the
+	 * grouppathkeys.
+	 */
+	if (bestpathkeys != NIL)
+		return bestpathkeys;
+
+	return grouppathkeys;
+}
+
 /*
  * Compute query_pathkeys and other pathkeys during plan generation
  */
@@ -3137,18 +3349,19 @@ standard_qp_callback(PlannerInfo *root, void *extra)
 	List	   *activeWindows = qp_extra->activeWindows;
 
 	/*
-	 * Calculate pathkeys that represent grouping/ordering requirements.  The
-	 * sortClause is certainly sort-able, but GROUP BY and DISTINCT might not
-	 * be, in which case we just leave their pathkeys empty.
+	 * Calculate pathkeys that represent grouping/ordering and/or ordered
+	 * aggregate requirements.
 	 */
-	if (qp_extra->groupClause &&
-		grouping_is_sortable(qp_extra->groupClause))
-		root->group_pathkeys =
-			make_pathkeys_for_sortclauses(root,
-										  qp_extra->groupClause,
-										  tlist);
+	if (qp_extra->groupClause != NIL || root->numOrderedAggs > 0)
+		root->group_pathkeys = make_pathkeys_for_groupagg(root,
+														  qp_extra->groupClause,
+														  tlist,
+														  &root->num_groupby_pathkeys);
 	else
+	{
 		root->group_pathkeys = NIL;
+		root->num_groupby_pathkeys = 0;
+	}
 
 	/* We consider only the first (bottom) window in pathkeys logic */
 	if (activeWindows != NIL)
@@ -6222,6 +6435,26 @@ add_paths_to_grouping_rel(PlannerInfo *root, RelOptInfo *input_rel,
 
 	if (can_sort)
 	{
+		List	   *group_pathkeys;
+		List	   *orderAggPathkeys;
+		int			numAggPathkeys;
+
+		numAggPathkeys = list_length(root->group_pathkeys) -
+			root->num_groupby_pathkeys;
+
+		if (numAggPathkeys > 0)
+		{
+			group_pathkeys = list_copy_head(root->group_pathkeys,
+											root->num_groupby_pathkeys);
+			orderAggPathkeys = list_copy_tail(root->group_pathkeys,
+											  root->num_groupby_pathkeys);
+		}
+		else
+		{
+			group_pathkeys = root->group_pathkeys;
+			orderAggPathkeys = NIL;
+		}
+
 		/*
 		 * Use any available suitably-sorted path as input, and also consider
 		 * sorting the cheapest-total path.
@@ -6234,7 +6467,6 @@ add_paths_to_grouping_rel(PlannerInfo *root, RelOptInfo *input_rel,
 
 			List	   *pathkey_orderings = NIL;
 
-			List	   *group_pathkeys = root->group_pathkeys;
 			List	   *group_clauses = parse->groupClause;
 
 			/* generate alternative group orderings that might be useful */
@@ -6242,7 +6474,8 @@ add_paths_to_grouping_rel(PlannerInfo *root, RelOptInfo *input_rel,
 																path->rows,
 																path->pathkeys,
 																group_pathkeys,
-																group_clauses);
+																group_clauses,
+																orderAggPathkeys);
 
 			Assert(list_length(pathkey_orderings) > 0);
 
@@ -6414,7 +6647,8 @@ add_paths_to_grouping_rel(PlannerInfo *root, RelOptInfo *input_rel,
 																	path->rows,
 																	path->pathkeys,
 																	group_pathkeys,
-																	group_clauses);
+																	group_clauses,
+																	orderAggPathkeys);
 
 				Assert(list_length(pathkey_orderings) > 0);
 
@@ -6717,6 +6951,26 @@ create_partial_grouping_paths(PlannerInfo *root,
 
 	if (can_sort && cheapest_total_path != NULL)
 	{
+		List	   *group_pathkeys;
+		List	   *orderAggPathkeys;
+		int			numAggPathkeys;
+
+		numAggPathkeys = list_length(root->group_pathkeys) -
+			root->num_groupby_pathkeys;
+
+		if (numAggPathkeys > 0)
+		{
+			group_pathkeys = list_copy_head(root->group_pathkeys,
+											root->num_groupby_pathkeys);
+			orderAggPathkeys = list_copy_tail(root->group_pathkeys,
+											  root->num_groupby_pathkeys);
+		}
+		else
+		{
+			group_pathkeys = root->group_pathkeys;
+			orderAggPathkeys = NIL;
+		}
+
 		/* This should have been checked previously */
 		Assert(parse->hasAggs || parse->groupClause);
 
@@ -6729,10 +6983,7 @@ create_partial_grouping_paths(PlannerInfo *root,
 			ListCell   *lc2;
 			Path	   *path = (Path *) lfirst(lc);
 			Path	   *path_save = path;
-
 			List	   *pathkey_orderings = NIL;
-
-			List	   *group_pathkeys = root->group_pathkeys;
 			List	   *group_clauses = parse->groupClause;
 
 			/* generate alternative group orderings that might be useful */
@@ -6740,7 +6991,8 @@ create_partial_grouping_paths(PlannerInfo *root,
 																path->rows,
 																path->pathkeys,
 																group_pathkeys,
-																group_clauses);
+																group_clauses,
+																orderAggPathkeys);
 
 			Assert(list_length(pathkey_orderings) > 0);
 
@@ -6856,16 +7108,33 @@ create_partial_grouping_paths(PlannerInfo *root,
 
 	if (can_sort && cheapest_partial_path != NULL)
 	{
+		List	   *group_pathkeys;
+		List	   *orderAggPathkeys;
+		int			numAggPathkeys;
+
+		numAggPathkeys = list_length(root->group_pathkeys) -
+			root->num_groupby_pathkeys;
+
+		if (numAggPathkeys > 0)
+		{
+			group_pathkeys = list_copy_head(root->group_pathkeys,
+											root->num_groupby_pathkeys);
+			orderAggPathkeys = list_copy_tail(root->group_pathkeys,
+											  root->num_groupby_pathkeys);
+		}
+		else
+		{
+			group_pathkeys = root->group_pathkeys;
+			orderAggPathkeys = NIL;
+		}
+
 		/* Similar to above logic, but for partial paths. */
 		foreach(lc, input_rel->partial_pathlist)
 		{
 			ListCell   *lc2;
 			Path	   *path = (Path *) lfirst(lc);
 			Path	   *path_original = path;
-
 			List	   *pathkey_orderings = NIL;
-
-			List	   *group_pathkeys = root->group_pathkeys;
 			List	   *group_clauses = parse->groupClause;
 
 			/* generate alternative group orderings that might be useful */
@@ -6873,7 +7142,8 @@ create_partial_grouping_paths(PlannerInfo *root,
 																path->rows,
 																path->pathkeys,
 																group_pathkeys,
-																group_clauses);
+																group_clauses,
+																orderAggPathkeys);
 
 			Assert(list_length(pathkey_orderings) > 0);
 

src/backend/optimizer/prep/prepagg.c

@@ -225,6 +225,7 @@ preprocess_aggref(Aggref *aggref, PlannerInfo *root)
 	{
 		AggInfo    *agginfo = list_nth_node(AggInfo, root->agginfos, aggno);
 
+		agginfo->aggrefs = lappend(agginfo->aggrefs, aggref);
 		transno = agginfo->transno;
 	}
 	else
@@ -232,7 +233,7 @@ preprocess_aggref(Aggref *aggref, PlannerInfo *root)
 		AggInfo    *agginfo = makeNode(AggInfo);
 
 		agginfo->finalfn_oid = aggfinalfn;
-		agginfo->representative_aggref = aggref;
+		agginfo->aggrefs = list_make1(aggref);
 		agginfo->shareable = shareable;
 
 		aggno = list_length(root->agginfos);
@@ -386,7 +387,7 @@ find_compatible_agg(PlannerInfo *root, Aggref *newagg,
 
 		aggno++;
 
-		existingRef = agginfo->representative_aggref;
+		existingRef = linitial_node(Aggref, agginfo->aggrefs);
 
 		/* all of the following must be the same or it's no match */
 		if (newagg->inputcollid != existingRef->inputcollid ||
@@ -648,7 +649,7 @@ get_agg_clause_costs(PlannerInfo *root, AggSplit aggsplit, AggClauseCosts *costs
 	foreach(lc, root->agginfos)
 	{
 		AggInfo    *agginfo = lfirst_node(AggInfo, lc);
-		Aggref	   *aggref = agginfo->representative_aggref;
+		Aggref	   *aggref = linitial_node(Aggref, agginfo->aggrefs);
 
 		/*
 		 * Add the appropriate component function execution costs to