diff --git a/server/analyzer/init.go b/server/analyzer/init.go index e0366af3a8..2b833f8ef5 100644 --- a/server/analyzer/init.go +++ b/server/analyzer/init.go @@ -148,6 +148,7 @@ var postgresOnlyAggregateFuncNames = map[string]bool{ "array_agg": true, "bool_and": true, "bool_or": true, + "json_agg": true, } // postgresOnlyWindowFuncNames holds Postgres functions that may only be used as window functions (i.e. diff --git a/server/functions/aggregate/init.go b/server/functions/aggregate/init.go index fe37348a9d..5fe99e4137 100755 --- a/server/functions/aggregate/init.go +++ b/server/functions/aggregate/init.go @@ -16,6 +16,7 @@ package aggregate func Init() { initBoolAggs() + initJsonAggs() initNumericAggs() initAvgAggs() initVarianceAggs() diff --git a/server/functions/aggregate/json_aggregates.go b/server/functions/aggregate/json_aggregates.go new file mode 100644 index 0000000000..278b88b03d --- /dev/null +++ b/server/functions/aggregate/json_aggregates.go @@ -0,0 +1,147 @@ +// Copyright 2026 Dolthub, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package aggregate + +import ( + "strings" + + "github.com/cockroachdb/errors" + "github.com/dolthub/go-mysql-server/sql" + + "github.com/dolthub/doltgresql/server/functions" + "github.com/dolthub/doltgresql/server/functions/framework" + pgtypes "github.com/dolthub/doltgresql/server/types" +) + +// initJsonAggs registers the JSON aggregate functions to the catalog. +func initJsonAggs() { + framework.RegisterAggregateFunction(jsonAgg) +} + +// jsonAgg represents PostgreSQL's json_agg(anyelement) aggregate. +var jsonAgg = framework.Func1Aggregate{ + Function1: framework.Function1{ + Name: "json_agg", + Return: pgtypes.Json, + Parameters: [1]*pgtypes.DoltgresType{pgtypes.AnyElement}, + // json_agg is deliberately not strict: an input SQL NULL contributes a + // JSON null element, while no input rows produce SQL NULL. + Strict: false, + Callable: func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val any) (any, error) { + return nil, nil + }, + }, + NewAggBuffer: newJsonAggBuffer, + NewAggWindowFunc: newJsonAggWindowFunction, +} + +// jsonAggBuffer accumulates the JSON representation of each input row. +type jsonAggBuffer struct { + expr sql.Expression + elemType *pgtypes.DoltgresType + elements []string +} + +var _ sql.AggregationBuffer = (*jsonAggBuffer)(nil) + +// newJsonAggBuffer creates an aggregation buffer for json_agg. +func newJsonAggBuffer(exprs []sql.Expression) (sql.AggregationBuffer, error) { + return &jsonAggBuffer{expr: exprs[0]}, nil +} + +// Dispose implements sql.AggregationBuffer. +func (b *jsonAggBuffer) Dispose(ctx *sql.Context) {} + +// Eval implements sql.AggregationBuffer. +func (b *jsonAggBuffer) Eval(ctx *sql.Context) (interface{}, error) { + if len(b.elements) == 0 { + return nil, nil + } + return joinJsonAggregateElements(b.elemType, b.elements), nil +} + +// Update implements sql.AggregationBuffer. +func (b *jsonAggBuffer) Update(ctx *sql.Context, row sql.Row) error { + value, include, err := framework.EvalAggregateArgument(ctx, b.expr, row) + if err != nil { + return err + } + if !include { + return nil + } + if b.elemType == nil { + var ok bool + b.elemType, ok = b.expr.Type(ctx).(*pgtypes.DoltgresType) + if !ok { + return errors.Errorf("json_agg: expected PostgreSQL argument type, got %T", b.expr.Type(ctx)) + } + } + raw, err := functions.ValueToJsonRaw(ctx, b.elemType, value) + if err != nil { + return err + } + b.elements = append(b.elements, string(raw)) + return nil +} + +// jsonAggWindowFunction computes json_agg over a window frame. +type jsonAggWindowFunction struct { + framework.WindowFramerState + expr sql.Expression +} + +var _ sql.WindowFunction = (*jsonAggWindowFunction)(nil) + +// newJsonAggWindowFunction creates a window-function implementation of json_agg. +func newJsonAggWindowFunction(exprs []sql.Expression, window *sql.WindowDefinition) (sql.WindowFunction, error) { + wf := &jsonAggWindowFunction{expr: exprs[0]} + if err := wf.BindFramer(window); err != nil { + return nil, err + } + return wf, nil +} + +// Compute implements sql.WindowFunction. +func (w *jsonAggWindowFunction) Compute(ctx *sql.Context, interval sql.WindowInterval, buffer sql.WindowBuffer) (interface{}, error) { + if interval.End <= interval.Start { + return nil, nil + } + elements := make([]string, 0, interval.End-interval.Start) + elemType, ok := w.expr.Type(ctx).(*pgtypes.DoltgresType) + if !ok { + return nil, errors.Errorf("json_agg: expected PostgreSQL argument type, got %T", w.expr.Type(ctx)) + } + for i := interval.Start; i < interval.End; i++ { + value, err := w.expr.Eval(ctx, buffer[i]) + if err != nil { + return nil, err + } + raw, err := functions.ValueToJsonRaw(ctx, elemType, value) + if err != nil { + return nil, err + } + elements = append(elements, string(raw)) + } + return joinJsonAggregateElements(elemType, elements), nil +} + +// joinJsonAggregateElements formats collected JSON values using PostgreSQL's aggregate layout. +func joinJsonAggregateElements(elemType *pgtypes.DoltgresType, elements []string) string { + separator := ", " + if elemType != nil && (elemType.IsArrayType() || elemType.IsCompositeType() || elemType.ID.TypeName() == "record") { + separator = ", \n " + } + return "[" + strings.Join(elements, separator) + "]" +} diff --git a/server/functions/array_to_json.go b/server/functions/array_to_json.go index 31f0dbc8d9..92e399c558 100644 --- a/server/functions/array_to_json.go +++ b/server/functions/array_to_json.go @@ -160,6 +160,13 @@ func valueToJsonRaw(ctx *sql.Context, elemType *pgtypes.DoltgresType, val any) ( } } +// ValueToJsonRaw converts a PostgreSQL value to its JSON representation. It is +// exported for aggregates such as json_agg, which share PostgreSQL's scalar +// value-to-JSON conversion rules with to_json and array_to_json. +func ValueToJsonRaw(ctx *sql.Context, elemType *pgtypes.DoltgresType, val any) (json.RawMessage, error) { + return valueToJsonRaw(ctx, elemType, val) +} + // marshalJsonNumber formats PostgreSQL numeric types through their output // functions. Non-finite float and numeric values are JSON strings because they // are not valid JSON number tokens. diff --git a/server/functions/framework/compiled_aggregate_function.go b/server/functions/framework/compiled_aggregate_function.go index ef6b401902..bac5035715 100644 --- a/server/functions/framework/compiled_aggregate_function.go +++ b/server/functions/framework/compiled_aggregate_function.go @@ -21,6 +21,8 @@ import ( "github.com/dolthub/go-mysql-server/sql" "github.com/dolthub/go-mysql-server/sql/expression" "github.com/dolthub/go-mysql-server/sql/transform" + + pgtypes "github.com/dolthub/doltgresql/server/types" ) // AggregateFunction is an expression that represents CompiledAggregateFunction @@ -152,7 +154,7 @@ func (c *CompiledAggregateFunction) NewBuffer(ctx *sql.Context) (sql.Aggregation // Buffers evaluate their argument expressions directly, without the GMS value conversion that // CompiledFunction.Eval performs, so any GMS-typed arguments (e.g. columns of the dolt_* system // tables) must be wrapped to convert their values. - return agg.NewBuffer(castGMSArguments(ctx, args)) + return agg.NewBuffer(withResolvedAggregateArgumentTypes(castGMSArguments(ctx, args), c.originalTypes)) } // Id implements the interface sql.Aggregation. @@ -182,7 +184,71 @@ func (c *CompiledAggregateFunction) NewWindowFunction(ctx *sql.Context) (sql.Win return nil, err } // See the comment in NewBuffer: GMS-typed arguments must convert their values. - return newWindowFunc(castGMSArguments(ctx, args), c.window) + return newWindowFunc(withResolvedAggregateArgumentTypes(castGMSArguments(ctx, args), c.originalTypes), c.window) +} + +// resolvedAggregateArgument preserves the concrete argument type selected for a +// polymorphic aggregate. Aggregate buffers evaluate cloned expressions directly +// and otherwise only see the declared pseudo-type (for example, anyarray), while +// conversion-sensitive aggregates such as json_agg need the concrete element +// type selected during overload resolution. +type resolvedAggregateArgument struct { + child sql.Expression + typ *pgtypes.DoltgresType +} + +var _ sql.Expression = (*resolvedAggregateArgument)(nil) + +// EvalAggregateArgument evaluates an argument and reports whether DISTINCT retained the value. +func EvalAggregateArgument(ctx *sql.Context, expr sql.Expression, row sql.Row) (interface{}, bool, error) { + if resolved, ok := expr.(*resolvedAggregateArgument); ok { + expr = resolved.child + } + if distinct, ok := expr.(*expression.DistinctExpression); ok { + return distinct.EvalDistinct(ctx, row) + } + value, err := expr.Eval(ctx, row) + return value, true, err +} + +// withResolvedAggregateArgumentTypes preserves the concrete types selected during overload resolution. +func withResolvedAggregateArgumentTypes(args []sql.Expression, types []*pgtypes.DoltgresType) []sql.Expression { + for i := range args { + if i < len(types) && types[i] != nil { + args[i] = &resolvedAggregateArgument{child: args[i], typ: types[i]} + } + } + return args +} + +// Resolved implements sql.Expression. +func (e *resolvedAggregateArgument) Resolved() bool { return e.child.Resolved() } + +// String implements sql.Expression. +func (e *resolvedAggregateArgument) String() string { return e.child.String() } + +// Type implements sql.Expression. +func (e *resolvedAggregateArgument) Type(ctx *sql.Context) sql.Type { return e.typ } + +// IsNullable implements sql.Expression. +func (e *resolvedAggregateArgument) IsNullable(ctx *sql.Context) bool { + return e.child.IsNullable(ctx) +} + +// Eval implements sql.Expression. +func (e *resolvedAggregateArgument) Eval(ctx *sql.Context, row sql.Row) (any, error) { + return e.child.Eval(ctx, row) +} + +// Children implements sql.Expression. +func (e *resolvedAggregateArgument) Children() []sql.Expression { return []sql.Expression{e.child} } + +// WithChildren implements sql.Expression. +func (e *resolvedAggregateArgument) WithChildren(ctx *sql.Context, children ...sql.Expression) (sql.Expression, error) { + if len(children) != 1 { + return nil, sql.ErrInvalidChildrenNumber.New(len(children), 1) + } + return &resolvedAggregateArgument{child: children[0], typ: e.typ}, nil } // cloneArguments returns a deep copy of args. Each partition/group gets its own AggregationBuffer or diff --git a/testing/go/functions_test.go b/testing/go/functions_test.go index 89d41199b5..6880cbc2dc 100644 --- a/testing/go/functions_test.go +++ b/testing/go/functions_test.go @@ -128,6 +128,127 @@ func TestAggregateFunctions(t *testing.T) { }, }, }, + { + Name: "json_agg", + SetUpScript: []string{ + `SET TIME ZONE 'UTC'`, + `CREATE TABLE json_agg_records (id int4, label text)`, + `INSERT INTO json_agg_records VALUES (1, 'one'), (2, NULL)`, + `CREATE TABLE json_agg_arrays (id int4 primary key, v int4[])`, + `INSERT INTO json_agg_arrays VALUES (1, ARRAY[1,NULL,3]), (2, ARRAY[4,5,NULL])`, + `CREATE TABLE json_agg_stored (id int4 primary key, amount numeric(40,20), payload json)`, + `INSERT INTO json_agg_stored VALUES (1, 12345678901234567890.12345678901234567890, '{"kind":"stored"}')`, + `CREATE DOMAIN json_agg_positive_int AS int4 CHECK (VALUE > 0)`, + }, + Assertions: []ScriptTestAssertion{ + { + Query: `SELECT json_agg(v) FROM (VALUES (1::int4),(2),(NULL)) AS t(v);`, + Expected: []sql.Row{{`[1, 2, null]`}}, + }, + { + Query: `SELECT json_agg(v) FROM (VALUES ('quote"slash\line'::text),(E'line\nnext')) AS t(v);`, + Expected: []sql.Row{{`["quote\"slash\\line", "line\nnext"]`}}, + }, + { + Query: `SELECT json_agg(v) FROM (VALUES (true),(false),(NULL::bool)) AS t(v);`, + Expected: []sql.Row{{`[true, false, null]`}}, + }, + { + Query: `SELECT json_agg(v) FROM (VALUES (1.2300::numeric),('-4.5'::numeric),('NaN'::numeric)) AS t(v);`, + Expected: []sql.Row{{`[1.2300, -4.5, "NaN"]`}}, + }, + { + Query: `SELECT json_agg(v) FROM (VALUES ('Infinity'::float8),('-Infinity'::float8),('NaN'::float8),(1.5::float8)) AS t(v);`, + Expected: []sql.Row{{`["Infinity", "-Infinity", "NaN", 1.5]`}}, + }, + { + Query: `SELECT json_agg(v) FROM (VALUES ('2024-02-29'::date),('0001-01-01 BC'::date)) AS t(v);`, + Expected: []sql.Row{{`["2024-02-29", "0001-01-01 BC"]`}}, + }, + { + Query: `SELECT json_agg(v) FROM (VALUES ('2024-02-29 12:34:56.123456'::timestamp),('2024-03-01 00:00:00'::timestamp)) AS t(v);`, + Expected: []sql.Row{{`["2024-02-29T12:34:56.123456", "2024-03-01T00:00:00"]`}}, + }, + { + Query: `SELECT json_agg(v) FROM (VALUES ('550e8400-e29b-41d4-a716-446655440000'::uuid),('00000000-0000-0000-0000-000000000000'::uuid)) AS t(v);`, + Expected: []sql.Row{{`["550e8400-e29b-41d4-a716-446655440000", "00000000-0000-0000-0000-000000000000"]`}}, + }, + { + Query: `SELECT json_agg(v) FROM (VALUES ('{"b": 2, "a": 1}'::json),('null'::json)) AS t(v);`, + Expected: []sql.Row{{`[{"b": 2, "a": 1}, null]`}}, + }, + { + Query: `SELECT json_agg(v) FROM (VALUES ('{"b": 2, "a": 1}'::jsonb),('[1, null]'::jsonb)) AS t(v);`, + Expected: []sql.Row{{`[{"a": 1, "b": 2}, [1, null]]`}}, + }, + { + Query: `SELECT json_agg(v) FROM (VALUES ('1'::json),('"two"'::json),('true'::json),('null'::json),('{"k":3}'::json),('[4]'::json)) AS t(v);`, + Expected: []sql.Row{{`[1, "two", true, null, {"k":3}, [4]]`}}, + }, + { + Query: `SELECT json_agg(v) FROM json_agg_arrays;`, + Expected: []sql.Row{{"[[1,null,3], \n [4,5,null]]"}}, + }, + { + Query: `SELECT json_agg(amount) FROM json_agg_stored;`, + Expected: []sql.Row{{`[12345678901234567890.12345678901234567890]`}}, + }, + { + Query: `SELECT json_agg(r) FROM (SELECT * FROM json_agg_stored ORDER BY id) r;`, + Expected: []sql.Row{{`[{"id":1,"amount":12345678901234567890.12345678901234567890,"payload":{"kind":"stored"}}]`}}, + }, + { + Query: `SELECT json_agg(NULL::text);`, + Expected: []sql.Row{{`[null]`}}, + }, + { + Query: `SELECT json_agg(v) FROM (SELECT 1::int AS v WHERE false) AS t;`, + Expected: []sql.Row{{nil}}, + }, + { + Query: `SELECT pg_typeof(json_agg(1));`, + Expected: []sql.Row{{"json"}}, + }, + { + Query: `SELECT json_agg(r) FROM (SELECT * FROM json_agg_records ORDER BY id) r;`, + Expected: []sql.Row{{`[{"id":1,"label":"one"}, ` + "\n " + `{"id":2,"label":null}]`}}, + }, + { + Query: `SELECT g, json_agg(v) FROM (VALUES ('a',1),('a',2),('b',3),('b',NULL)) AS t(g,v) GROUP BY g ORDER BY g;`, + Expected: []sql.Row{ + {"a", `[1, 2]`}, + {"b", `[3, null]`}, + }, + }, + { + Query: `SELECT json_agg(v) OVER (ORDER BY id ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) FROM (VALUES (1,10),(2,NULL),(3,30)) AS t(id,v) ORDER BY id;`, + Expected: []sql.Row{ + {`[10]`}, + {`[10, null]`}, + {`[10, null, 30]`}, + }, + }, + { + Query: `SELECT json_agg(v) OVER (ORDER BY id ROWS BETWEEN 1 PRECEDING AND 1 PRECEDING) FROM (VALUES (1,10),(2,20)) AS t(id,v) ORDER BY id;`, + Expected: []sql.Row{ + {nil}, + {`[10]`}, + }, + }, + { + Query: `SELECT json_agg(v) FROM (VALUES ('\x00ff'::bytea),(NULL::bytea)) AS t(v);`, + Expected: []sql.Row{{`["\\x00ff", null]`}}, + }, + { + Query: `SELECT json_agg(v) FROM (VALUES (1::json_agg_positive_int),(2::json_agg_positive_int)) AS t(v);`, + Expected: []sql.Row{{`[1, 2]`}}, + }, + { + Query: `SELECT json_agg(DISTINCT v) FROM (VALUES (1),(1),(NULL),(NULL)) AS t(v);`, + Expected: []sql.Row{{`[1, null]`}}, + }, + }, + }, { Name: "array_agg", SetUpScript: []string{