expression: implement vectorized evaluation for builtinJSONObjectSig (

pingcap#12663)
XiaTianliang · Dec 21, 2019 · e000a35 · e000a35
1 parent 5248403
commit e000a35
Show file tree

Hide file tree

Showing 2 changed files with 106 additions and 7 deletions.
diff --git a/expression/builtin_json_vec.go b/expression/builtin_json_vec.go
@@ -15,6 +15,7 @@ package expression
 
 import (
 	"github.com/pingcap/errors"
+	"github.com/pingcap/parser/ast"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/types/json"
 	"github.com/pingcap/tidb/util/chunk"
@@ -154,11 +155,77 @@ func (b *builtinJSONSetSig) vecEvalJSON(input *chunk.Chunk, result *chunk.Column
 }
 
 func (b *builtinJSONObjectSig) vectorized() bool {
-	return false
+	return true
 }
 
 func (b *builtinJSONObjectSig) vecEvalJSON(input *chunk.Chunk, result *chunk.Column) error {
-	return errors.Errorf("not implemented")
+	nr := input.NumRows()
+	if len(b.args)&1 == 1 {
+		err := ErrIncorrectParameterCount.GenWithStackByArgs(ast.JSONObject)
+		return err
+	}
+
+	jsons := make([]map[string]interface{}, nr)
+	for i := 0; i < nr; i++ {
+		jsons[i] = make(map[string]interface{}, len(b.args)>>1)
+	}
+
+	argBuffers := make([]*chunk.Column, len(b.args))
+	var err error
+	for i := 0; i < len(b.args); i++ {
+		if i&1 == 0 {
+			if argBuffers[i], err = b.bufAllocator.get(types.ETString, nr); err != nil {
+				return err
+			}
+			defer func(buf *chunk.Column) {
+				b.bufAllocator.put(buf)
+			}(argBuffers[i])
+
+			if err = b.args[i].VecEvalString(b.ctx, input, argBuffers[i]); err != nil {
+				return err
+			}
+		} else {
+			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 = b.args[i].VecEvalJSON(b.ctx, input, argBuffers[i]); err != nil {
+				return err
+			}
+		}
+	}
+
+	result.ReserveJSON(nr)
+	for i := 0; i < len(b.args); i++ {
+		if i&1 == 1 {
+			keyCol := argBuffers[i-1]
+			valueCol := argBuffers[i]
+
+			var key string
+			var value json.BinaryJSON
+			for j := 0; j < nr; j++ {
+				if keyCol.IsNull(j) {
+					err := errors.New("JSON documents may not contain NULL member names")
+					return err
+				}
+				key = keyCol.GetString(j)
+				if valueCol.IsNull(j) {
+					value = json.CreateBinary(nil)
+				} else {
+					value = valueCol.GetJSON(j)
+				}
+				jsons[j][key] = value
+			}
+		}
+	}
+
+	for i := 0; i < nr; i++ {
+		result.AppendJSON(json.CreateBinary(jsons[i]))
+	}
+	return nil
 }
 
 func (b *builtinJSONArrayInsertSig) vectorized() bool {

diff --git a/expression/builtin_json_vec_test.go b/expression/builtin_json_vec_test.go
@@ -33,11 +33,43 @@ var vecBuiltinJSONCases = map[string][]vecExprBenchCase{
 	ast.JSONArray:        {},
 	ast.JSONArrayInsert:  {},
 	ast.JSONContains:     {},
-	ast.JSONObject:       {},
-	ast.JSONSet:          {},
-	ast.JSONSearch:       {},
-	ast.JSONReplace:      {},
-	ast.JSONDepth:        {{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETJson}}},
+	ast.JSONObject: {
+		{
+			retEvalType: types.ETJson,
+			childrenTypes: []types.EvalType{
+				types.ETString, types.ETJson,
+				types.ETString, types.ETJson,
+				types.ETString, types.ETJson,
+				types.ETString, types.ETJson,
+				types.ETString, types.ETJson,
+				types.ETString, types.ETJson,
+				types.ETString, types.ETJson,
+				types.ETString, types.ETJson,
+				types.ETString, types.ETJson,
+				types.ETString, types.ETJson,
+				types.ETString, types.ETJson,
+				types.ETString, types.ETJson,
+			},
+			geners: []dataGenerator{
+				&randLenStrGener{10, 20}, nil,
+				&randLenStrGener{10, 20}, nil,
+				&randLenStrGener{10, 20}, nil,
+				&randLenStrGener{10, 20}, nil,
+				&randLenStrGener{10, 20}, nil,
+				&randLenStrGener{10, 20}, nil,
+				&randLenStrGener{10, 20}, nil,
+				&randLenStrGener{10, 20}, nil,
+				&randLenStrGener{10, 20}, nil,
+				&randLenStrGener{10, 20}, nil,
+				&randLenStrGener{10, 20}, nil,
+				&randLenStrGener{10, 20}, nil,
+			},
+		},
+	},
+	ast.JSONSet:     {},
+	ast.JSONSearch:  {},
+	ast.JSONReplace: {},
+	ast.JSONDepth:   {{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETJson}}},
 	ast.JSONUnquote: {
 		{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETString}, geners: []dataGenerator{&jsonStringGener{}}},
 	},