Add JSON_OBJECTAGG() and JSON_ARRAYAGG() transformation

Author: Nikita Glukhov

src/backend/nodes/copyfuncs.c

@@ -2193,6 +2193,26 @@ _copyJsonCtorOpts(const JsonCtorOpts *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;
+}
+
 /*
  * _copyJsonArrayCtor
  */
@@ -2209,6 +2229,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;
+}
+
 /* ****************************************************************
  *						relation.h copy functions
  *
@@ -5073,9 +5112,15 @@ copyObjectImpl(const void *from)
 		case T_JsonObjectCtor:
 			retval = _copyJsonObjectCtor(from);
 			break;
+		case T_JsonObjectAgg:
+			retval = _copyJsonObjectAgg(from);
+			break;
 		case T_JsonArrayCtor:
 			retval = _copyJsonArrayCtor(from);
 			break;
+		case T_JsonArrayAgg:
+			retval = _copyJsonArrayAgg(from);
+			break;
 
 			/*
 			 * RELATION NODES

src/backend/nodes/nodeFuncs.c

@@ -3740,6 +3740,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"
@@ -124,6 +126,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,
@@ -380,6 +384,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));
@@ -3857,6 +3869,162 @@ transformJsonObjectCtor(ParseState *pstate, JsonObjectCtor *ctor)
 	return coerceJsonFuncExpr(pstate, (Node *) fexpr, &returning, true);
 }
 
+/*
+ * Common code for JSON_OBJECTAGG and JSON_ARRAYAGG transformation.
+ */
+static Node *
+transformJsonAggCtor(ParseState *pstate, JsonAggCtor *agg_ctor,
+					 JsonReturning *returning, List *args, Oid aggfnoid,
+					 Oid aggtype, FuncFormat format, JsonCtorOpts *formatopts)
+{
+	Node	   *node;
+	Expr	   *aggfilter = agg_ctor->agg_filter ? (Expr *)
+		transformWhereClause(pstate, agg_ctor->agg_filter,
+							 EXPR_KIND_FILTER, "FILTER") : NULL;
+
+	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->winformat = format;
+		wfunc->winformatopts = (Node *) formatopts;
+		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->aggformat = format;
+		aggref->aggformatopts = (Node *) formatopts;
+		aggref->location = agg_ctor->location;
+
+		transformAggregateCall(pstate, aggref, args, agg_ctor->agg_order, false);
+
+		node = (Node *) aggref;
+	}
+
+	return coerceJsonFuncExpr(pstate, node, 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;
+	Oid			aggfnoid;
+	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)
+	{
+		aggfnoid = F_JSONB_OBJECTAGG;
+		aggtype = JSONBOID;
+	}
+	else
+	{
+		aggfnoid = F_JSON_OBJECTAGG;
+		aggtype = JSONOID;
+	}
+
+	return transformJsonAggCtor(pstate, &agg->ctor, &returning, args, aggfnoid,
+								aggtype, FUNCFMT_JSON_OBJECTAGG,
+								makeJsonCtorOpts(&returning,
+												 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;
+	Oid			aggfnoid;
+	Oid			aggtype;
+
+	transformJsonOutput(pstate, agg->ctor.output, true, &returning);
+
+	arg = transformJsonValueExpr(pstate, agg->arg, JS_FORMAT_DEFAULT);
+
+	if (returning.format.type == JS_FORMAT_JSONB)
+	{
+		aggfnoid = agg->absent_on_null ? F_JSONB_AGG_STRICT : F_JSONB_AGG;
+		aggtype = JSONBOID;
+	}
+	else
+	{
+		aggfnoid = agg->absent_on_null ? F_JSON_AGG_STRICT : F_JSON_AGG;
+		aggtype = JSONOID;
+	}
+
+	return transformJsonAggCtor(pstate, &agg->ctor, &returning, list_make1(arg),
+								aggfnoid, aggtype, FUNCFMT_JSON_ARRAYAGG,
+								makeJsonCtorOpts(&returning,
+												 false, agg->absent_on_null));
+}
+
 /*
  * Transform JSON_ARRAY() constructor.
  *