expression: implement vectorized evaluation for JSONInsert JSONSet JS…

…ONReplace (pingcap#13298)
matthewli968 · Nov 12, 2019 · 8404a43 · 8404a43
1 parent 5bc709f
commit 8404a43
Show file tree

Hide file tree

Showing 2 changed files with 95 additions and 11 deletions.
diff --git 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
@@ -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}},
 	},