open-dsi
/
postgres
mirror of https://github.com/postgres/postgres
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Remove fundamentally-redundant processing in jsonb_agg() et al.

Browse Source 
The various variants of jsonb_agg() operate as follows,
for each aggregate input value:

1. Build a JsonbValue tree representation of the input value.
2. Flatten the JsonbValue tree into a Jsonb in on-disk format.
3. Iterate through the Jsonb, building a JsonbValue that is part
of the aggregate's state stored in aggcontext, but is otherwise
identical to what phase 1 built.

This is very slightly less silly than it sounds, because phase 1
involves calling non-JSONB code such as datatype output functions,
which are likely to leak memory, and we don't want to leak into the
aggcontext.  Nonetheless, phases 2 and 3 are accomplishing exactly
nothing that is useful if we can make phase 1 put the JsonbValue
tree where we need it.  We could probably do that with a bunch of
MemoryContextSwitchTo's, but what seems more robust is to give
pushJsonbValue the responsibility of building the JsonbValue tree
in a specified non-current memory context.  The previous patch
created the infrastructure for that, and this patch simply makes
the aggregate functions use it and then rips out phases 2 and 3.

For me, this makes jsonb_agg() with a text column as input run
about 2X faster than before.  It's not yet on par with json_agg(),
but this removes a whole lot of the difference.

Author: Tom Lane <tgl@sss.pgh.pa.us>
Reviewed-by: jian he <jian.universality@gmail.com>
Reviewed-by: Chao Li <li.evan.chao@gmail.com>
Discussion: https://postgr.es/m/1060917.1753202222@sss.pgh.pa.us
pull/256/head
Tom Lane 1 week ago
parent 0986e95161
commit b61aa76e45
1 changed files with 45 additions and 258 deletions
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Show Stats Download Patch File Download Diff File
  1. 303
      src/backend/utils/adt/jsonb.c

303
src/backend/utils/adt/jsonb.c
Unescape View File

@ -27,7 +27,7 @@
typedef struct JsonbAggState
{
JsonbInState *res;
JsonbInState pstate;
JsonTypeCategory key_category;
Oid key_output_func;
JsonTypeCategory val_category;
@ -54,7 +54,6 @@ static void datum_to_jsonb_internal(Datum val, bool is_null, JsonbInState *resul
bool key_scalar);
static void add_jsonb(Datum val, bool is_null, JsonbInState *result,
Oid val_type, bool key_scalar);
static JsonbParseState *clone_parse_state(JsonbParseState *state);
static char *JsonbToCStringWorker(StringInfo out, JsonbContainer *in, int estimated_len, bool indent);
static void add_indent(StringInfo out, bool indent, int level);
@ -1454,54 +1453,16 @@ close_object:
/*
* shallow clone of a parse state, suitable for use in aggregate
* final functions that will only append to the values rather than
* change them.
* Functions for jsonb_agg, jsonb_object_agg, and variants
*/
static JsonbParseState *
clone_parse_state(JsonbParseState *state)
{
JsonbParseState *result,
*icursor,
*ocursor;
if (state == NULL)
return NULL;
result = palloc(sizeof(JsonbParseState));
icursor = state;
ocursor = result;
for (;;)
{
ocursor->contVal = icursor->contVal;
ocursor->size = icursor->size;
ocursor->unique_keys = icursor->unique_keys;
ocursor->skip_nulls = icursor->skip_nulls;
icursor = icursor->next;
if (icursor == NULL)
break;
ocursor->next = palloc(sizeof(JsonbParseState));
ocursor = ocursor->next;
}
ocursor->next = NULL;
return result;
}
static Datum
jsonb_agg_transfn_worker(FunctionCallInfo fcinfo, bool absent_on_null)
{
MemoryContext oldcontext,
aggcontext;
MemoryContext aggcontext;
JsonbAggState *state;
JsonbInState elem;
Datum val;
JsonbInState *result;
bool single_scalar = false;
JsonbIterator *it;
Jsonb *jbelem;
JsonbValue v;
JsonbIteratorToken type;
if (!AggCheckCallContext(fcinfo, &aggcontext))
{
@ -1520,12 +1481,10 @@ jsonb_agg_transfn_worker(FunctionCallInfo fcinfo, bool absent_on_null)
(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
errmsg("could not determine input data type")));
oldcontext = MemoryContextSwitchTo(aggcontext);
state = palloc(sizeof(JsonbAggState));
result = palloc0(sizeof(JsonbInState));
state->res = result;
state = MemoryContextAllocZero(aggcontext, sizeof(JsonbAggState));
result = &state->pstate;
result->outcontext = aggcontext;
pushJsonbValue(result, WJB_BEGIN_ARRAY, NULL);
MemoryContextSwitchTo(oldcontext);
json_categorize_type(arg_type, true, &state->val_category,
&state->val_output_func);
@ -1533,74 +1492,23 @@ jsonb_agg_transfn_worker(FunctionCallInfo fcinfo, bool absent_on_null)
else
{
state = (JsonbAggState *) PG_GETARG_POINTER(0);
result = state->res;
result = &state->pstate;
}
if (absent_on_null && PG_ARGISNULL(1))
PG_RETURN_POINTER(state);
/* turn the argument into jsonb in the normal function context */
/*
* We run this code in the normal function context, so that we don't leak
* any cruft from datatype output functions and such into the aggcontext.
* But the "result" JsonbValue will be constructed in aggcontext, so that
* it remains available across calls.
*/
val = PG_ARGISNULL(1) ? (Datum) 0 : PG_GETARG_DATUM(1);
memset(&elem, 0, sizeof(JsonbInState));
datum_to_jsonb_internal(val, PG_ARGISNULL(1), &elem, state->val_category,
datum_to_jsonb_internal(val, PG_ARGISNULL(1), result, state->val_category,
state->val_output_func, false);
jbelem = JsonbValueToJsonb(elem.result);
/* switch to the aggregate context for accumulation operations */
oldcontext = MemoryContextSwitchTo(aggcontext);
it = JsonbIteratorInit(&jbelem->root);
while ((type = JsonbIteratorNext(&it, &v, false)) != WJB_DONE)
{
switch (type)
{
case WJB_BEGIN_ARRAY:
if (v.val.array.rawScalar)
single_scalar = true;
else
pushJsonbValue(result, type, NULL);
break;
case WJB_END_ARRAY:
if (!single_scalar)
pushJsonbValue(result, type, NULL);
break;
case WJB_BEGIN_OBJECT:
case WJB_END_OBJECT:
pushJsonbValue(result, type, NULL);
break;
case WJB_ELEM:
case WJB_KEY:
case WJB_VALUE:
if (v.type == jbvString)
{
/* copy string values in the aggregate context */
char *buf = palloc(v.val.string.len + 1);
snprintf(buf, v.val.string.len + 1, "%s", v.val.string.val);
v.val.string.val = buf;
}
else if (v.type == jbvNumeric)
{
/* same for numeric */
v.val.numeric =
DatumGetNumeric(DirectFunctionCall1(numeric_uplus,
NumericGetDatum(v.val.numeric)));
}
pushJsonbValue(result, type, &v);
break;
default:
elog(ERROR, "unknown jsonb iterator token type");
}
}
MemoryContextSwitchTo(oldcontext);
PG_RETURN_POINTER(state);
}
@ -1638,17 +1546,18 @@ jsonb_agg_finalfn(PG_FUNCTION_ARGS)
arg = (JsonbAggState *) PG_GETARG_POINTER(0);
/*
* We need to do a shallow clone of the argument in case the final
* function is called more than once, so we avoid changing the argument. A
* shallow clone is sufficient as we aren't going to change any of the
* values, just add the final array end marker.
* The final function can be called more than once, so we must not change
* the stored JsonbValue data structure. Fortunately, the WJB_END_ARRAY
* action will only change fields in the JsonbInState struct itself, so we
* can simply invoke pushJsonbValue on a local copy of that.
*/
memset(&result, 0, sizeof(JsonbInState));
result.parseState = clone_parse_state(arg->res->parseState);
result = arg->pstate;
pushJsonbValue(&result, WJB_END_ARRAY, NULL);
/* We expect result.parseState == NULL after closing the array */
Assert(result.parseState == NULL);
out = JsonbValueToJsonb(result.result);
PG_RETURN_POINTER(out);
@ -1658,18 +1567,10 @@ static Datum
jsonb_object_agg_transfn_worker(FunctionCallInfo fcinfo,
bool absent_on_null, bool unique_keys)
{
MemoryContext oldcontext,
aggcontext;
JsonbInState elem;
MemoryContext aggcontext;
JsonbAggState *state;
Datum val;
JsonbInState *result;
bool single_scalar;
JsonbIterator *it;
Jsonb *jbkey,
*jbval;
JsonbValue v;
JsonbIteratorToken type;
bool skip;
if (!AggCheckCallContext(fcinfo, &aggcontext))
@ -1684,16 +1585,13 @@ jsonb_object_agg_transfn_worker(FunctionCallInfo fcinfo,
{
Oid arg_type;
oldcontext = MemoryContextSwitchTo(aggcontext);
state = palloc(sizeof(JsonbAggState));
result = palloc0(sizeof(JsonbInState));
state->res = result;
state = MemoryContextAllocZero(aggcontext, sizeof(JsonbAggState));
result = &state->pstate;
result->outcontext = aggcontext;
pushJsonbValue(result, WJB_BEGIN_OBJECT, NULL);
result->parseState->unique_keys = unique_keys;
result->parseState->skip_nulls = absent_on_null;
MemoryContextSwitchTo(oldcontext);
arg_type = get_fn_expr_argtype(fcinfo->flinfo, 1);
if (arg_type == InvalidOid)
@ -1717,11 +1615,9 @@ jsonb_object_agg_transfn_worker(FunctionCallInfo fcinfo,
else
{
state = (JsonbAggState *) PG_GETARG_POINTER(0);
result = state->res;
result = &state->pstate;
}
/* turn the argument into jsonb in the normal function context */
if (PG_ARGISNULL(1))
ereport(ERROR,
(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
@ -1736,133 +1632,22 @@ jsonb_object_agg_transfn_worker(FunctionCallInfo fcinfo,
if (skip && !unique_keys)
PG_RETURN_POINTER(state);
/*
* We run this code in the normal function context, so that we don't leak
* any cruft from datatype output functions and such into the aggcontext.
* But the "result" JsonbValue will be constructed in aggcontext, so that
* it remains available across calls.
*/
val = PG_GETARG_DATUM(1);
memset(&elem, 0, sizeof(JsonbInState));
datum_to_jsonb_internal(val, false, &elem, state->key_category,
datum_to_jsonb_internal(val, false, result, state->key_category,
state->key_output_func, true);
jbkey = JsonbValueToJsonb(elem.result);
val = PG_ARGISNULL(2) ? (Datum) 0 : PG_GETARG_DATUM(2);
memset(&elem, 0, sizeof(JsonbInState));
datum_to_jsonb_internal(val, PG_ARGISNULL(2), &elem, state->val_category,
datum_to_jsonb_internal(val, PG_ARGISNULL(2), result, state->val_category,
state->val_output_func, false);
jbval = JsonbValueToJsonb(elem.result);
it = JsonbIteratorInit(&jbkey->root);
/* switch to the aggregate context for accumulation operations */
oldcontext = MemoryContextSwitchTo(aggcontext);
/*
* keys should be scalar, and we should have already checked for that
* above when calling datum_to_jsonb, so we only need to look for these
* things.
*/
while ((type = JsonbIteratorNext(&it, &v, false)) != WJB_DONE)
{
switch (type)
{
case WJB_BEGIN_ARRAY:
if (!v.val.array.rawScalar)
elog(ERROR, "unexpected structure for key");
break;
case WJB_ELEM:
if (v.type == jbvString)
{
/* copy string values in the aggregate context */
char *buf = palloc(v.val.string.len + 1);
snprintf(buf, v.val.string.len + 1, "%s", v.val.string.val);
v.val.string.val = buf;
}
else
{
ereport(ERROR,
(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
errmsg("object keys must be strings")));
}
pushJsonbValue(result, WJB_KEY, &v);
if (skip)
{
v.type = jbvNull;
pushJsonbValue(result, WJB_VALUE, &v);
MemoryContextSwitchTo(oldcontext);
PG_RETURN_POINTER(state);
}
break;
case WJB_END_ARRAY:
break;
default:
elog(ERROR, "unexpected structure for key");
break;
}
}
it = JsonbIteratorInit(&jbval->root);
single_scalar = false;
/*
* values can be anything, including structured and null, so we treat them
* as in json_agg_transfn, except that single scalars are always pushed as
* WJB_VALUE items.
*/
while ((type = JsonbIteratorNext(&it, &v, false)) != WJB_DONE)
{
switch (type)
{
case WJB_BEGIN_ARRAY:
if (v.val.array.rawScalar)
single_scalar = true;
else
pushJsonbValue(result, type, NULL);
break;
case WJB_END_ARRAY:
if (!single_scalar)
pushJsonbValue(result, type, NULL);
break;
case WJB_BEGIN_OBJECT:
case WJB_END_OBJECT:
pushJsonbValue(result, type, NULL);
break;
case WJB_ELEM:
case WJB_KEY:
case WJB_VALUE:
if (v.type == jbvString)
{
/* copy string values in the aggregate context */
char *buf = palloc(v.val.string.len + 1);
snprintf(buf, v.val.string.len + 1, "%s", v.val.string.val);
v.val.string.val = buf;
}
else if (v.type == jbvNumeric)
{
/* same for numeric */
v.val.numeric =
DatumGetNumeric(DirectFunctionCall1(numeric_uplus,
NumericGetDatum(v.val.numeric)));
}
pushJsonbValue(result, single_scalar ? WJB_VALUE : type, &v);
break;
default:
elog(ERROR, "unknown jsonb iterator token type");
}
}
MemoryContextSwitchTo(oldcontext);
PG_RETURN_POINTER(state);
}
@ -1919,18 +1704,20 @@ jsonb_object_agg_finalfn(PG_FUNCTION_ARGS)
arg = (JsonbAggState *) PG_GETARG_POINTER(0);
/*
* We need to do a shallow clone of the argument's res field in case the
* final function is called more than once, so we avoid changing the
* aggregate state value. A shallow clone is sufficient as we aren't
* going to change any of the values, just add the final object end
* marker.
* The final function can be called more than once, so we must not change
* the stored JsonbValue data structure. Fortunately, the WJB_END_OBJECT
* action will only destructively change fields in the JsonbInState struct
* itself, so we can simply invoke pushJsonbValue on a local copy of that.
* (This technique results in running uniqueifyJsonbObject each time, but
* for now we won't bother trying to avoid that.)
*/
memset(&result, 0, sizeof(JsonbInState));
result.parseState = clone_parse_state(arg->res->parseState);
result = arg->pstate;
pushJsonbValue(&result, WJB_END_OBJECT, NULL);
/* We expect result.parseState == NULL after closing the object */
Assert(result.parseState == NULL);
out = JsonbValueToJsonb(result.result);
PG_RETURN_POINTER(out);

Write Preview
Loading…
Cancel
Save