Aggregate function API versus grouping sets

diff --git a/src/backend/executor/nodeAgg.c b/src/backend/executor/nodeAgg.c
index 09ff035..0388735 100644
--- a/src/backend/executor/nodeAgg.c
+++ b/src/backend/executor/nodeAgg.c
@@ -2199,44 +2199,56 @@ AggGetAggref(FunctionCallInfo fcinfo)
 }
 
 /*
- * AggGetPerTupleEContext - fetch per-input-tuple ExprContext
+ * AggGetTempMemoryContext - fetch short-term memory context
  *
- * This is useful in agg final functions; the econtext returned is the
- * same per-tuple context that the transfn was called in (which can
- * safely get reset during the final function).
+ * This is useful in agg final functions; the context returned is one that the
+ * final function can safely reset as desired. This isn't useful for transition
+ * functions, since the context returned MAY (we don't promise) be the same as
+ * the context those are called in.
  *
  * As above, this is currently not useful for aggs called as window functions.
  */
-ExprContext *
-AggGetPerTupleEContext(FunctionCallInfo fcinfo)
+MemoryContext
+AggGetTempMemoryContext(FunctionCallInfo fcinfo)
 {
 	if (fcinfo->context && IsA(fcinfo->context, AggState))
 	{
 		AggState   *aggstate = (AggState *) fcinfo->context;
 
-		return aggstate->tmpcontext;
+		return aggstate->tmpcontext->ecxt_per_tuple_memory;
 	}
 	return NULL;
 }
 
 /*
- * AggGetPerAggEContext - fetch per-output-tuple ExprContext
+ * AggRegisterCallback - register a cleanup callback linked to aggcontext
  *
  * This is useful for aggs to register shutdown callbacks, which will ensure
- * that non-memory resources are freed.
+ * that non-memory resources are freed.  These callbacks will occur just before
+ * the associated aggcontext (as returned by AggCheckCallContext) is reset,
+ * either between groups or as a result of rescanning the query.  They will NOT
+ * be called on error paths.  The primary intent is to allow for freeing of
+ * tuplestores or tuplesorts maintained in aggcontext, or pins held by slots
+ * created by the agg functions.  (They will not be called until after the
+ * result of the finalfn is no longer needed, so it's safe for the finalfn to
+ * return data that will be freed by the callback.)
  *
  * As above, this is currently not useful for aggs called as window functions.
  */
-ExprContext *
-AggGetPerAggEContext(FunctionCallInfo fcinfo)
+void
+AggRegisterCallback(FunctionCallInfo fcinfo,
+					ExprContextCallbackFunction func,
+					Datum arg)
 {
 	if (fcinfo->context && IsA(fcinfo->context, AggState))
 	{
 		AggState   *aggstate = (AggState *) fcinfo->context;
 
-		return aggstate->ss.ps.ps_ExprContext;
+		RegisterExprContextCallback(aggstate->ss.ps.ps_ExprContext, func, arg);
+
+		return;
 	}
-	return NULL;
+	elog(ERROR,"Aggregate function cannot register a callback in this context");
 }
 
 
diff --git a/src/backend/utils/adt/orderedsetaggs.c b/src/backend/utils/adt/orderedsetaggs.c
index efb0411..d116a63 100644
--- a/src/backend/utils/adt/orderedsetaggs.c
+++ b/src/backend/utils/adt/orderedsetaggs.c
@@ -49,8 +49,6 @@ typedef struct OSAPerQueryState
 	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: */
 
@@ -117,7 +115,6 @@ ordered_set_startup(FunctionCallInfo fcinfo, bool use_tuples)
 		Aggref	   *aggref;
 		MemoryContext qcontext;
 		MemoryContext gcontext;
-		ExprContext *peraggecontext;
 		List	   *sortlist;
 		int			numSortCols;
 
@@ -133,10 +130,6 @@ ordered_set_startup(FunctionCallInfo fcinfo, bool use_tuples)
 			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");
 
 		/*
 		 * Prepare per-query structures in the fn_mcxt, which we assume is the
@@ -150,7 +143,6 @@ ordered_set_startup(FunctionCallInfo fcinfo, bool use_tuples)
 		qstate->aggref = aggref;
 		qstate->qcontext = qcontext;
 		qstate->gcontext = gcontext;
-		qstate->peraggecontext = peraggecontext;
 
 		/* Extract the sort information */
 		sortlist = aggref->aggorder;
@@ -291,9 +283,9 @@ ordered_set_startup(FunctionCallInfo fcinfo, bool use_tuples)
 	osastate->number_of_rows = 0;
 
 	/* Now register a shutdown callback to clean things up */
-	RegisterExprContextCallback(qstate->peraggecontext,
-								ordered_set_shutdown,
-								PointerGetDatum(osastate));
+	AggRegisterCallback(fcinfo,
+						ordered_set_shutdown,
+						PointerGetDatum(osastate));
 
 	MemoryContextSwitchTo(oldcontext);
 
@@ -1310,7 +1302,7 @@ hypothetical_dense_rank_final(PG_FUNCTION_ARGS)
 	sortColIdx = osastate->qstate->sortColIdx;
 
 	/* Get short-term context we can use for execTuplesMatch */
-	tmpcontext = AggGetPerTupleEContext(fcinfo)->ecxt_per_tuple_memory;
+	tmpcontext = AggGetTempMemoryContext(fcinfo);
 
 	/* insert the hypothetical row into the sort */
 	slot = osastate->qstate->tupslot;
diff --git a/src/include/fmgr.h b/src/include/fmgr.h
index 267403c..9324dba 100644
--- a/src/include/fmgr.h
+++ b/src/include/fmgr.h
@@ -24,6 +24,7 @@ typedef struct Aggref *fmAggrefPtr;
 
 /* Likewise, avoid including execnodes.h here */
 typedef struct ExprContext *fmExprContextPtr;
+typedef void (*fmExprContextCallbackFunction) (Datum arg);
 
 /* Likewise, avoid including stringinfo.h here */
 typedef struct StringInfoData *fmStringInfo;
@@ -656,8 +657,10 @@ extern void **find_rendezvous_variable(const char *varName);
 extern int AggCheckCallContext(FunctionCallInfo fcinfo,
 					MemoryContext *aggcontext);
 extern fmAggrefPtr AggGetAggref(FunctionCallInfo fcinfo);
-extern fmExprContextPtr AggGetPerTupleEContext(FunctionCallInfo fcinfo);
-extern fmExprContextPtr AggGetPerAggEContext(FunctionCallInfo fcinfo);
+extern MemoryContext AggGetTempMemoryContext(FunctionCallInfo fcinfo);
+extern void AggRegisterCallback(FunctionCallInfo fcinfo,
+					fmExprContextCallbackFunction func,
+					Datum arg);
 
 /*
  * We allow plugin modules to hook function entry/exit.  This is intended

diff --git a/src/backend/utils/adt/orderedsetaggs.c b/src/backend/utils/adt/orderedsetaggs.c
index d116a63..92fa9f3 100644
--- a/src/backend/utils/adt/orderedsetaggs.c
+++ b/src/backend/utils/adt/orderedsetaggs.c
@@ -47,8 +47,6 @@ typedef struct OSAPerQueryState
 	Aggref	   *aggref;
 	/* Memory context containing this struct and other per-query data: */
 	MemoryContext qcontext;
-	/* Memory context containing per-group data: */
-	MemoryContext gcontext;
 
 	/* These fields are used only when accumulating tuples: */
 
@@ -86,6 +84,8 @@ typedef struct OSAPerGroupState
 {
 	/* Link to the per-query state for this aggregate: */
 	OSAPerQueryState *qstate;
+	/* MemoryContext for per-group data */
+	MemoryContext gcontext;
 	/* Sort object we're accumulating data in: */
 	Tuplesortstate *sortstate;
 	/* Number of normal rows inserted into sortstate: */
@@ -104,6 +104,15 @@ ordered_set_startup(FunctionCallInfo fcinfo, bool use_tuples)
 	OSAPerGroupState *osastate;
 	OSAPerQueryState *qstate;
 	MemoryContext oldcontext;
+	MemoryContext gcontext;
+
+	/*
+	 * Check we're called as aggregate (and not a window function), and
+	 * get the Agg node's group-lifespan context (which might not be the
+	 * same throughout the query, but will be the same for each transval)
+	 */
+	if (AggCheckCallContext(fcinfo, &gcontext) != AGG_CONTEXT_AGGREGATE)
+		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,
@@ -114,16 +123,9 @@ ordered_set_startup(FunctionCallInfo fcinfo, bool use_tuples)
 	{
 		Aggref	   *aggref;
 		MemoryContext qcontext;
-		MemoryContext gcontext;
 		List	   *sortlist;
 		int			numSortCols;
 
-		/*
-		 * Check we're called as aggregate (and not a window function), 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)
@@ -142,7 +144,6 @@ ordered_set_startup(FunctionCallInfo fcinfo, bool use_tuples)
 		qstate = (OSAPerQueryState *) palloc0(sizeof(OSAPerQueryState));
 		qstate->aggref = aggref;
 		qstate->qcontext = qcontext;
-		qstate->gcontext = gcontext;
 
 		/* Extract the sort information */
 		sortlist = aggref->aggorder;
@@ -259,10 +260,11 @@ ordered_set_startup(FunctionCallInfo fcinfo, bool use_tuples)
 	}
 
 	/* Now build the stuff we need in group-lifespan context */
-	oldcontext = MemoryContextSwitchTo(qstate->gcontext);
+	oldcontext = MemoryContextSwitchTo(gcontext);
 
 	osastate = (OSAPerGroupState *) palloc(sizeof(OSAPerGroupState));
 	osastate->qstate = qstate;
+	osastate->gcontext = gcontext;
 
 	/* Initialize tuplesort object */
 	if (use_tuples)

diff --git a/src/backend/utils/adt/orderedsetaggs.c b/src/backend/utils/adt/orderedsetaggs.c
index 92fa9f3..bbad0e5 100644
--- a/src/backend/utils/adt/orderedsetaggs.c
+++ b/src/backend/utils/adt/orderedsetaggs.c
@@ -16,6 +16,8 @@
 
 #include <math.h>
 
+#include "utils/orderedset.h"
+
 #include "catalog/pg_aggregate.h"
 #include "catalog/pg_operator.h"
 #include "catalog/pg_type.h"
@@ -30,71 +32,8 @@
 #include "utils/tuplesort.h"
 
 
-/*
- * 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 OSAPerQueryState
-{
-	/* Aggref for this aggregate: */
-	Aggref	   *aggref;
-	/* Memory context containing this struct and other per-query data: */
-	MemoryContext qcontext;
-
-	/* These fields are used only when accumulating tuples: */
-
-	/* Tuple descriptor for tuples inserted into sortstate: */
-	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			sortColType;
-	int16		typLen;
-	bool		typByVal;
-	char		typAlign;
-	/* Info about sort ordering: */
-	Oid			sortOperator;
-	Oid			eqOperator;
-	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;
-	/* MemoryContext for per-group data */
-	MemoryContext gcontext;
-	/* 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);
 
-
 /*
  * Set up working state for an ordered-set aggregate
  */
diff --git a/src/include/utils/orderedset.h b/src/include/utils/orderedset.h
new file mode 100644
index 0000000..ab1113e
--- /dev/null
+++ b/src/include/utils/orderedset.h
@@ -0,0 +1,82 @@
+/*-------------------------------------------------------------------------
+ *
+ * orderedset.h
+ *	  Generic support for ordered-set aggregates
+ *
+ * Ordered-set aggregates are free to supply their own transition functions,
+ * but it is expected that the same transition function will serve for almost
+ * all imaginable cases including the built-in ones.  We therefore expose here
+ * the structures used by the builtin ordered_set_transition and
+ * ordered_set_transition_multi functions.
+ *
+ * User-defined ordered-set aggs can therefore use the stock transition
+ * function and supply their own final function to interpret the sorted data.
+ *
+ * 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().
+ */
+#ifndef ORDEREDSET_H
+#define ORDEREDSET_H
+
+#include "nodes/primnodes.h"
+#include "access/tupdesc.h"
+#include "executor/tuptable.h"
+#include "utils/tuplesort.h"
+#include "fmgr.h"
+
+typedef struct OSAPerQueryState
+{
+	/* Aggref for this aggregate: */
+	Aggref	   *aggref;
+	/* Memory context containing this struct and other per-query data: */
+	MemoryContext qcontext;
+
+	/* These fields are used only when accumulating tuples: */
+
+	/* Tuple descriptor for tuples inserted into sortstate: */
+	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			sortColType;
+	int16		typLen;
+	bool		typByVal;
+	char		typAlign;
+	/* Info about sort ordering: */
+	Oid			sortOperator;
+	Oid			eqOperator;
+	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;
+	/* MemoryContext for per-group data */
+	MemoryContext gcontext;
+	/* Sort object we're accumulating data in: */
+	Tuplesortstate *sortstate;
+	/* Number of normal rows inserted into sortstate: */
+	int64		number_of_rows;
+} OSAPerGroupState;
+
+#endif   /* ORDEREDSET_H */