diff --git a/expression/builtin_json_vec.go b/expression/builtin_json_vec.go index 297ac95dda69c..97c3f2adeb935 100644 --- a/expression/builtin_json_vec.go +++ b/expression/builtin_json_vec.go @@ -16,12 +16,87 @@ package expression import ( "github.com/pingcap/errors" "github.com/pingcap/parser/ast" + "github.com/pingcap/tidb/sessionctx" "github.com/pingcap/tidb/types" "github.com/pingcap/tidb/types/json" "github.com/pingcap/tidb/util/chunk" "github.com/pingcap/tipb/go-tipb" ) +func vecJSONModify(ctx sessionctx.Context, args []Expression, bufAllocator columnBufferAllocator, input *chunk.Chunk, result *chunk.Column, mt json.ModifyType) error { + nr := input.NumRows() + jsonBuf, err := bufAllocator.get(types.ETJson, nr) + if err != nil { + return err + } + defer bufAllocator.put(jsonBuf) + if err := args[0].VecEvalJSON(ctx, input, jsonBuf); err != nil { + return err + } + + strBufs := make([]*chunk.Column, (len(args)-1)/2) + for i := 1; i < len(args); i += 2 { + strBufs[(i-1)/2], err = bufAllocator.get(types.ETString, nr) + if err != nil { + return err + } + defer bufAllocator.put(strBufs[(i-1)/2]) + if err := args[i].VecEvalString(ctx, input, strBufs[(i-1)/2]); err != nil { + return err + } + } + valueBufs := make([]*chunk.Column, (len(args)-1)/2+1) + for i := 2; i < len(args); i += 2 { + valueBufs[i/2-1], err = bufAllocator.get(types.ETJson, nr) + if err != nil { + return err + } + defer bufAllocator.put(valueBufs[i/2-1]) + if err := args[i].VecEvalJSON(ctx, input, valueBufs[i/2-1]); err != nil { + return err + } + } + result.ReserveJSON(nr) + for i := 0; i < nr; i++ { + if jsonBuf.IsNull(i) { + result.AppendNull() + continue + } + pathExprs := make([]json.PathExpression, 0, (len(args)-1)/2+1) + values := make([]json.BinaryJSON, 0, (len(args)-1)/2+1) + var pathExpr json.PathExpression + isNull := false + for j := 1; j < len(args); j += 2 { + if strBufs[(j-1)/2].IsNull(i) { + isNull = true + break + } + pathExpr, err = json.ParseJSONPathExpr(strBufs[(j-1)/2].GetString(i)) + if err != nil { + return err + } + pathExprs = append(pathExprs, pathExpr) + } + for j := 2; j < len(args); j += 2 { + if valueBufs[j/2-1].IsNull(i) { + values = append(values, json.CreateBinary(nil)) + } else { + values = append(values, valueBufs[j/2-1].GetJSON(i)) + } + } + if isNull { + result.AppendNull() + } else { + res, err := jsonBuf.GetJSON(i).Modify(pathExprs, values, mt) + if err != nil { + return err + } + result.AppendJSON(res) + } + } + return nil +} + func (b *builtinJSONDepthSig) vectorized() bool { return true } @@ -82,19 +157,21 @@ func (b *builtinJSONKeysSig) vecEvalJSON(input *chunk.Chunk, result *chunk.Colum } func (b *builtinJSONInsertSig) vectorized() bool { - return false + return true } func (b *builtinJSONInsertSig) vecEvalJSON(input *chunk.Chunk, result *chunk.Column) error { - return errors.Errorf("not implemented") + err := vecJSONModify(b.ctx, b.args, b.bufAllocator, input, result, json.ModifyInsert) + return err } func (b *builtinJSONReplaceSig) vectorized() bool { - return false + return true } func (b *builtinJSONReplaceSig) vecEvalJSON(input *chunk.Chunk, result *chunk.Column) error { - return errors.Errorf("not implemented") + err := vecJSONModify(b.ctx, b.args, b.bufAllocator, input, result, json.ModifyReplace) + return err } func (b *builtinJSONArraySig) vectorized() bool { @@ -251,11 +328,12 @@ func (b *builtinJSONSearchSig) vecEvalJSON(input *chunk.Chunk, result *chunk.Col } func (b *builtinJSONSetSig) vectorized() bool { - return false + return true } func (b *builtinJSONSetSig) vecEvalJSON(input *chunk.Chunk, result *chunk.Column) error { - return errors.Errorf("not implemented") + err := vecJSONModify(b.ctx, b.args, b.bufAllocator, input, result, json.ModifySet) + return err } func (b *builtinJSONObjectSig) vectorized() bool { diff --git a/expression/builtin_json_vec_test.go b/expression/builtin_json_vec_test.go index 5ecce455a0355..55197aae6c473 100644 --- a/expression/builtin_json_vec_test.go +++ b/expression/builtin_json_vec_test.go @@ -82,16 +82,22 @@ var vecBuiltinJSONCases = map[string][]vecExprBenchCase{ }, }, }, - ast.JSONSet: {}, - ast.JSONSearch: {}, - ast.JSONReplace: {}, - ast.JSONDepth: {{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETJson}}}, + ast.JSONSet: { + {retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson, types.ETString, types.ETJson, types.ETString, types.ETJson}, geners: []dataGenerator{nil, &constStrGener{"$.key"}, nil, &constStrGener{"$.aaa"}, nil}}, + }, + ast.JSONSearch: {}, + ast.JSONReplace: { + {retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson, types.ETString, types.ETJson}, geners: []dataGenerator{nil, &constStrGener{"$.key"}, nil}}, + }, + ast.JSONDepth: {{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETJson}}}, ast.JSONUnquote: { {retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{&jsonStringGener{}}}, }, ast.JSONRemove: {}, ast.JSONMerge: {{retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson, types.ETJson, types.ETJson, types.ETJson, types.ETJson}}}, - ast.JSONInsert: {}, + ast.JSONInsert: { + {retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson, types.ETString, types.ETJson, types.ETString, types.ETJson}, geners: []dataGenerator{nil, &constStrGener{"$.aaa"}, nil, &constStrGener{"$.bbb"}, nil}}, + }, ast.JSONQuote: { {retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETJson}}, },