citusdata / citusdata/citus

Further support for JSON_ARRAYAGG() and JSON_OBJECTAGG()

Open
#7,273 0 comments 0 reactions 0 assignees View on GitHub
enhancement pg16_support
Dominant language
C
Stars
12.8k
Forks
794
Avg merge
2d 14h
Merged PRs (30d)
31

Description

Setup:
```SQL
CREATE TABLE agg_test(country text, city text);
SELECT create_distributed_table('agg_test', 'country', shard_count := 2);
INSERT INTO agg_test VALUES ('Albania', 'Tirana'), ('Albania', 'Shkodra'), ('Albania', 'Elbasan');
INSERT INTO agg_test VALUES ('Turkey', 'Ankara'), ('Turkey', 'Istanbul');
SET client_min_messages TO debug4;
SET citus.log_remote_commands TO on;
```

1. Currently, for PG16's new SQL/JSON standard conforming constructors for JSON types, namely `json_arrayagg()` and `json_objectagg()`, we push down the execution to the workers when the aggregate is grouped by the table's distribution column:
```SQL
SELECT JSON_ARRAYAGG(city) AS cities_aggregated FROM agg_test GROUP BY country;
DEBUG: combine query: SELECT cities_aggregated FROM pg_catalog.citus_extradata_container(10, NULL::cstring(0), NULL::cstring(0), '(i 1)'::cstring(0)) remote_scan(cities_aggregated json, worker_column_2 text)
NOTICE: issuing SELECT JSON_ARRAYAGG(city RETURNING json) AS cities_aggregated, country AS worker_column_2 FROM public.agg_test_102008 agg_test WHERE true GROUP BY country
DETAIL: on server naisila@localhost:9701 connectionId: 3
NOTICE: issuing SELECT JSON_ARRAYAGG(city RETURNING json) AS cities_aggregated, country AS worker_column_2 FROM public.agg_test_102009 agg_test WHERE true GROUP BY country
DETAIL: on server naisila@localhost:9702 connectionId: 4
```

2. If not grouped on the distribution column, it is currently pulling all rows from the workers and performing the aggregation on the coordinator node:
```sql
SELECT JSON_ARRAYAGG(city) AS cities_aggregated FROM agg_test;
DEBUG: combine query: SELECT JSON_ARRAYAGG(cities_aggregated RETURNING json) AS cities_aggregated FROM pg_catalog.citus_extradata_container(10, NULL::cstring(0), NULL::cstring(0), '(i 1)'::cstring(0)) remote_scan(cities_aggregated text)
NOTICE: issuing SELECT city AS cities_aggregated FROM public.agg_test_102008 agg_test WHERE true
DETAIL: on server naisila@localhost:9701 connectionId: 5
NOTICE: issuing SELECT city AS cities_aggregated FROM public.agg_test_102009 agg_test WHERE true
DETAIL: on server naisila@localhost:9702 connectionId: 6
```

**We can improve the behavior in (2) by adding `json_arrayagg()` and `json_objectagg()` to our list of [special-case aggregates](https://docs.citusdata.com/en/v12.1/develop/reference_sql.html#aggregate-functions). See also `AggregateType` struct in our codebase.**

It could be a mimic of how the classic `json_agg` and `json_object_agg` work, utilizing the Citus-defined function `json_cat_agg`:
```SQL
SELECT json_agg(city) AS cities_aggregated FROM agg_test;
DEBUG: combine query: SELECT pg_catalog.json_cat_agg(cities_aggregated) AS cities_aggregated FROM pg_catalog.citus_extradata_container(10, NULL::cstring(0), NULL::cstring(0), '(i 1)'::cstring(0)) remote_scan(cities_aggregated json)
NOTICE: issuing SELECT json_agg(city) AS cities_aggregated FROM public.agg_test_102008 agg_test WHERE true
DETAIL: on server naisila@localhost:9701 connectionId: 5
NOTICE: issuing SELECT json_agg(city) AS cities_aggregated FROM public.agg_test_102009 agg_test WHERE true
DETAIL: on server naisila@localhost:9702 connectionId: 6
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.