Cache catalog lookup data across groups in ordered-set aggregates.

The initial commit of ordered-set aggregates just did all the setup work
afresh each time the aggregate function is started up.  But in a GROUP BY
query, the catalog lookups need not be repeated for each group, since the
column datatypes and sort information won't change.  When there are many
small groups, this makes for a useful, though not huge, performance
improvement.  Per suggestion from Andrew Gierth.

Profiling of these cases suggests that it might be profitable to avoid
duplicate lookups within tuplesort startup as well; but changing the
tuplesort APIs would have much broader impact, so I left that for
another day.

src/backend/utils/adt/orderedsetaggs.c

@@ -32,16 +32,25 @@
 
 /*
  * Generic support for ordered-set aggregates
+ *
+ * The state for an ordered-set aggregate is divided into a per-group struct
+ * (which is the internal-type transition state datum returned to nodeAgg.c)
+ * and a per-query struct, which contains data and sub-objects that we can
+ * create just once per query because they will not change across groups.
+ * The per-query struct and subsidiary data live in the executor's per-query
+ * memory context, and go away implicitly at ExecutorEnd().
  */
 
-typedef struct OrderedSetAggState
+typedef struct OSAPerQueryState
 {
 	/* Aggref for this aggregate: */
 	Aggref	   *aggref;
-	/* Sort object we're accumulating data in: */
-	Tuplesortstate *sortstate;
-	/* Number of normal rows inserted into sortstate: */
-	int64		number_of_rows;
+	/* Memory context containing this struct and other per-query data: */
+	MemoryContext qcontext;
+	/* Memory context containing per-group data: */
+	MemoryContext gcontext;
+	/* Agg plan node's output econtext: */
+	ExprContext *peraggecontext;
 
 	/* These fields are used only when accumulating tuples: */
 
@@ -49,17 +58,41 @@ typedef struct OrderedSetAggState
 	TupleDesc	tupdesc;
 	/* Tuple slot we can use for inserting/extracting tuples: */
 	TupleTableSlot *tupslot;
+	/* Per-sort-column sorting information */
+	int			numSortCols;
+	AttrNumber *sortColIdx;
+	Oid		   *sortOperators;
+	Oid		   *eqOperators;
+	Oid		   *sortCollations;
+	bool	   *sortNullsFirsts;
+	/* Equality operator call info, created only if needed: */
+	FmgrInfo   *equalfns;
 
 	/* These fields are used only when accumulating datums: */
 
 	/* Info about datatype of datums being sorted: */
-	Oid			datumtype;
+	Oid			sortColType;
 	int16		typLen;
 	bool		typByVal;
 	char		typAlign;
-	/* Info about equality operator associated with sort operator: */
+	/* Info about sort ordering: */
+	Oid			sortOperator;
 	Oid			eqOperator;
-} OrderedSetAggState;
+	Oid			sortCollation;
+	bool		sortNullsFirst;
+	/* Equality operator call info, created only if needed: */
+	FmgrInfo	equalfn;
+} OSAPerQueryState;
+
+typedef struct OSAPerGroupState
+{
+	/* Link to the per-query state for this aggregate: */
+	OSAPerQueryState *qstate;
+	/* Sort object we're accumulating data in: */
+	Tuplesortstate *sortstate;
+	/* Number of normal rows inserted into sortstate: */
+	int64		number_of_rows;
+} OSAPerGroupState;
 
 static void ordered_set_shutdown(Datum arg);
 
@@ -67,175 +100,208 @@ static void ordered_set_shutdown(Datum arg);
 /*
  * Set up working state for an ordered-set aggregate
  */
-static OrderedSetAggState *
+static OSAPerGroupState *
 ordered_set_startup(FunctionCallInfo fcinfo, bool use_tuples)
 {
-	OrderedSetAggState *osastate;
-	Aggref	   *aggref;
-	ExprContext *peraggecontext;
-	MemoryContext aggcontext;
+	OSAPerGroupState *osastate;
+	OSAPerQueryState *qstate;
 	MemoryContext oldcontext;
-	List	   *sortlist;
-	int			numSortCols;
 
-	/* Must be called as aggregate; get the Agg node's query-lifespan context */
-	if (AggCheckCallContext(fcinfo, &aggcontext) != AGG_CONTEXT_AGGREGATE)
-		elog(ERROR, "ordered-set aggregate called in non-aggregate context");
-	/* Need the Aggref as well */
-	aggref = AggGetAggref(fcinfo);
-	if (!aggref)
-		elog(ERROR, "ordered-set aggregate called in non-aggregate context");
-	if (!AGGKIND_IS_ORDERED_SET(aggref->aggkind))
-		elog(ERROR, "ordered-set aggregate support function called for non-ordered-set aggregate");
-	/* Also get output exprcontext so we can register shutdown callback */
-	peraggecontext = AggGetPerAggEContext(fcinfo);
-	if (!peraggecontext)
-		elog(ERROR, "ordered-set aggregate called in non-aggregate context");
+	/*
+	 * We keep a link to the per-query state in fn_extra; if it's not there,
+	 * create it, and do the per-query setup we need.
+	 */
+	qstate = (OSAPerQueryState *) fcinfo->flinfo->fn_extra;
+	if (qstate == NULL)
+	{
+		Aggref	   *aggref;
+		MemoryContext qcontext;
+		MemoryContext gcontext;
+		ExprContext *peraggecontext;
+		List	   *sortlist;
+		int			numSortCols;
 
-	/* Initialize working-state object in the aggregate-lifespan context */
-	osastate = (OrderedSetAggState *)
-		MemoryContextAllocZero(aggcontext, sizeof(OrderedSetAggState));
-	osastate->aggref = aggref;
+		/*
+		 * Check we're called as aggregate, and get the Agg node's
+		 * group-lifespan context
+		 */
+		if (AggCheckCallContext(fcinfo, &gcontext) != AGG_CONTEXT_AGGREGATE)
+			elog(ERROR, "ordered-set aggregate called in non-aggregate context");
+		/* Need the Aggref as well */
+		aggref = AggGetAggref(fcinfo);
+		if (!aggref)
+			elog(ERROR, "ordered-set aggregate called in non-aggregate context");
+		if (!AGGKIND_IS_ORDERED_SET(aggref->aggkind))
+			elog(ERROR, "ordered-set aggregate support function called for non-ordered-set aggregate");
+		/* Also get output exprcontext so we can register shutdown callback */
+		peraggecontext = AggGetPerAggEContext(fcinfo);
+		if (!peraggecontext)
+			elog(ERROR, "ordered-set aggregate called in non-aggregate context");
 
-	/* Extract the sort information */
-	sortlist = aggref->aggorder;
-	numSortCols = list_length(sortlist);
+		/*
+		 * Prepare per-query structures in the fn_mcxt, which we assume is the
+		 * executor's per-query context; in any case it's the right place to
+		 * keep anything found via fn_extra.
+		 */
+		qcontext = fcinfo->flinfo->fn_mcxt;
+		oldcontext = MemoryContextSwitchTo(qcontext);
 
-	if (use_tuples)
-	{
-		bool		ishypothetical = (aggref->aggkind == AGGKIND_HYPOTHETICAL);
-		AttrNumber *sortColIdx;
-		Oid		   *sortOperators;
-		Oid		   *sortCollations;
-		bool	   *sortNullsFirst;
-		ListCell   *lc;
-		int			i;
-
-		if (ishypothetical)
-			numSortCols++;		/* make space for flag column */
-		/* these arrays are made in short-lived context */
-		sortColIdx = (AttrNumber *) palloc(numSortCols * sizeof(AttrNumber));
-		sortOperators = (Oid *) palloc(numSortCols * sizeof(Oid));
-		sortCollations = (Oid *) palloc(numSortCols * sizeof(Oid));
-		sortNullsFirst = (bool *) palloc(numSortCols * sizeof(bool));
-
-		i = 0;
-		foreach(lc, sortlist)
+		qstate = (OSAPerQueryState *) palloc0(sizeof(OSAPerQueryState));
+		qstate->aggref = aggref;
+		qstate->qcontext = qcontext;
+		qstate->gcontext = gcontext;
+		qstate->peraggecontext = peraggecontext;
+
+		/* Extract the sort information */
+		sortlist = aggref->aggorder;
+		numSortCols = list_length(sortlist);
+
+		if (use_tuples)
 		{
-			SortGroupClause *sortcl = (SortGroupClause *) lfirst(lc);
-			TargetEntry *tle = get_sortgroupclause_tle(sortcl, aggref->args);
+			bool		ishypothetical = (aggref->aggkind == AGGKIND_HYPOTHETICAL);
+			AttrNumber *sortColIdx;
+			Oid		   *sortOperators;
+			Oid		   *eqOperators;
+			Oid		   *sortCollations;
+			bool	   *sortNullsFirsts;
+			ListCell   *lc;
+			int			i;
+
+			if (ishypothetical)
+				numSortCols++;	/* make space for flag column */
+			qstate->numSortCols = numSortCols;
+			qstate->sortColIdx = sortColIdx =
+				(AttrNumber *) palloc(numSortCols * sizeof(AttrNumber));
+			qstate->sortOperators = sortOperators =
+				(Oid *) palloc(numSortCols * sizeof(Oid));
+			qstate->eqOperators = eqOperators =
+				(Oid *) palloc(numSortCols * sizeof(Oid));
+			qstate->sortCollations = sortCollations =
+				(Oid *) palloc(numSortCols * sizeof(Oid));
+			qstate->sortNullsFirsts = sortNullsFirsts =
+				(bool *) palloc(numSortCols * sizeof(bool));
+
+			i = 0;
+			foreach(lc, sortlist)
+			{
+				SortGroupClause *sortcl = (SortGroupClause *) lfirst(lc);
+				TargetEntry *tle = get_sortgroupclause_tle(sortcl,
+														   aggref->args);
+
+				/* the parser should have made sure of this */
+				Assert(OidIsValid(sortcl->sortop));
+
+				sortColIdx[i] = tle->resno;
+				sortOperators[i] = sortcl->sortop;
+				eqOperators[i] = sortcl->eqop;
+				sortCollations[i] = exprCollation((Node *) tle->expr);
+				sortNullsFirsts[i] = sortcl->nulls_first;
+				i++;
+			}
 
-			/* the parser should have made sure of this */
-			Assert(OidIsValid(sortcl->sortop));
+			if (ishypothetical)
+			{
+				/* Add an integer flag column as the last sort column */
+				sortColIdx[i] = list_length(aggref->args) + 1;
+				sortOperators[i] = Int4LessOperator;
+				eqOperators[i] = Int4EqualOperator;
+				sortCollations[i] = InvalidOid;
+				sortNullsFirsts[i] = false;
+				i++;
+			}
 
-			sortColIdx[i] = tle->resno;
-			sortOperators[i] = sortcl->sortop;
-			sortCollations[i] = exprCollation((Node *) tle->expr);
-			sortNullsFirst[i] = sortcl->nulls_first;
-			i++;
-		}
+			Assert(i == numSortCols);
 
-		if (ishypothetical)
-		{
-			/* Add an integer flag column as the last sort column */
-			sortColIdx[i] = list_length(aggref->args) + 1;
-			sortOperators[i] = Int4LessOperator;
-			sortCollations[i] = InvalidOid;
-			sortNullsFirst[i] = false;
-			i++;
+			/*
+			 * Get a tupledesc corresponding to the aggregated inputs
+			 * (including sort expressions) of the agg.
+			 */
+			qstate->tupdesc = ExecTypeFromTL(aggref->args, false);
+
+			/* If we need a flag column, hack the tupledesc to include that */
+			if (ishypothetical)
+			{
+				TupleDesc	newdesc;
+				int			natts = qstate->tupdesc->natts;
+
+				newdesc = CreateTemplateTupleDesc(natts + 1, false);
+				for (i = 1; i <= natts; i++)
+					TupleDescCopyEntry(newdesc, i, qstate->tupdesc, i);
+
+				TupleDescInitEntry(newdesc,
+								   (AttrNumber) ++natts,
+								   "flag",
+								   INT4OID,
+								   -1,
+								   0);
+
+				FreeTupleDesc(qstate->tupdesc);
+				qstate->tupdesc = newdesc;
+			}
+
+			/* Create slot we'll use to store/retrieve rows */
+			qstate->tupslot = MakeSingleTupleTableSlot(qstate->tupdesc);
 		}
+		else
+		{
+			/* Sort single datums */
+			SortGroupClause *sortcl;
+			TargetEntry *tle;
 
-		Assert(i == numSortCols);
+			if (numSortCols != 1 || aggref->aggkind == AGGKIND_HYPOTHETICAL)
+				elog(ERROR, "ordered-set aggregate support function does not support multiple aggregated columns");
 
-		/* Now build the stuff we need in aggregate-lifespan context */
-		oldcontext = MemoryContextSwitchTo(aggcontext);
+			sortcl = (SortGroupClause *) linitial(sortlist);
+			tle = get_sortgroupclause_tle(sortcl, aggref->args);
 
-		/*
-		 * Get a tupledesc corresponding to the aggregated inputs (including
-		 * sort expressions) of the agg.
-		 */
-		osastate->tupdesc = ExecTypeFromTL(aggref->args, false);
+			/* the parser should have made sure of this */
+			Assert(OidIsValid(sortcl->sortop));
 
-		/* If we need a flag column, hack the tupledesc to include that */
-		if (ishypothetical)
-		{
-			TupleDesc	newdesc;
-			int			natts = osastate->tupdesc->natts;
-
-			newdesc = CreateTemplateTupleDesc(natts + 1, false);
-			for (i = 1; i <= natts; i++)
-				TupleDescCopyEntry(newdesc, i, osastate->tupdesc, i);
-
-			TupleDescInitEntry(newdesc,
-							   (AttrNumber) ++natts,
-							   "flag",
-							   INT4OID,
-							   -1,
-							   0);
-
-			FreeTupleDesc(osastate->tupdesc);
-			osastate->tupdesc = newdesc;
+			/* Save sort ordering info */
+			qstate->sortColType = exprType((Node *) tle->expr);
+			qstate->sortOperator = sortcl->sortop;
+			qstate->eqOperator = sortcl->eqop;
+			qstate->sortCollation = exprCollation((Node *) tle->expr);
+			qstate->sortNullsFirst = sortcl->nulls_first;
+
+			/* Save datatype info */
+			get_typlenbyvalalign(qstate->sortColType,
+								 &qstate->typLen,
+								 &qstate->typByVal,
+								 &qstate->typAlign);
 		}
 
-		/* Initialize tuplesort object */
-		osastate->sortstate = tuplesort_begin_heap(osastate->tupdesc,
-												   numSortCols,
-												   sortColIdx,
-												   sortOperators,
-												   sortCollations,
-												   sortNullsFirst,
-												   work_mem, false);
+		fcinfo->flinfo->fn_extra = (void *) qstate;
 
-		/* Create slot we'll use to store/retrieve rows */
-		osastate->tupslot = MakeSingleTupleTableSlot(osastate->tupdesc);
+		MemoryContextSwitchTo(oldcontext);
 	}
+
+	/* Now build the stuff we need in group-lifespan context */
+	oldcontext = MemoryContextSwitchTo(qstate->gcontext);
+
+	osastate = (OSAPerGroupState *) palloc(sizeof(OSAPerGroupState));
+	osastate->qstate = qstate;
+
+	/* Initialize tuplesort object */
+	if (use_tuples)
+		osastate->sortstate = tuplesort_begin_heap(qstate->tupdesc,
+												   qstate->numSortCols,
+												   qstate->sortColIdx,
+												   qstate->sortOperators,
+												   qstate->sortCollations,
+												   qstate->sortNullsFirsts,
+												   work_mem, false);
 	else
-	{
-		/* Sort single datums */
-		SortGroupClause *sortcl;
-		TargetEntry *tle;
-		Oid			sortColType;
-		Oid			sortOperator;
-		Oid			eqOperator;
-		Oid			sortCollation;
-		bool		sortNullsFirst;
-
-		if (numSortCols != 1 || aggref->aggkind == AGGKIND_HYPOTHETICAL)
-			elog(ERROR, "ordered-set aggregate support function does not support multiple aggregated columns");
-
-		sortcl = (SortGroupClause *) linitial(sortlist);
-		tle = get_sortgroupclause_tle(sortcl, aggref->args);
-
-		/* the parser should have made sure of this */
-		Assert(OidIsValid(sortcl->sortop));
-
-		sortColType = exprType((Node *) tle->expr);
-		sortOperator = sortcl->sortop;
-		eqOperator = sortcl->eqop;
-		sortCollation = exprCollation((Node *) tle->expr);
-		sortNullsFirst = sortcl->nulls_first;
-
-		/* Save datatype info */
-		osastate->datumtype = sortColType;
-		get_typlenbyvalalign(sortColType,
-							 &osastate->typLen,
-							 &osastate->typByVal,
-							 &osastate->typAlign);
-		osastate->eqOperator = eqOperator;
-
-		/* Now build the stuff we need in aggregate-lifespan context */
-		oldcontext = MemoryContextSwitchTo(aggcontext);
-
-		/* Initialize tuplesort object */
-		osastate->sortstate = tuplesort_begin_datum(sortColType,
-													sortOperator,
-													sortCollation,
-													sortNullsFirst,
+		osastate->sortstate = tuplesort_begin_datum(qstate->sortColType,
+													qstate->sortOperator,
+													qstate->sortCollation,
+													qstate->sortNullsFirst,
 													work_mem, false);
-	}
 
-	/* Now register a shutdown callback to clean it all up */
-	RegisterExprContextCallback(peraggecontext,
+	osastate->number_of_rows = 0;
+
+	/* Now register a shutdown callback to clean things up */
+	RegisterExprContextCallback(qstate->peraggecontext,
 								ordered_set_shutdown,
 								PointerGetDatum(osastate));
 
@@ -247,9 +313,11 @@ ordered_set_startup(FunctionCallInfo fcinfo, bool use_tuples)
 /*
  * Clean up when evaluation of an ordered-set aggregate is complete.
  *
- * We don't need to bother freeing objects in the aggcontext memory context,
- * since that will get reset anyway by nodeAgg.c, but we should take care to
- * release any potential non-memory resources.
+ * We don't need to bother freeing objects in the per-group memory context,
+ * since that will get reset anyway by nodeAgg.c; nor should we free anything
+ * in the per-query context, which will get cleared (if this was the last
+ * group) by ExecutorEnd.  But we must take care to release any potential
+ * non-memory resources.
  *
  * This callback is arguably unnecessary, since we don't support use of
  * ordered-set aggs in AGG_HASHED mode and there is currently no non-error
@@ -263,16 +331,15 @@ ordered_set_startup(FunctionCallInfo fcinfo, bool use_tuples)
 static void
 ordered_set_shutdown(Datum arg)
 {
-	OrderedSetAggState *osastate = (OrderedSetAggState *) DatumGetPointer(arg);
+	OSAPerGroupState *osastate = (OSAPerGroupState *) DatumGetPointer(arg);
 
 	/* Tuplesort object might have temp files. */
 	if (osastate->sortstate)
 		tuplesort_end(osastate->sortstate);
 	osastate->sortstate = NULL;
 	/* The tupleslot probably can't be holding a pin, but let's be safe. */
-	if (osastate->tupslot)
-		ExecDropSingleTupleTableSlot(osastate->tupslot);
-	osastate->tupslot = NULL;
+	if (osastate->qstate->tupslot)
+		ExecClearTuple(osastate->qstate->tupslot);
 }
 
 
@@ -283,7 +350,7 @@ ordered_set_shutdown(Datum arg)
 Datum
 ordered_set_transition(PG_FUNCTION_ARGS)
 {
-	OrderedSetAggState *osastate;
+	OSAPerGroupState *osastate;
 
 	/* If first call, create the transition state workspace */
 	if (PG_ARGISNULL(0))
@@ -293,7 +360,7 @@ ordered_set_transition(PG_FUNCTION_ARGS)
 		/* safety check */
 		if (AggCheckCallContext(fcinfo, NULL) != AGG_CONTEXT_AGGREGATE)
 			elog(ERROR, "ordered-set aggregate called in non-aggregate context");
-		osastate = (OrderedSetAggState *) PG_GETARG_POINTER(0);
+		osastate = (OSAPerGroupState *) PG_GETARG_POINTER(0);
 	}
 
 	/* Load the datum into the tuplesort object, but only if it's not null */
@@ -313,7 +380,7 @@ ordered_set_transition(PG_FUNCTION_ARGS)
 Datum
 ordered_set_transition_multi(PG_FUNCTION_ARGS)
 {
-	OrderedSetAggState *osastate;
+	OSAPerGroupState *osastate;
 	TupleTableSlot *slot;
 	int			nargs;
 	int			i;
@@ -326,11 +393,11 @@ ordered_set_transition_multi(PG_FUNCTION_ARGS)
 		/* safety check */
 		if (AggCheckCallContext(fcinfo, NULL) != AGG_CONTEXT_AGGREGATE)
 			elog(ERROR, "ordered-set aggregate called in non-aggregate context");
-		osastate = (OrderedSetAggState *) PG_GETARG_POINTER(0);
+		osastate = (OSAPerGroupState *) PG_GETARG_POINTER(0);
 	}
 
 	/* Form a tuple from all the other inputs besides the transition value */
-	slot = osastate->tupslot;
+	slot = osastate->qstate->tupslot;
 	ExecClearTuple(slot);
 	nargs = PG_NARGS() - 1;
 	for (i = 0; i < nargs; i++)
@@ -338,7 +405,7 @@ ordered_set_transition_multi(PG_FUNCTION_ARGS)
 		slot->tts_values[i] = PG_GETARG_DATUM(i + 1);
 		slot->tts_isnull[i] = PG_ARGISNULL(i + 1);
 	}
-	if (osastate->aggref->aggkind == AGGKIND_HYPOTHETICAL)
+	if (osastate->qstate->aggref->aggkind == AGGKIND_HYPOTHETICAL)
 	{
 		/* Add a zero flag value to mark this row as a normal input row */
 		slot->tts_values[i] = Int32GetDatum(0);
@@ -362,7 +429,7 @@ ordered_set_transition_multi(PG_FUNCTION_ARGS)
 Datum
 percentile_disc_final(PG_FUNCTION_ARGS)
 {
-	OrderedSetAggState *osastate;
+	OSAPerGroupState *osastate;
 	double		percentile;
 	Datum		val;
 	bool		isnull;
@@ -388,7 +455,7 @@ percentile_disc_final(PG_FUNCTION_ARGS)
 	if (PG_ARGISNULL(0))
 		PG_RETURN_NULL();
 
-	osastate = (OrderedSetAggState *) PG_GETARG_POINTER(0);
+	osastate = (OSAPerGroupState *) PG_GETARG_POINTER(0);
 
 	/* number_of_rows could be zero if we only saw NULL input values */
 	if (osastate->number_of_rows == 0)
@@ -465,7 +532,7 @@ percentile_cont_final_common(FunctionCallInfo fcinfo,
 							 Oid expect_type,
 							 LerpFunc lerpfunc)
 {
-	OrderedSetAggState *osastate;
+	OSAPerGroupState *osastate;
 	double		percentile;
 	int64		first_row = 0;
 	int64		second_row = 0;
@@ -495,13 +562,13 @@ percentile_cont_final_common(FunctionCallInfo fcinfo,
 	if (PG_ARGISNULL(0))
 		PG_RETURN_NULL();
 
-	osastate = (OrderedSetAggState *) PG_GETARG_POINTER(0);
+	osastate = (OSAPerGroupState *) PG_GETARG_POINTER(0);
 
 	/* number_of_rows could be zero if we only saw NULL input values */
 	if (osastate->number_of_rows == 0)
 		PG_RETURN_NULL();
 
-	Assert(expect_type == osastate->datumtype);
+	Assert(expect_type == osastate->qstate->sortColType);
 
 	/* Finish the sort */
 	tuplesort_performsort(osastate->sortstate);
@@ -672,7 +739,7 @@ setup_pct_info(int num_percentiles,
 Datum
 percentile_disc_multi_final(PG_FUNCTION_ARGS)
 {
-	OrderedSetAggState *osastate;
+	OSAPerGroupState *osastate;
 	ArrayType  *param;
 	Datum	   *percentiles_datum;
 	bool	   *percentiles_null;
@@ -693,7 +760,7 @@ percentile_disc_multi_final(PG_FUNCTION_ARGS)
 	if (PG_ARGISNULL(0))
 		PG_RETURN_NULL();
 
-	osastate = (OrderedSetAggState *) PG_GETARG_POINTER(0);
+	osastate = (OSAPerGroupState *) PG_GETARG_POINTER(0);
 
 	/* number_of_rows could be zero if we only saw NULL input values */
 	if (osastate->number_of_rows == 0)
@@ -712,7 +779,7 @@ percentile_disc_multi_final(PG_FUNCTION_ARGS)
 					  &num_percentiles);
 
 	if (num_percentiles == 0)
-		PG_RETURN_POINTER(construct_empty_array(osastate->datumtype));
+		PG_RETURN_POINTER(construct_empty_array(osastate->qstate->sortColType));
 
 	pct_info = setup_pct_info(num_percentiles,
 							  percentiles_datum,
@@ -779,10 +846,10 @@ percentile_disc_multi_final(PG_FUNCTION_ARGS)
 										 ARR_NDIM(param),
 										 ARR_DIMS(param),
 										 ARR_LBOUND(param),
-										 osastate->datumtype,
-										 osastate->typLen,
-										 osastate->typByVal,
-										 osastate->typAlign));
+										 osastate->qstate->sortColType,
+										 osastate->qstate->typLen,
+										 osastate->qstate->typByVal,
+										 osastate->qstate->typAlign));
 }
 
 /*
@@ -794,7 +861,7 @@ percentile_cont_multi_final_common(FunctionCallInfo fcinfo,
 								   int16 typLen, bool typByVal, char typAlign,
 								   LerpFunc lerpfunc)
 {
-	OrderedSetAggState *osastate;
+	OSAPerGroupState *osastate;
 	ArrayType  *param;
 	Datum	   *percentiles_datum;
 	bool	   *percentiles_null;
@@ -816,13 +883,13 @@ percentile_cont_multi_final_common(FunctionCallInfo fcinfo,
 	if (PG_ARGISNULL(0))
 		PG_RETURN_NULL();
 
-	osastate = (OrderedSetAggState *) PG_GETARG_POINTER(0);
+	osastate = (OSAPerGroupState *) PG_GETARG_POINTER(0);
 
 	/* number_of_rows could be zero if we only saw NULL input values */
 	if (osastate->number_of_rows == 0)
 		PG_RETURN_NULL();
 
-	Assert(expect_type == osastate->datumtype);
+	Assert(expect_type == osastate->qstate->sortColType);
 
 	/* Deconstruct the percentile-array input */
 	if (PG_ARGISNULL(1))
@@ -837,7 +904,7 @@ percentile_cont_multi_final_common(FunctionCallInfo fcinfo,
 					  &num_percentiles);
 
 	if (num_percentiles == 0)
-		PG_RETURN_POINTER(construct_empty_array(osastate->datumtype));
+		PG_RETURN_POINTER(construct_empty_array(osastate->qstate->sortColType));
 
 	pct_info = setup_pct_info(num_percentiles,
 							  percentiles_datum,
@@ -967,7 +1034,7 @@ percentile_cont_interval_multi_final(PG_FUNCTION_ARGS)
 Datum
 mode_final(PG_FUNCTION_ARGS)
 {
-	OrderedSetAggState *osastate;
+	OSAPerGroupState *osastate;
 	Datum		val;
 	bool		isnull;
 	Datum		mode_val = (Datum) 0;
@@ -975,7 +1042,7 @@ mode_final(PG_FUNCTION_ARGS)
 	Datum		last_val = (Datum) 0;
 	int64		last_val_freq = 0;
 	bool		last_val_is_mode = false;
-	FmgrInfo	equalfn;
+	FmgrInfo   *equalfn;
 	bool		shouldfree;
 
 	/* safety check */
@@ -986,16 +1053,19 @@ mode_final(PG_FUNCTION_ARGS)
 	if (PG_ARGISNULL(0))
 		PG_RETURN_NULL();
 
-	osastate = (OrderedSetAggState *) PG_GETARG_POINTER(0);
+	osastate = (OSAPerGroupState *) PG_GETARG_POINTER(0);
 
 	/* number_of_rows could be zero if we only saw NULL input values */
 	if (osastate->number_of_rows == 0)
 		PG_RETURN_NULL();
 
-	/* Look up the equality function for the datatype */
-	fmgr_info(get_opcode(osastate->eqOperator), &equalfn);
+	/* Look up the equality function for the datatype, if we didn't already */
+	equalfn = &(osastate->qstate->equalfn);
+	if (!OidIsValid(equalfn->fn_oid))
+		fmgr_info_cxt(get_opcode(osastate->qstate->eqOperator), equalfn,
+					  osastate->qstate->qcontext);
 
-	shouldfree = !(osastate->typByVal);
+	shouldfree = !(osastate->qstate->typByVal);
 
 	/* Finish the sort */
 	tuplesort_performsort(osastate->sortstate);
@@ -1014,7 +1084,7 @@ mode_final(PG_FUNCTION_ARGS)
 			mode_freq = last_val_freq = 1;
 			last_val_is_mode = true;
 		}
-		else if (DatumGetBool(FunctionCall2(&equalfn, val, last_val)))
+		else if (DatumGetBool(FunctionCall2(equalfn, val, last_val)))
 		{
 			/* value equal to previous value, count it */
 			if (last_val_is_mode)
@@ -1099,7 +1169,7 @@ hypothetical_rank_common(FunctionCallInfo fcinfo, int flag,
 {
 	int			nargs = PG_NARGS() - 1;
 	int64		rank = 1;
-	OrderedSetAggState *osastate;
+	OSAPerGroupState *osastate;
 	TupleTableSlot *slot;
 	int			i;
 
@@ -1114,7 +1184,7 @@ hypothetical_rank_common(FunctionCallInfo fcinfo, int flag,
 		return 1;
 	}
 
-	osastate = (OrderedSetAggState *) PG_GETARG_POINTER(0);
+	osastate = (OSAPerGroupState *) PG_GETARG_POINTER(0);
 	*number_of_rows = osastate->number_of_rows;
 
 	/* Adjust nargs to be the number of direct (or aggregated) args */
@@ -1122,10 +1192,10 @@ hypothetical_rank_common(FunctionCallInfo fcinfo, int flag,
 		elog(ERROR, "wrong number of arguments in hypothetical-set function");
 	nargs /= 2;
 
-	hypothetical_check_argtypes(fcinfo, nargs, osastate->tupdesc);
+	hypothetical_check_argtypes(fcinfo, nargs, osastate->qstate->tupdesc);
 
 	/* insert the hypothetical row into the sort */
-	slot = osastate->tupslot;
+	slot = osastate->qstate->tupslot;
 	ExecClearTuple(slot);
 	for (i = 0; i < nargs; i++)
 	{
@@ -1225,8 +1295,7 @@ hypothetical_dense_rank_final(PG_FUNCTION_ARGS)
 	int			nargs = PG_NARGS() - 1;
 	int64		rank = 1;
 	int64		duplicate_count = 0;
-	OrderedSetAggState *osastate;
-	List	   *sortlist;
+	OSAPerGroupState *osastate;
 	int			numDistinctCols;
 	AttrNumber *sortColIdx;
 	FmgrInfo   *equalfns;
@@ -1234,7 +1303,6 @@ hypothetical_dense_rank_final(PG_FUNCTION_ARGS)
 	TupleTableSlot *extraslot;
 	TupleTableSlot *slot2;
 	MemoryContext tmpcontext;
-	ListCell   *lc;
 	int			i;
 
 	/* safety check */
@@ -1245,41 +1313,44 @@ hypothetical_dense_rank_final(PG_FUNCTION_ARGS)
 	if (PG_ARGISNULL(0))
 		PG_RETURN_INT64(rank);
 
-	osastate = (OrderedSetAggState *) PG_GETARG_POINTER(0);
+	osastate = (OSAPerGroupState *) PG_GETARG_POINTER(0);
 
 	/* Adjust nargs to be the number of direct (or aggregated) args */
 	if (nargs % 2 != 0)
 		elog(ERROR, "wrong number of arguments in hypothetical-set function");
 	nargs /= 2;
 
-	hypothetical_check_argtypes(fcinfo, nargs, osastate->tupdesc);
+	hypothetical_check_argtypes(fcinfo, nargs, osastate->qstate->tupdesc);
 
 	/*
-	 * Construct list of columns to compare for uniqueness.  We can omit the
-	 * flag column since we will only compare rows with flag == 0.
+	 * When comparing tuples, we can omit the flag column since we will only
+	 * compare rows with flag == 0.
 	 */
-	sortlist = osastate->aggref->aggorder;
-	numDistinctCols = list_length(sortlist);
-	sortColIdx = (AttrNumber *) palloc(numDistinctCols * sizeof(AttrNumber));
-	equalfns = (FmgrInfo *) palloc(numDistinctCols * sizeof(FmgrInfo));
+	numDistinctCols = osastate->qstate->numSortCols - 1;
 
-	i = 0;
-	foreach(lc, sortlist)
+	/* Look up the equality function(s), if we didn't already */
+	equalfns = osastate->qstate->equalfns;
+	if (equalfns == NULL)
 	{
-		SortGroupClause *sortcl = (SortGroupClause *) lfirst(lc);
-		TargetEntry *tle = get_sortgroupclause_tle(sortcl,
-												   osastate->aggref->args);
+		MemoryContext qcontext = osastate->qstate->qcontext;
 
-		sortColIdx[i] = tle->resno;
-		fmgr_info(get_opcode(sortcl->eqop), &equalfns[i]);
-		i++;
+		equalfns = (FmgrInfo *)
+			MemoryContextAlloc(qcontext, numDistinctCols * sizeof(FmgrInfo));
+		for (i = 0; i < numDistinctCols; i++)
+		{
+			fmgr_info_cxt(get_opcode(osastate->qstate->eqOperators[i]),
+						  &equalfns[i],
+						  qcontext);
+		}
+		osastate->qstate->equalfns = equalfns;
 	}
+	sortColIdx = osastate->qstate->sortColIdx;
 
 	/* Get short-term context we can use for execTuplesMatch */
 	tmpcontext = AggGetPerTupleEContext(fcinfo)->ecxt_per_tuple_memory;
 
 	/* insert the hypothetical row into the sort */
-	slot = osastate->tupslot;
+	slot = osastate->qstate->tupslot;
 	ExecClearTuple(slot);
 	for (i = 0; i < nargs; i++)
 	{
@@ -1296,11 +1367,11 @@ hypothetical_dense_rank_final(PG_FUNCTION_ARGS)
 	tuplesort_performsort(osastate->sortstate);
 
 	/*
-	 * We alternate fetching into osastate->tupslot and extraslot so that we
-	 * have the previous row available for comparisons.  This is accomplished
-	 * by swapping the slot pointer variables after each row.
+	 * We alternate fetching into tupslot and extraslot so that we have the
+	 * previous row available for comparisons.	This is accomplished by
+	 * swapping the slot pointer variables after each row.
 	 */
-	extraslot = MakeSingleTupleTableSlot(osastate->tupdesc);
+	extraslot = MakeSingleTupleTableSlot(osastate->qstate->tupdesc);
 	slot2 = extraslot;
 
 	/* iterate till we find the hypothetical row */

src/include/catalog/pg_operator.h

@@ -128,6 +128,7 @@ DATA(insert OID =  95 ( "<"		   PGNSP PGUID b f f	21	21	16 520 524 int2lt scalar
 DESCR("less than");
 DATA(insert OID =  96 ( "="		   PGNSP PGUID b t t	23	23	16	96 518 int4eq eqsel eqjoinsel ));
 DESCR("equal");
+#define Int4EqualOperator	96
 DATA(insert OID =  97 ( "<"		   PGNSP PGUID b f f	23	23	16 521 525 int4lt scalarltsel scalarltjoinsel ));
 DESCR("less than");
 #define Int4LessOperator	97