diff --git a/expression/builtin_json.go b/expression/builtin_json.go index f2f2bb0198d4c..de1626886ac55 100644 --- a/expression/builtin_json.go +++ b/expression/builtin_json.go @@ -868,53 +868,53 @@ func (b *builtinJSONArrayAppendSig) evalJSON(row chunk.Row) (res json.BinaryJSON for i := 1; i < len(b.args)-1; i += 2 { // If JSON path is NULL, MySQL breaks and returns NULL. - s, isNull, err := b.args[i].EvalString(b.ctx, row) - if isNull || err != nil { + s, sNull, err := b.args[i].EvalString(b.ctx, row) + if sNull || err != nil { return res, true, err } - - // We should do the following checks to get correct values in res.Extract - pathExpr, err := json.ParseJSONPathExpr(s) - if err != nil { - return res, true, json.ErrInvalidJSONPath.GenWithStackByArgs(s) - } - if pathExpr.ContainsAnyAsterisk() { - return res, true, json.ErrInvalidJSONPathWildcard.GenWithStackByArgs(s) - } - - obj, exists := res.Extract([]json.PathExpression{pathExpr}) - if !exists { - // If path not exists, just do nothing and no errors. - continue - } - - if obj.TypeCode != json.TypeCodeArray { - // res.Extract will return a json object instead of an array if there is an object at path pathExpr. - // JSON_ARRAY_APPEND({"a": "b"}, "$", {"b": "c"}) => [{"a": "b"}, {"b", "c"}] - // We should wrap them to a single array first. - obj = json.CreateBinary([]interface{}{obj}) - } - - value, isnull, err := b.args[i+1].EvalJSON(b.ctx, row) + value, vNull, err := b.args[i+1].EvalJSON(b.ctx, row) if err != nil { return res, true, err } - - if isnull { + if vNull { value = json.CreateBinary(nil) } - - obj = json.MergeBinary([]json.BinaryJSON{obj, value}) - res, err = res.Modify([]json.PathExpression{pathExpr}, []json.BinaryJSON{obj}, json.ModifySet) - if err != nil { - // We checked pathExpr in the same way as res.Modify do. - // So err should always be nil, the function should never return here. - return res, true, err + res, isNull, err = b.appendJSONArray(res, s, value) + if isNull || err != nil { + return res, isNull, err } } return res, false, nil } +func (b *builtinJSONArrayAppendSig) appendJSONArray(res json.BinaryJSON, p string, v json.BinaryJSON) (json.BinaryJSON, bool, error) { + // We should do the following checks to get correct values in res.Extract + pathExpr, err := json.ParseJSONPathExpr(p) + if err != nil { + return res, true, json.ErrInvalidJSONPath.GenWithStackByArgs(p) + } + if pathExpr.ContainsAnyAsterisk() { + return res, true, json.ErrInvalidJSONPathWildcard.GenWithStackByArgs(p) + } + + obj, exists := res.Extract([]json.PathExpression{pathExpr}) + if !exists { + // If path not exists, just do nothing and no errors. + return res, false, nil + } + + if obj.TypeCode != json.TypeCodeArray { + // res.Extract will return a json object instead of an array if there is an object at path pathExpr. + // JSON_ARRAY_APPEND({"a": "b"}, "$", {"b": "c"}) => [{"a": "b"}, {"b", "c"}] + // We should wrap them to a single array first. + obj = json.CreateBinary([]interface{}{obj}) + } + + obj = json.MergeBinary([]json.BinaryJSON{obj, v}) + res, err = res.Modify([]json.PathExpression{pathExpr}, []json.BinaryJSON{obj}, json.ModifySet) + return res, false, err +} + type jsonArrayInsertFunctionClass struct { baseFunctionClass } diff --git a/expression/builtin_json_vec.go b/expression/builtin_json_vec.go index 0a6213788f56c..ce7507a90502b 100644 --- a/expression/builtin_json_vec.go +++ b/expression/builtin_json_vec.go @@ -967,11 +967,87 @@ func (b *builtinJSONContainsPathSig) vecEvalInt(input *chunk.Chunk, result *chun } func (b *builtinJSONArrayAppendSig) vectorized() bool { - return false + return true } func (b *builtinJSONArrayAppendSig) vecEvalJSON(input *chunk.Chunk, result *chunk.Column) error { - return errors.Errorf("not implemented") + n := input.NumRows() + m := (len(b.args) - 1) / 2 + + jsonBufs, err := b.bufAllocator.get(types.ETJson, n) + if err != nil { + return err + } + defer b.bufAllocator.put(jsonBufs) + if err := b.args[0].VecEvalJSON(b.ctx, input, jsonBufs); err != nil { + return err + } + + pathBufs := make([]*chunk.Column, 0, m) + valBufs := make([]*chunk.Column, 0, m) + defer func() { + for _, buf := range pathBufs { + b.bufAllocator.put(buf) + } + for _, buf := range valBufs { + b.bufAllocator.put(buf) + } + }() + for i := 1; i < len(b.args)-1; i += 2 { + pathBuf, err := b.bufAllocator.get(types.ETString, n) + if err != nil { + return err + } + pathBufs = append(pathBufs, pathBuf) + if err := b.args[i].VecEvalString(b.ctx, input, pathBuf); err != nil { + return err + } + valBuf, err := b.bufAllocator.get(types.ETJson, n) + if err != nil { + return err + } + if err := b.args[i+1].VecEvalJSON(b.ctx, input, valBuf); err != nil { + return err + } + valBufs = append(valBufs, valBuf) + } + + result.ReserveJSON(n) + for i := 0; i < n; i++ { + if jsonBufs.IsNull(i) { + result.AppendNull() + continue + } + res := jsonBufs.GetJSON(i) + isNull := false + for j := 0; j < m; j++ { + if pathBufs[j].IsNull(i) { + isNull = true + break + } + s := pathBufs[j].GetString(i) + v, vNull := json.BinaryJSON{}, valBufs[j].IsNull(i) + if !vNull { + v = valBufs[j].GetJSON(i) + } + if vNull { + v = json.CreateBinary(nil) + } + res, isNull, err = b.appendJSONArray(res, s, v) + if err != nil { + return err + } + if isNull { + break + } + } + if isNull { + result.AppendNull() + } else { + result.AppendJSON(res) + } + } + return nil } func (b *builtinJSONUnquoteSig) vectorized() bool { diff --git a/expression/builtin_json_vec_test.go b/expression/builtin_json_vec_test.go index 663659582c384..e8cf08f4675a7 100644 --- a/expression/builtin_json_vec_test.go +++ b/expression/builtin_json_vec_test.go @@ -26,7 +26,26 @@ var vecBuiltinJSONCases = map[string][]vecExprBenchCase{ {retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson}}, {retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson, types.ETString}, geners: []dataGenerator{&constJSONGener{"{\"a\": {\"c\": 3}, \"b\": 2}"}, &constStrGener{"$.a"}}}, }, - ast.JSONArrayAppend: {}, + ast.JSONArrayAppend: { + {retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson, types.ETString, types.ETJson}, + geners: []dataGenerator{newNullWrappedGener(0.1, &constJSONGener{"{\"a\": {\"c\": 3}, \"b\": 2}"}), + newNullWrappedGener(0.1, &constStrGener{"$.a"}), + newNullWrappedGener(0.1, &constJSONGener{"{\"b\": 2}"})}}, + {retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson, types.ETString, types.ETJson, types.ETString, types.ETJson}, + geners: []dataGenerator{newNullWrappedGener(0.1, &constJSONGener{"{\"a\": {\"c\": 3}, \"b\": 2}"}), + newNullWrappedGener(0.1, &constStrGener{"$.a"}), + newNullWrappedGener(0.1, &constJSONGener{"{\"b\": 2}"}), + newNullWrappedGener(0.1, &constStrGener{"$.b"}), + newNullWrappedGener(0.1, &constJSONGener{"{\"x\": 3}"}), + }}, + {retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson, types.ETString, types.ETJson, types.ETString, types.ETJson}, + geners: []dataGenerator{newNullWrappedGener(0.1, &constJSONGener{"{\"a\": {\"c\": 3}, \"b\": 2}"}), + newNullWrappedGener(0.1, &constStrGener{"$.a"}), + newNullWrappedGener(0.1, &constJSONGener{"{\"b\": 2}"}), + newNullWrappedGener(0.1, &constStrGener{"$.x"}), // not exists + newNullWrappedGener(0.1, &constJSONGener{"{\"x\": 3}"}), + }}, + }, ast.JSONContainsPath: { {retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETJson, types.ETString, types.ETString}, geners: []dataGenerator{&constJSONGener{"{\"a\": {\"c\": {\"d\": 4}}, \"b\": 2}"}, &constStrGener{"one"}, &constStrGener{"$.c"}}}, {retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETJson, types.ETString, types.ETString}, geners: []dataGenerator{&constJSONGener{"{\"a\": {\"c\": {\"d\": 4}}, \"b\": 2}"}, &constStrGener{"all"}, &constStrGener{"$.a"}, &constStrGener{"$.c"}}},