Further support for JSON_ARRAYAGG() and JSON_OBJECTAGG()
- 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
Assessment
This issue has not been assessed yet.