Add support for json_agg() pushdown

This patch adds a collection function for json_agg() aggregate. Also use a
specific json_agg_state type for the internal agg state so that corresponding
in/out functions can be specified for transition values to be passed around
from one node to another

diff --git a/src/backend/utils/adt/json.c b/src/backend/utils/adt/json.c
index af97fc1eff4cdbd5ac0989331c0d0601805d9f73..bf2bb725450757f104f0fd5ce132d4c66fbc0167 100644 (file)
--- a/src/backend/utils/adt/json.c
+++ b/src/backend/utils/adt/json.c
@@ -276,6 +276,45 @@ json_recv(PG_FUNCTION_ARGS)
        PG_RETURN_TEXT_P(cstring_to_text_with_len(str, nbytes));
 }
 
+#ifdef XCP
+Datum
+json_agg_state_in(PG_FUNCTION_ARGS)
+{
+       char       *str = pstrdup(PG_GETARG_CSTRING(0));
+       JsonAggState *state;
+       char *token, *freestr;
+
+       state = (JsonAggState *) palloc0(sizeof (JsonAggState));
+       state->str = makeStringInfo();
+
+       freestr = str;
+
+       token = strsep(&str, ":");
+       state->val_category = atoi(token);
+       appendStringInfoString(state->str, str);
+
+       pfree(freestr);
+
+       PG_RETURN_POINTER(state);
+}
+
+/*
+ * json_agg_collectfn only needs the 'val_category' for formatting purposes. So
+ * only output that along with the json string
+ */
+Datum
+json_agg_state_out(PG_FUNCTION_ARGS)
+{
+       JsonAggState *state = (JsonAggState *) PG_GETARG_POINTER(0);
+       char *result;
+       int len = 15 + strlen(state->str->data);
+
+       result = (char *) palloc0(len);
+       sprintf(result, "%d:%s", state->val_category, state->str->data);
+
+       PG_RETURN_CSTRING(result);
+}
+#endif
 /*
  * makeJsonLexContext
  *
@@ -1888,7 +1927,13 @@ json_agg_transfn(PG_FUNCTION_ARGS)
                state->str = makeStringInfo();
                MemoryContextSwitchTo(oldcontext);
 
+#ifndef XCP
+               /* 
+                * Do not add the start marker. It will be done by the
+                * json_agg_collectfn on receiving the first result
+                */
                appendStringInfoChar(state->str, '[');
+#endif
                json_categorize_type(arg_type, &state->val_category,
                                                         &state->val_output_func);
        }
@@ -1927,6 +1972,70 @@ json_agg_transfn(PG_FUNCTION_ARGS)
        PG_RETURN_POINTER(state);
 }
 
+#ifdef XCP
+/*
+ * json_agg collection function
+ */
+Datum
+json_agg_collectfn(PG_FUNCTION_ARGS)
+{
+       MemoryContext aggcontext,
+                               oldcontext;
+       JsonAggState *collectstate;
+       JsonAggState *transstate;
+
+       if (!AggCheckCallContext(fcinfo, &aggcontext))
+       {
+               /* cannot be called directly because of internal-type argument */
+               elog(ERROR, "json_agg_collectfn called in non-aggregate context");
+       }
+
+
+       /* cannot be called directly because of internal-type argument */
+       Assert(AggCheckCallContext(fcinfo, NULL));
+
+       if (PG_ARGISNULL(0))
+       {
+               /*
+                * Make this state object in a context where it will persist for the
+                * duration of the aggregate call.  MemoryContextSwitchTo is only
+                * needed the first time, as the StringInfo routines make sure they
+                * use the right context to enlarge the object if necessary.
+                */
+               oldcontext = MemoryContextSwitchTo(aggcontext);
+               collectstate = (JsonAggState *) palloc(sizeof(JsonAggState));
+               collectstate->str = makeStringInfo();
+               MemoryContextSwitchTo(oldcontext);
+
+               appendStringInfoChar(collectstate->str, '[');
+       }
+       else
+       {
+               collectstate = (JsonAggState *) PG_GETARG_POINTER(0);
+               if (!PG_ARGISNULL(1))
+                       appendStringInfoString(collectstate->str, ", ");
+       }
+
+       /* fast path for NULLs */
+       if (PG_ARGISNULL(1))
+               PG_RETURN_POINTER(collectstate);
+
+       transstate = (JsonAggState *) PG_GETARG_POINTER(1);
+
+       /* add some whitespace if structured type and not first item */
+       if (!PG_ARGISNULL(0) &&
+               (transstate->val_category == JSONTYPE_ARRAY ||
+                transstate->val_category == JSONTYPE_COMPOSITE))
+       {
+               appendStringInfoString(collectstate->str, "\n ");
+       }
+
+       appendStringInfoString(collectstate->str, transstate->str->data);
+
+       PG_RETURN_POINTER(collectstate);
+}
+#endif
+
 /*
  * json_agg final function
  */

diff --git a/src/include/catalog/pg_aggregate.h b/src/include/catalog/pg_aggregate.h
index ebbd118d883c77648227f8a212c49bad622d8dcc..6d8ade4156e43ea28778cf799f14cc70a67920e2 100644 (file)
--- a/src/include/catalog/pg_aggregate.h
+++ b/src/include/catalog/pg_aggregate.h
@@ -306,7 +306,7 @@ DATA(insert ( 3538  n 0 string_agg_transfn  -       string_agg_finalfn      -                               -                               -                               f
 DATA(insert ( 3545     n 0 bytea_string_agg_transfn    -       bytea_string_agg_finalfn        -                               -                               -               f f 0   2281    0       0       0               0       _null_  _null_ _null_ ));
 
 /* json */
-DATA(insert ( 3175     n 0 json_agg_transfn    -       json_agg_finalfn                        -                               -                               -                               f f 0   2281    0       0       0               0       _null_ _null_ _null_ ));
+DATA(insert ( 3175     n 0 json_agg_transfn    json_agg_collectfn json_agg_finalfn                     -                               -                               - f f 0 7028    7028    0       0       0       _null_ _null_ _null_ ));
 DATA(insert ( 3197     n 0 json_object_agg_transfn -   json_object_agg_finalfn -                               -                               -                               f f 0   2281    0       0       0               0       _null_ _null_ _null_ ));
 
 /* jsonb */

diff --git a/src/include/catalog/pg_proc.h b/src/include/catalog/pg_proc.h
index ede11ebc18abde8df447318b83282e7ee6f189c8..14214df39f67a21817e2b6044693c5f1ef3a4bdb 100644 (file)
--- a/src/include/catalog/pg_proc.h
+++ b/src/include/catalog/pg_proc.h
@@ -4403,9 +4403,11 @@ DATA(insert OID = 3155 (  row_to_json       PGNSP PGUID 12 1 0 0 0 f f f f t f s 1
 DESCR("map row to json");
 DATA(insert OID = 3156 (  row_to_json     PGNSP PGUID 12 1 0 0 0 f f f f t f s 2 0 114 "2249 16" _null_ _null_ _null_ _null_ _null_ row_to_json_pretty _null_ _null_ _null_ ));
 DESCR("map row to json with optional pretty printing");
-DATA(insert OID = 3173 (  json_agg_transfn      PGNSP PGUID 12 1 0 0 0 f f f f f f s 2 0 2281 "2281 2283" _null_ _null_ _null_ _null_ _null_ json_agg_transfn _null_ _null_ _null_ ));
+DATA(insert OID = 3173 (  json_agg_transfn      PGNSP PGUID 12 1 0 0 0 f f f f f f s 2 0 7028 "7028 2283" _null_ _null_ _null_ _null_ _null_ json_agg_transfn _null_ _null_ _null_ ));
 DESCR("json aggregate transition function");
-DATA(insert OID = 3174 (  json_agg_finalfn      PGNSP PGUID 12 1 0 0 0 f f f f f f i 1 0 114 "2281" _null_ _null_ _null_ _null_ _null_ json_agg_finalfn _null_ _null_ _null_ ));
+DATA(insert OID = 7029 (  json_agg_collectfn    PGNSP PGUID 12 1 0 0 0 f f f f f f s 2 0 7028 "7028 7028" _null_ _null_ _null_ _null_ _null_ json_agg_collectfn _null_ _null_ _null_ ));
+DESCR("json aggregate collection function");
+DATA(insert OID = 3174 (  json_agg_finalfn      PGNSP PGUID 12 1 0 0 0 f f f f f f i 1 0 114 "7028" _null_ _null_ _null_ _null_ _null_ json_agg_finalfn _null_ _null_ _null_ ));
 DESCR("json aggregate final function");
 DATA(insert OID = 3175 (  json_agg                PGNSP PGUID 12 1 0 0 0 t f f f f f s 1 0 114 "2283" _null_ _null_ _null_ _null_ _null_ aggregate_dummy _null_ _null_ _null_ ));
 DESCR("aggregate input into json");
@@ -5318,6 +5320,10 @@ DATA(insert OID = 7022 (  numeric_poly_agg_state_recv               PGNSP PGUID 12 1 0 0 0
 DESCR("I/O");
 DATA(insert OID = 7023 (  numeric_poly_agg_state_send             PGNSP PGUID 12 1 0 0 0 f f f f t f i 1 0 17 "7019" _null_ _null_ _null_ _null_ _null_ numeric_poly_agg_state_send _null_ _null_ _null_ ));
 DESCR("I/O");
+DATA(insert OID = 7030 ( json_agg_state_in                             PGNSP PGUID 12 1 0 0 0 f f f f t f i 1 0 7028 "2275" _null_ _null_ _null_ _null_ _null_ json_agg_state_in _null_ _null_ _null_ ));
+DESCR("I/O");
+DATA(insert OID = 7025 ( json_agg_state_out                    PGNSP PGUID 12 1 0 0 0 f f f f t f i 1 0 2275 "7028" _null_ _null_ _null_ _null_ _null_ json_agg_state_out _null_ _null_ _null_ ));
+DESCR("I/O");
 #endif
 /* pg_upgrade support */
 DATA(insert OID = 3582 ( binary_upgrade_set_next_pg_type_oid PGNSP PGUID  12 1 0 0 0 f f f f t f v 1 0 2278 "26" _null_ _null_ _null_ _null_ _null_ binary_upgrade_set_next_pg_type_oid _null_ _null_ _null_ ));

diff --git a/src/include/catalog/pg_type.h b/src/include/catalog/pg_type.h
index 571b5039073dcc80d95f13654bc10104cff58bd0..855c98c64db99b5fe68db9bebf6731ef94b52888 100644 (file)
--- a/src/include/catalog/pg_type.h
+++ b/src/include/catalog/pg_type.h
@@ -707,6 +707,10 @@ DATA(insert OID = 7019 ( numeric_poly_agg_state       PGNSP PGUID SIZEOF_POINTER t
 DESCR("numeric_poly_agg_state - internal type used for numeric/int8 aggregation");
 #define NUMERIC_POLY_AGG_STATE_OID             7019
 
+DATA(insert OID = 7028 ( json_agg_state           PGNSP PGUID SIZEOF_POINTER t p C f t \054 0 0 0 json_agg_state_in json_agg_state_out - - - - - ALIGNOF_POINTER p f 0 -1 0 0 _null_ _null_ _null_ ));
+DESCR("json_agg_state - internal type used for json aggregation");
+#define JSON_AGG_STATE_OID             7028
+
 /*
  * macros
  */

diff --git a/src/include/utils/json.h b/src/include/utils/json.h
index 6f69403f3dfa43dc543d14f04ff8ef7d09258872..11960bbf2c6bc8351d76f048d2005a9b7c76dad3 100644 (file)
--- a/src/include/utils/json.h
+++ b/src/include/utils/json.h
@@ -22,6 +22,10 @@ extern Datum json_in(PG_FUNCTION_ARGS);
 extern Datum json_out(PG_FUNCTION_ARGS);
 extern Datum json_recv(PG_FUNCTION_ARGS);
 extern Datum json_send(PG_FUNCTION_ARGS);
+#ifdef XCP
+extern Datum json_agg_state_in(PG_FUNCTION_ARGS);
+extern Datum json_agg_state_out(PG_FUNCTION_ARGS);
+#endif
 extern Datum array_to_json(PG_FUNCTION_ARGS);
 extern Datum array_to_json_pretty(PG_FUNCTION_ARGS);
 extern Datum row_to_json(PG_FUNCTION_ARGS);