From 2762b631f37f396a9dfbb82965daa71f4423716a Mon Sep 17 00:00:00 2001 From: Aaron Zinger Date: Mon, 23 May 2022 17:48:20 -0400 Subject: [PATCH] sql: implement array functions for JSON arrays My answer on https://stackoverflow.com/questions/72347170/delete-element-from-jsonb-array-in-cockaroachdb/72351955#72351955 made me sad, and it's not the first time I've had trouble manipulating a JSONB array. There are a few different, non-mutually-exclusive ways we could make this easier. We could allow aggregate(generator) functions in scalar contexts. We could implement more JSON functions as needed. This PR instead just overloads every function that takes a json[] argument (which is a datum type we don't allow in columns) to also be able to take a json array argument, and alters the return type accordingly. It's not particularly performance-optimized but we can replace individual implementations if needed. That said, this is an unauthorized spike, what do you think? Release note (sql change): Built-in array functions array_append, array_prepend, array_cat, array_remove, array_replace, array_position, and array_positions may now be used with jsonb arrays. --- docs/generated/sql/functions.md | 14 ++++ pkg/sql/logictest/testdata/logic_test/array | 5 ++ pkg/sql/sem/builtins/builtins.go | 77 ++++++++++++++++++++- 3 files changed, 95 insertions(+), 1 deletion(-) diff --git a/docs/generated/sql/functions.md b/docs/generated/sql/functions.md index c074d68acdb3..f27b9165c39e 100644 --- a/docs/generated/sql/functions.md +++ b/docs/generated/sql/functions.md @@ -37,6 +37,8 @@ array_append(array: geometry[], elem: geometry) → geometry[]

Appends elem to array, returning the result.

 +array_append(array: jsonb, elem: jsonb) → jsonb

Appends elem to array, returning the result.

+ array_append(array: jsonb[], elem: jsonb) → jsonb[]

Appends elem to array, returning the result.

 array_append(array: oid[], elem: oid) → oid[]

Appends elem to array, returning the result.

@@ -81,6 +83,8 @@ array_cat(left: geometry[], right: geometry[]) → geometry[]

Appends two arrays.

 +array_cat(left: jsonb, right: jsonb) → jsonb

Appends two arrays.

+ array_cat(left: jsonb[], right: jsonb[]) → jsonb[]

Appends two arrays.

 array_cat(left: oid[], right: oid[]) → oid[]

Appends two arrays.

@@ -129,6 +133,8 @@ array_position(array: geometry[], elem: geometry) → int

Return the index of the first occurrence of elem in array.

 +array_position(array: jsonb, elem: jsonb) → int

Return the index of the first occurrence of elem in array.

+ array_position(array: jsonb[], elem: jsonb) → int

Return the index of the first occurrence of elem in array.

 array_position(array: oid[], elem: oid) → int

Return the index of the first occurrence of elem in array.

@@ -173,6 +179,8 @@ array_positions(array: geometry[], elem: geometry) → int[]

Returns and array of indexes of all occurrences of elem in array.

 +array_positions(array: jsonb, elem: jsonb) → int[]

Returns and array of indexes of all occurrences of elem in array.

+ array_positions(array: jsonb[], elem: jsonb) → int[]

Returns and array of indexes of all occurrences of elem in array.

 array_positions(array: oid[], elem: oid) → int[]

Returns and array of indexes of all occurrences of elem in array.

@@ -217,6 +225,8 @@ array_prepend(elem: geometry, array: geometry[]) → geometry[]

Prepends elem to array, returning the result.

 +array_prepend(elem: jsonb, array: jsonb) → jsonb

Prepends elem to array, returning the result.

+ array_prepend(elem: jsonb, array: jsonb[]) → jsonb[]

Prepends elem to array, returning the result.

 array_prepend(elem: oid, array: oid[]) → oid[]

Prepends elem to array, returning the result.

@@ -261,6 +271,8 @@ array_remove(array: geometry[], elem: geometry) → geometry[]

Remove from array all elements equal to elem.

 +array_remove(array: jsonb, elem: jsonb) → jsonb

Remove from array all elements equal to elem.

+ array_remove(array: jsonb[], elem: jsonb) → jsonb[]

Remove from array all elements equal to elem.

 array_remove(array: oid[], elem: oid) → oid[]

Remove from array all elements equal to elem.

@@ -305,6 +317,8 @@ array_replace(array: geometry[], toreplace: geometry, replacewith: geometry) → geometry[]

Replace all occurrences of toreplace in array with replacewith.

 +array_replace(array: jsonb, toreplace: jsonb, replacewith: jsonb) → jsonb

Replace all occurrences of toreplace in array with replacewith.

+ array_replace(array: jsonb[], toreplace: jsonb, replacewith: jsonb) → jsonb[]

Replace all occurrences of toreplace in array with replacewith.

 array_replace(array: oid[], toreplace: oid, replacewith: oid) → oid[]

Replace all occurrences of toreplace in array with replacewith.

diff --git a/pkg/sql/logictest/testdata/logic_test/array b/pkg/sql/logictest/testdata/logic_test/array index 097097afccf4..3b940b0c3b86 100644 --- a/pkg/sql/logictest/testdata/logic_test/array +++ b/pkg/sql/logictest/testdata/logic_test/array @@ -1070,6 +1070,11 @@ SELECT array_remove(NULL::INT[], NULL::INT) ---- NULL +query T +SELECT array_remove('[1,2,3,2]'::jsonb, '2'::jsonb) +---- +[1, 3] + # ARRAY_REPLACE function query T diff --git a/pkg/sql/sem/builtins/builtins.go b/pkg/sql/sem/builtins/builtins.go index d3df4489f7ae..ac7bce93ce12 100644 --- a/pkg/sql/sem/builtins/builtins.go +++ b/pkg/sql/sem/builtins/builtins.go @@ -8008,12 +8008,18 @@ var similarOverloads = []tree.Overload{ } func arrayBuiltin(impl func(*types.T) tree.Overload) builtinDefinition { - overloads := make([]tree.Overload, 0, len(types.Scalar)+2) + overloads := make([]tree.Overload, 0, len(types.Scalar)+3) for _, typ := range append(types.Scalar, types.AnyEnum) { if ok, _ := types.IsValidArrayElementType(typ); ok { overloads = append(overloads, impl(typ)) } } + + arrayOfJSONImpl := impl(types.Jsonb) + if _, ok := arrayOfJSONImpl.Types.(tree.ArgTypes); ok && arrayOfJSONImpl.Fn != nil { + overloads = append(overloads, makeJSONArrayOverload(arrayOfJSONImpl)) + } + // Prevent usage in DistSQL because it cannot handle arrays of untyped tuples. tupleOverload := impl(types.AnyTuple) tupleOverload.DistsqlBlocklist = true @@ -8024,6 +8030,75 @@ func arrayBuiltin(impl func(*types.T) tree.Overload) builtinDefinition { } } +// Convert a json[] function to a json_array function. +func makeJSONArrayOverload(arrayOfJSONImpl tree.Overload) tree.Overload { + isJSONArray := func(t *types.T) bool { + return t != nil && t.ArrayContents() == types.Jsonb + } + castToJSON := func(t *types.T) *types.T { + if isJSONArray(t) { + return types.Jsonb + } + return t + } + newImpl := arrayOfJSONImpl + newImpl.Types = make(tree.ArgTypes, len(arrayOfJSONImpl.Types.(tree.ArgTypes))) + argsToCast := make([]int, 0) + // TODO: Don't panic if it's not an ArgTypes + for i, t := range arrayOfJSONImpl.Types.(tree.ArgTypes) { + if isJSONArray(t.Typ) { + argsToCast = append(argsToCast, i) + newImpl.Types.(tree.ArgTypes)[i] = struct { + Name string + Typ *types.T + }{t.Name, types.Jsonb} + } else { + newImpl.Types.(tree.ArgTypes)[i] = t + } + } + newImpl.ReturnType = func(args []tree.TypedExpr) *types.T { + return castToJSON(arrayOfJSONImpl.ReturnType(args)) + } + newImpl.Fn = func(ctx *eval.Context, args tree.Datums) (tree.Datum, error) { + newArgs := append(tree.Datums{}, args...) + for _, i := range argsToCast { + j := tree.MustBeDJSON(args[i]) + if j.Type() != json.ArrayJSONType { + return nil, pgerror.Newf(pgcode.InvalidParameterValue, + "cannot call array function on non-array json value") + } + ary := tree.NewDArray(types.Jsonb) + for idx := 0; idx < j.Len(); idx++ { + json, err := j.FetchValIdx(idx) + if err != nil { + return nil, err + } + d, err := tree.MakeDJSON(json) + if err != nil { + return nil, err + } + if err := ary.Append(d); err != nil { + return nil, err + } + } + newArgs[i] = ary + } + ret, err := arrayOfJSONImpl.Fn.(eval.FnOverload)(ctx, newArgs) + if err == nil && isJSONArray(ret.ResolvedType()) { + dArray := tree.MustBeDArray(ret).Array + builder := json.NewArrayBuilder(len(dArray)) + for _, elt := range dArray { + builder.Add(tree.MustBeDJSON(elt).JSON) + } + return tree.NewDJSON(builder.Build()), nil + } + return ret, err + } + + return newImpl + +} + func setProps(props tree.FunctionProperties, d builtinDefinition) builtinDefinition { d.props = props return d