Redesign API presented by nodeAgg.c for ordered-set and similar aggregates.

tglsfdc · tglsfdc · commit 6f5034eda05c · 2014-07-03T18:25:33.000-04:00
The previous design exposed the input and output ExprContexts of the
Agg plan node, but work on grouping sets has suggested that we'll regret
doing that.  Instead provide more narrowly-defined APIs that can be
implemented in multiple ways, namely a way to get a short-term memory
context and a way to register an aggregate shutdown callback.

Back-patch to 9.4 where the bad APIs were introduced, since we don't
want third-party code using these APIs and then having to change in 9.5.

Andrew Gierth
diff --git 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 for aggregates
  *
- * 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 for an aggregate
  *
  * This is useful for aggs to register shutdown callbacks, which will ensure
- * that non-memory resources are freed.
+ * that non-memory resources are freed.  The callback will occur just before
+ * the associated aggcontext (as returned by AggCheckCallContext) is reset,
+ * either between groups or as a result of rescanning the query.  The callback
+ * will NOT be called on error paths.  The typical use-case is for freeing of
+ * tuplestores or tuplesorts maintained in aggcontext, or pins held by slots
+ * created by the agg functions.  (The callback 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
@@ -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
@@ -23,7 +23,7 @@ typedef struct Node *fmNodePtr;
 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 +656,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

Original file line number	Diff line number	Diff line change
`@@ -2199,44 +2199,56 @@ AggGetAggref(FunctionCallInfo fcinfo)`
`2199`	`2199`	`}`
`2200`	`2200`
`2201`	`2201`	`/*`
`2202`		`- * AggGetPerTupleEContext - fetch per-input-tuple ExprContext`
	`2202`	`+ * AggGetTempMemoryContext - fetch short-term memory context for aggregates`
`2203`	`2203`	`*`
`2204`		`- * This is useful in agg final functions; the econtext returned is the`
`2205`		`- * same per-tuple context that the transfn was called in (which can`
`2206`		`- * safely get reset during the final function).`
	`2204`	`+ * This is useful in agg final functions; the context returned is one that`
	`2205`	`+ * the final function can safely reset as desired. This isn't useful for`
	`2206`	`+ * transition functions, since the context returned MAY (we don't promise)`
	`2207`	`+ * be the same as the context those are called in.`
`2207`	`2208`	`*`
`2208`	`2209`	`* As above, this is currently not useful for aggs called as window functions.`
`2209`	`2210`	`*/`
`2210`		`-ExprContext *`
`2211`		`-AggGetPerTupleEContext(FunctionCallInfo fcinfo)`
	`2211`	`+MemoryContext`
	`2212`	`+AggGetTempMemoryContext(FunctionCallInfo fcinfo)`
`2212`	`2213`	`{`
`2213`	`2214`	`if (fcinfo->context && IsA(fcinfo->context, AggState))`
`2214`	`2215`	`{`
`2215`	`2216`	`AggState aggstate = (AggState ) fcinfo->context;`
`2216`	`2217`
`2217`		`- return aggstate->tmpcontext;`
	`2218`	`+ return aggstate->tmpcontext->ecxt_per_tuple_memory;`
`2218`	`2219`	`}`
`2219`	`2220`	`return NULL;`
`2220`	`2221`	`}`
`2221`	`2222`
`2222`	`2223`	`/*`
`2223`		`- * AggGetPerAggEContext - fetch per-output-tuple ExprContext`
	`2224`	`+ * AggRegisterCallback - register a cleanup callback for an aggregate`
`2224`	`2225`	`*`
`2225`	`2226`	`* This is useful for aggs to register shutdown callbacks, which will ensure`
`2226`		`- * that non-memory resources are freed.`
	`2227`	`+ * that non-memory resources are freed. The callback will occur just before`
	`2228`	`+ * the associated aggcontext (as returned by AggCheckCallContext) is reset,`
	`2229`	`+ * either between groups or as a result of rescanning the query. The callback`
	`2230`	`+ * will NOT be called on error paths. The typical use-case is for freeing of`
	`2231`	`+ * tuplestores or tuplesorts maintained in aggcontext, or pins held by slots`
	`2232`	`+ * created by the agg functions. (The callback will not be called until after`
	`2233`	`+ * the result of the finalfn is no longer needed, so it's safe for the finalfn`
	`2234`	`+ * to return data that will be freed by the callback.)`
`2227`	`2235`	`*`
`2228`	`2236`	`* As above, this is currently not useful for aggs called as window functions.`
`2229`	`2237`	`*/`
`2230`		`-ExprContext *`
`2231`		`-AggGetPerAggEContext(FunctionCallInfo fcinfo)`
	`2238`	`+void`
	`2239`	`+AggRegisterCallback(FunctionCallInfo fcinfo,`
	`2240`	`+ ExprContextCallbackFunction func,`
	`2241`	`+ Datum arg)`
`2232`	`2242`	`{`
`2233`	`2243`	`if (fcinfo->context && IsA(fcinfo->context, AggState))`
`2234`	`2244`	`{`
`2235`	`2245`	`AggState aggstate = (AggState ) fcinfo->context;`
`2236`	`2246`
`2237`		`- return aggstate->ss.ps.ps_ExprContext;`
	`2247`	`+ RegisterExprContextCallback(aggstate->ss.ps.ps_ExprContext, func, arg);`
	`2248`	`+`
	`2249`	`+ return;`
`2238`	`2250`	`}`
`2239`		`- return NULL;`
	`2251`	`+ elog(ERROR, "aggregate function cannot register a callback in this context");`
`2240`	`2252`	`}`
`2241`	`2253`
`2242`	`2254`