Don't cache per-group context across the whole query in orderedsetaggs.c.

Although nodeAgg.c currently uses the same per-group memory context for
all groups of a query, that might change in future.  Avoid assuming it.
This costs us an extra AggCheckCallContext() call per group, but that's
pretty cheap and is probably good from a safety standpoint anyway.

Back-patch to 9.4 in case any third-party code copies this logic.

Andrew Gierth

diff --git a/src/backend/utils/adt/orderedsetaggs.c b/src/backend/utils/adt/orderedsetaggs.c
index d116a630355ccb06e6b9a6bdd54b06d95c38b6b3..135a14b02288489238e16e52335773c7bf85edf3 100644 (file)
--- 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;
+       /* Memory context containing per-group data: */
+       MemoryContext gcontext;
        /* Sort object we're accumulating data in: */
        Tuplesortstate *sortstate;
        /* Number of normal rows inserted into sortstate: */
@@ -103,8 +103,17 @@ ordered_set_startup(FunctionCallInfo fcinfo, bool use_tuples)
 {
        OSAPerGroupState *osastate;
        OSAPerQueryState *qstate;
+       MemoryContext gcontext;
        MemoryContext oldcontext;
 
+       /*
+        * Check we're called as aggregate (and not a window function), and get
+        * the Agg node's group-lifespan context (which might change from group to
+        * group, so we shouldn't cache it in the per-query state).
+        */
+       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,
         * create it, and do the per-query setup we need.
@@ -114,17 +123,10 @@ 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 */
+               /* Get the Aggref so we can examine aggregate's arguments */
                aggref = AggGetAggref(fcinfo);
                if (!aggref)
                        elog(ERROR, "ordered-set aggregate called in non-aggregate context");
@@ -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)
@@ -282,7 +284,7 @@ ordered_set_startup(FunctionCallInfo fcinfo, bool use_tuples)
 
        osastate->number_of_rows = 0;
 
-       /* Now register a shutdown callback to clean things up */
+       /* Now register a shutdown callback to clean things up at end of group */
        AggRegisterCallback(fcinfo,
                                                ordered_set_shutdown,
                                                PointerGetDatum(osastate));