expression: implement vectorized evaluation for builtinJSONArrayInser…

…tSig (pingcap#13667)
XiaTianliang · Dec 21, 2019 · 64e25e9 · 64e25e9
1 parent d9a68de
commit 64e25e9
Show file tree

Hide file tree

Showing 2 changed files with 77 additions and 3 deletions.
diff --git a/expression/builtin_json_vec.go b/expression/builtin_json_vec.go
@@ -411,11 +411,82 @@ func (b *builtinJSONObjectSig) vecEvalJSON(input *chunk.Chunk, result *chunk.Col
 }
 
 func (b *builtinJSONArrayInsertSig) vectorized() bool {
-	return false
+	return true
 }
 
 func (b *builtinJSONArrayInsertSig) vecEvalJSON(input *chunk.Chunk, result *chunk.Column) error {
-	return errors.Errorf("not implemented")
+	nr := input.NumRows()
+	buf, err := b.bufAllocator.get(types.ETJson, nr)
+	if err != nil {
+		return err
+	}
+	defer b.bufAllocator.put(buf)
+	if err := b.args[0].VecEvalJSON(b.ctx, input, buf); err != nil {
+		return err
+	}
+	pathBufs := make([]*chunk.Column, (len(b.args)-1)/2)
+	valueBufs := make([]*chunk.Column, (len(b.args)-1)/2)
+	for i := 1; i < len(b.args); i++ {
+		if i&1 == 0 {
+			valueBufs[i/2-1], err = b.bufAllocator.get(types.ETJson, nr)
+			if err != nil {
+				return err
+			}
+			defer b.bufAllocator.put(valueBufs[i/2-1])
+			if err := b.args[i].VecEvalJSON(b.ctx, input, valueBufs[i/2-1]); err != nil {
+				return err
+			}
+		} else {
+			pathBufs[(i-1)/2], err = b.bufAllocator.get(types.ETString, nr)
+			if err != nil {
+				return err
+			}
+			defer b.bufAllocator.put(pathBufs[(i-1)/2])
+			if err := b.args[i].VecEvalString(b.ctx, input, pathBufs[(i-1)/2]); err != nil {
+				return err
+			}
+		}
+	}
+	var pathExpr json.PathExpression
+	var value json.BinaryJSON
+	result.ReserveJSON(nr)
+	for i := 0; i < nr; i++ {
+		if buf.IsNull(i) {
+			result.AppendNull()
+			continue
+		}
+
+		res := buf.GetJSON(i)
+		isnull := false
+		for j := 0; j < (len(b.args)-1)/2; j++ {
+			if pathBufs[j].IsNull(i) {
+				isnull = true
+				break
+			}
+			pathExpr, err = json.ParseJSONPathExpr(pathBufs[j].GetString(i))
+			if err != nil {
+				return json.ErrInvalidJSONPath.GenWithStackByArgs(pathBufs[j].GetString(i))
+			}
+			if pathExpr.ContainsAnyAsterisk() {
+				return json.ErrInvalidJSONPathWildcard.GenWithStackByArgs(pathBufs[j].GetString(i))
+			}
+			if valueBufs[j].IsNull(i) {
+				value = json.CreateBinary(nil)
+			} else {
+				value = valueBufs[j].GetJSON(i)
+			}
+			res, err = res.ArrayInsert(pathExpr, value)
+			if err != nil {
+				return err
+			}
+		}
+		if isnull {
+			result.AppendNull()
+			continue
+		}
+		result.AppendJSON(res)
+	}
+	return nil
 }
 
 func (b *builtinJSONKeys2ArgsSig) vectorized() bool {

diff --git a/expression/builtin_json_vec_test.go b/expression/builtin_json_vec_test.go
@@ -43,7 +43,10 @@ var vecBuiltinJSONCases = map[string][]vecExprBenchCase{
 		{retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson, types.ETJson}},
 		{retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson, types.ETJson, types.ETJson}},
 	},
-	ast.JSONArrayInsert: {},
+	ast.JSONArrayInsert: {
+		{retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson, types.ETString, types.ETJson, types.ETString, types.ETJson}, geners: []dataGenerator{&constJSONGener{"[\"a\", {\"b\": [1, 2]}, [3, 4]]"}, &constStrGener{"$[1]"}, nil, &constStrGener{"$[2][1]"}, nil}},
+		{retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson, types.ETString, types.ETJson, types.ETString, types.ETJson}, geners: []dataGenerator{&constJSONGener{"[\"a\", {\"b\": [1, 2]}, [3, 4]]"}, &constStrGener{"$[1]"}, &constJSONGener{"[1,2,3,4,5]"}, &constStrGener{"$[2][1]"}, &constJSONGener{"{\"abc\":1,\"def\":2,\"ghi\":[1,2,3,4]}"}}},
+	},
 	ast.JSONContains: {
 		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETJson, types.ETJson}},
 		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETJson, types.ETJson, types.ETString}, geners: []dataGenerator{nil, nil, &constStrGener{"$.abc"}}},