Add JSON_OBJECTAGG() and JSON_ARRAYAGG() transformation

Author: Nikita Glukhov

src/backend/executor/execExpr.c

@@ -2138,6 +2138,12 @@ ExecInitExprRec(Expr *node, ExprState *state,
 				int			nargs = list_length(args);
 				int			argno = 0;
 
+				if (ctor->func)
+				{
+					ExecInitExprRec(ctor->func, state, resv, resnull);
+					break;
+				}
+
 				scratch.opcode = EEOP_JSON_CTOR;
 				scratch.d.json_ctor.ctor = ctor;
 				scratch.d.json_ctor.arg_values = palloc(sizeof(Datum) * nargs);

src/backend/nodes/copyfuncs.c

@@ -2344,6 +2344,26 @@ _copyJsonObjectCtor(const JsonObjectCtor *from)
 	return newnode;
 }
 
+/*
+ * _copyJsonObjectAgg
+ */
+static JsonObjectAgg *
+_copyJsonObjectAgg(const JsonObjectAgg *from)
+{
+	JsonObjectAgg *newnode = makeNode(JsonObjectAgg);
+
+	COPY_NODE_FIELD(ctor.output);
+	COPY_NODE_FIELD(ctor.agg_filter);
+	COPY_NODE_FIELD(ctor.agg_order);
+	COPY_NODE_FIELD(ctor.over);
+	COPY_LOCATION_FIELD(ctor.location);
+	COPY_NODE_FIELD(arg);
+	COPY_SCALAR_FIELD(absent_on_null);
+	COPY_SCALAR_FIELD(unique);
+
+	return newnode;
+}
+
 /*
  * _copyJsonOutput
  */
@@ -2374,6 +2394,25 @@ _copyJsonArrayCtor(const JsonArrayCtor *from)
 	return newnode;
 }
 
+/*
+ * _copyJsonArrayAgg
+ */
+static JsonArrayAgg *
+_copyJsonArrayAgg(const JsonArrayAgg *from)
+{
+	JsonArrayAgg *newnode = makeNode(JsonArrayAgg);
+
+	COPY_NODE_FIELD(ctor.output);
+	COPY_NODE_FIELD(ctor.agg_filter);
+	COPY_NODE_FIELD(ctor.agg_order);
+	COPY_NODE_FIELD(ctor.over);
+	COPY_LOCATION_FIELD(ctor.location);
+	COPY_NODE_FIELD(arg);
+	COPY_SCALAR_FIELD(absent_on_null);
+
+	return newnode;
+}
+
 /* ****************************************************************
  *						pathnodes.h copy functions
  *
@@ -5289,12 +5328,18 @@ copyObjectImpl(const void *from)
 		case T_JsonObjectCtor:
 			retval = _copyJsonObjectCtor(from);
 			break;
+		case T_JsonObjectAgg:
+			retval = _copyJsonObjectAgg(from);
+			break;
 		case T_JsonOutput:
 			retval = _copyJsonOutput(from);
 			break;
 		case T_JsonArrayCtor:
 			retval = _copyJsonArrayCtor(from);
 			break;
+		case T_JsonArrayAgg:
+			retval = _copyJsonArrayAgg(from);
+			break;
 
 			/*
 			 * RELATION NODES

src/backend/nodes/nodeFuncs.c

@@ -3254,7 +3254,7 @@ expression_tree_mutator(Node *node,
 
 				FLATCOPY(newnode, jve, JsonCtorExpr);
 				MUTATE(newnode->args, jve->args, List *);
-				MUTATE(newnode->func, jve->func, FuncExpr *);
+				MUTATE(newnode->func, jve->func, Expr *);
 				MUTATE(newnode->returning, jve->returning, JsonReturning *);
 
 				return (Node *) newnode;
@@ -4020,6 +4020,38 @@ raw_expression_tree_walker(Node *node,
 					return true;
 			}
 			break;
+		case T_JsonObjectAgg:
+			{
+				JsonObjectAgg *joa = (JsonObjectAgg *) node;
+
+				if (walker(joa->ctor.output, context))
+					return true;
+				if (walker(joa->ctor.agg_order, context))
+					return true;
+				if (walker(joa->ctor.agg_filter, context))
+					return true;
+				if (walker(joa->ctor.over, context))
+					return true;
+				if (walker(joa->arg, context))
+					return true;
+			}
+			break;
+		case T_JsonArrayAgg:
+			{
+				JsonArrayAgg *jaa = (JsonArrayAgg *) node;
+
+				if (walker(jaa->ctor.output, context))
+					return true;
+				if (walker(jaa->ctor.agg_order, context))
+					return true;
+				if (walker(jaa->ctor.agg_filter, context))
+					return true;
+				if (walker(jaa->ctor.over, context))
+					return true;
+				if (walker(jaa->arg, context))
+					return true;
+			}
+			break;
 		default:
 			elog(ERROR, "unrecognized node type: %d",
 				 (int) nodeTag(node));

src/backend/parser/parse_expr.c

@@ -15,6 +15,8 @@
 
 #include "postgres.h"
 
+#include "catalog/pg_aggregate.h"
+#include "catalog/pg_proc.h"
 #include "catalog/pg_type.h"
 #include "commands/dbcommands.h"
 #include "miscadmin.h"
@@ -123,6 +125,8 @@ static Node *transformTypeCast(ParseState *pstate, TypeCast *tc);
 static Node *transformCollateClause(ParseState *pstate, CollateClause *c);
 static Node *transformJsonObjectCtor(ParseState *pstate, JsonObjectCtor *ctor);
 static Node *transformJsonArrayCtor(ParseState *pstate, JsonArrayCtor *ctor);
+static Node *transformJsonObjectAgg(ParseState *pstate, JsonObjectAgg *agg);
+static Node *transformJsonArrayAgg(ParseState *pstate, JsonArrayAgg *agg);
 static Node *make_row_comparison_op(ParseState *pstate, List *opname,
 									List *largs, List *rargs, int location);
 static Node *make_row_distinct_op(ParseState *pstate, List *opname,
@@ -379,6 +383,14 @@ transformExprRecurse(ParseState *pstate, Node *expr)
 			result = transformJsonArrayCtor(pstate, (JsonArrayCtor *) expr);
 			break;
 
+		case T_JsonObjectAgg:
+			result = transformJsonObjectAgg(pstate, (JsonObjectAgg *) expr);
+			break;
+
+		case T_JsonArrayAgg:
+			result = transformJsonArrayAgg(pstate, (JsonArrayAgg *) expr);
+			break;
+
 		default:
 			/* should not reach here */
 			elog(ERROR, "unrecognized node type: %d", (int) nodeTag(expr));
@@ -3991,7 +4003,6 @@ transformJsonObjectCtor(ParseState *pstate, JsonObjectCtor *ctor)
 
 	jsctor = makeNode(JsonCtorExpr);
 	jsctor->args = args;
-	jsctor->func = NULL;
 	jsctor->type = JSCTOR_JSON_OBJECT;
 	jsctor->returning = transformJsonCtorOutput(pstate, ctor->output, args);
 	jsctor->unique = ctor->unique;
@@ -4001,6 +4012,171 @@ transformJsonObjectCtor(ParseState *pstate, JsonObjectCtor *ctor)
 	return coerceJsonFuncExpr(pstate, (Node *) jsctor, jsctor->returning, true);
 }
 
+/*
+ * Common code for JSON_OBJECTAGG and JSON_ARRAYAGG transformation.
+ */
+static Node *
+transformJsonAggCtor(ParseState *pstate, JsonAggCtor *agg_ctor,
+					 JsonReturning *returning, List *args, const char *aggfn,
+					 Oid aggtype, JsonCtorType ctor_type,
+					 bool unique, bool absent_on_null)
+{
+	Oid			aggfnoid;
+	JsonCtorExpr *jsctor;
+	Node	   *node;
+	Expr	   *aggfilter = agg_ctor->agg_filter ? (Expr *)
+		transformWhereClause(pstate, agg_ctor->agg_filter,
+							 EXPR_KIND_FILTER, "FILTER") : NULL;
+
+	aggfnoid = DatumGetInt32(DirectFunctionCall1(regprocin,
+												 CStringGetDatum(aggfn)));
+
+	if (agg_ctor->over)
+	{
+		/* window function */
+		WindowFunc *wfunc = makeNode(WindowFunc);
+
+		wfunc->winfnoid = aggfnoid;
+		wfunc->wintype = aggtype;
+		/* wincollid and inputcollid will be set by parse_collate.c */
+		wfunc->args = args;
+		/* winref will be set by transformWindowFuncCall */
+		wfunc->winstar = false;
+		wfunc->winagg = true;
+		wfunc->aggfilter = aggfilter;
+		wfunc->location = agg_ctor->location;
+
+		/*
+		 * ordered aggs not allowed in windows yet
+		 */
+		if (agg_ctor->agg_order != NIL)
+			ereport(ERROR,
+					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+					 errmsg("aggregate ORDER BY is not implemented for window functions"),
+					 parser_errposition(pstate, agg_ctor->location)));
+
+		/* parse_agg.c does additional window-func-specific processing */
+		transformWindowFuncCall(pstate, wfunc, agg_ctor->over);
+
+		node = (Node *) wfunc;
+	}
+	else
+	{
+		Aggref	   *aggref = makeNode(Aggref);
+
+		aggref->aggfnoid = aggfnoid;
+		aggref->aggtype = aggtype;
+
+		/* aggcollid and inputcollid will be set by parse_collate.c */
+		aggref->aggtranstype = InvalidOid;		/* will be set by planner */
+		/* aggargtypes will be set by transformAggregateCall */
+		/* aggdirectargs and args will be set by transformAggregateCall */
+		/* aggorder and aggdistinct will be set by transformAggregateCall */
+		aggref->aggfilter = aggfilter;
+		aggref->aggstar = false;
+		aggref->aggvariadic = false;
+		aggref->aggkind = AGGKIND_NORMAL;
+		/* agglevelsup will be set by transformAggregateCall */
+		aggref->aggsplit = AGGSPLIT_SIMPLE;		/* planner might change this */
+		aggref->location = agg_ctor->location;
+
+		transformAggregateCall(pstate, aggref, args, agg_ctor->agg_order, false);
+
+		node = (Node *) aggref;
+	}
+
+	jsctor = makeNode(JsonCtorExpr);
+	jsctor->func = (Expr *) node;
+	jsctor->type = ctor_type;
+	jsctor->returning = *returning;
+	jsctor->unique = unique;
+	jsctor->absent_on_null = absent_on_null;
+	jsctor->location = agg_ctor->location;
+
+	return coerceJsonFuncExpr(pstate, (Node *) jsctor, returning, true);
+}
+
+/*
+ * Transform JSON_OBJECTAGG() aggregate function.
+ *
+ * JSON_OBJECTAGG() is transformed into
+ * json[b]_objectagg(key, value, absent_on_null, check_unique) call depending on
+ * the output JSON format.  Then the function call result is coerced to the
+ * target output type.
+ */
+static Node *
+transformJsonObjectAgg(ParseState *pstate, JsonObjectAgg *agg)
+{
+	JsonReturning returning;
+	Node	   *key;
+	Node	   *val;
+	List	   *args;
+	const char *aggfnname;
+	Oid			aggtype;
+
+	transformJsonOutput(pstate, agg->ctor.output, true, &returning);
+
+	key = transformExprRecurse(pstate, (Node *) agg->arg->key);
+	val = transformJsonValueExpr(pstate, agg->arg->value, JS_FORMAT_DEFAULT);
+
+	args = list_make4(key,
+					  val,
+					  makeBoolConst(agg->absent_on_null, false),
+					  makeBoolConst(agg->unique, false));
+
+	if (returning.format.type == JS_FORMAT_JSONB)
+	{
+		aggfnname = "pg_catalog.jsonb_objectagg"; /* F_JSONB_OBJECTAGG */
+		aggtype = JSONBOID;
+	}
+	else
+	{
+		aggfnname = "pg_catalog.json_objectagg"; /* F_JSON_OBJECTAGG; */
+		aggtype = JSONOID;
+	}
+
+	return transformJsonAggCtor(pstate, &agg->ctor, &returning, args, aggfnname,
+								aggtype, JSCTOR_JSON_OBJECTAGG,
+								agg->unique, agg->absent_on_null);
+}
+
+/*
+ * Transform JSON_ARRAYAGG() aggregate function.
+ *
+ * JSON_ARRAYAGG() is transformed into json[b]_agg[_strict]() call depending
+ * on the output JSON format and absent_on_null.  Then the function call result
+ * is coerced to the target output type.
+ */
+static Node *
+transformJsonArrayAgg(ParseState *pstate, JsonArrayAgg *agg)
+{
+	JsonReturning returning;
+	Node	   *arg;
+	const char *aggfnname;
+	Oid			aggtype;
+
+	transformJsonOutput(pstate, agg->ctor.output, true, &returning);
+
+	arg = transformJsonValueExpr(pstate, agg->arg, JS_FORMAT_DEFAULT);
+
+	if (returning.format.type == JS_FORMAT_JSONB)
+	{
+		aggfnname = agg->absent_on_null ?
+			"pg_catalog.jsonb_agg_strict" : "pg_catalog.jsonb_agg";
+		aggtype = JSONBOID;
+	}
+	else
+	{
+		aggfnname = agg->absent_on_null ?
+			"pg_catalog.json_agg_strict" : "pg_catalog.json_agg";
+		aggtype = JSONOID;
+	}
+
+	return transformJsonAggCtor(pstate, &agg->ctor, &returning, list_make1(arg),
+								aggfnname, aggtype, JSCTOR_JSON_ARRAYAGG,
+								false, agg->absent_on_null);
+}
+
 /*
  * Transform JSON_ARRAY() constructor.
  *

src/backend/parser/parse_target.c

@@ -1937,6 +1937,12 @@ FigureColnameInternal(Node *node, char **name)
 		case T_JsonArrayCtor:
 			*name = "json_array";
 			return 2;
+		case T_JsonObjectAgg:
+			*name = "json_objectagg";
+			return 2;
+		case T_JsonArrayAgg:
+			*name = "json_arrayagg";
+			return 2;
 		default:
 			break;
 	}