expression: implement vectorized evaluation for builtinJSONMergeSig (…

…#12916)
pingcap · Oct 30, 2019 · a13bb53 · a13bb53
1 parent 98e066c
commit a13bb53
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 3 deletions.
diff --git a/expression/builtin_json_vec.go b/expression/builtin_json_vec.go
@@ -19,6 +19,7 @@ import (
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/types/json"
 	"github.com/pingcap/tidb/util/chunk"
+	"github.com/pingcap/tipb/go-tipb"
 )
 
 func (b *builtinJSONDepthSig) vectorized() bool {
@@ -398,11 +399,67 @@ func (b *builtinJSONRemoveSig) vecEvalJSON(input *chunk.Chunk, result *chunk.Col
 }
 
 func (b *builtinJSONMergeSig) vectorized() bool {
-	return false
+	return true
 }
 
 func (b *builtinJSONMergeSig) vecEvalJSON(input *chunk.Chunk, result *chunk.Column) error {
-	return errors.Errorf("not implemented")
+	nr := input.NumRows()
+	argBuffers := make([]*chunk.Column, len(b.args))
+	var err error
+	for i, arg := range b.args {
+		if argBuffers[i], err = b.bufAllocator.get(types.ETJson, nr); err != nil {
+			return err
+		}
+		defer func(buf *chunk.Column) {
+			b.bufAllocator.put(buf)
+		}(argBuffers[i])
+
+		if err := arg.VecEvalJSON(b.ctx, input, argBuffers[i]); err != nil {
+			return err
+		}
+	}
+
+	jsonValues := make([][]json.BinaryJSON, nr)
+
+	for i := 0; i < nr; i++ {
+		isNullFlag := false
+		for j := 0; j < len(b.args); j++ {
+			isNullFlag = isNullFlag || argBuffers[j].IsNull(i)
+		}
+		if isNullFlag {
+			jsonValues[i] = nil
+		} else {
+			jsonValues[i] = make([]json.BinaryJSON, 0, len(b.args))
+		}
+	}
+	for i := 0; i < len(b.args); i++ {
+		for j := 0; j < nr; j++ {
+			if jsonValues[j] == nil {
+				continue
+			}
+			jsonValues[j] = append(jsonValues[j], argBuffers[i].GetJSON(j))
+		}
+	}
+
+	result.ReserveJSON(nr)
+	for i := 0; i < nr; i++ {
+		if jsonValues[i] == nil {
+			result.AppendNull()
+			continue
+		}
+		result.AppendJSON(json.MergeBinary(jsonValues[i]))
+	}
+
+	if b.pbCode == tipb.ScalarFuncSig_JsonMergeSig {
+		for i := 0; i < nr; i++ {
+			if result.IsNull(i) {
+				continue
+			}
+			b.ctx.GetSessionVars().StmtCtx.AppendWarning(errDeprecatedSyntaxNoReplacement.GenWithStackByArgs("JSON_MERGE"))
+		}
+	}
+
+	return nil
 }
 
 func (b *builtinJSONContainsPathSig) vectorized() bool {

diff --git a/expression/builtin_json_vec_test.go b/expression/builtin_json_vec_test.go
@@ -82,7 +82,7 @@ var vecBuiltinJSONCases = map[string][]vecExprBenchCase{
 		{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{&jsonStringGener{}}},
 	},
 	ast.JSONRemove: {},
-	ast.JSONMerge:  {},
+	ast.JSONMerge:  {{retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson, types.ETJson, types.ETJson, types.ETJson, types.ETJson}}},
 	ast.JSONInsert: {},
 	ast.JSONQuote: {
 		{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETJson}},