expression: implement vectorized evaluation for `JSONArrayAppe… (#16379)

pingcap · Apr 15, 2020 · 3d30871 · 3d30871
1 parent 830e723
commit 3d30871
Show file tree

Hide file tree

Showing 3 changed files with 133 additions and 38 deletions.
diff --git a/expression/builtin_json.go b/expression/builtin_json.go
@@ -868,53 +868,53 @@ func (b *builtinJSONArrayAppendSig) evalJSON(row chunk.Row) (res json.BinaryJSON
 
 	for i := 1; i < len(b.args)-1; i += 2 {
 		// If JSON path is NULL, MySQL breaks and returns NULL.
-		s, isNull, err := b.args[i].EvalString(b.ctx, row)
-		if isNull || err != nil {
+		s, sNull, err := b.args[i].EvalString(b.ctx, row)
+		if sNull || err != nil {
 			return res, true, err
 		}
-
-		// We should do the following checks to get correct values in res.Extract
-		pathExpr, err := json.ParseJSONPathExpr(s)
-		if err != nil {
-			return res, true, json.ErrInvalidJSONPath.GenWithStackByArgs(s)
-		}
-		if pathExpr.ContainsAnyAsterisk() {
-			return res, true, json.ErrInvalidJSONPathWildcard.GenWithStackByArgs(s)
-		}
-
-		obj, exists := res.Extract([]json.PathExpression{pathExpr})
-		if !exists {
-			// If path not exists, just do nothing and no errors.
-			continue
-		}
-
-		if obj.TypeCode != json.TypeCodeArray {
-			// res.Extract will return a json object instead of an array if there is an object at path pathExpr.
-			// JSON_ARRAY_APPEND({"a": "b"}, "$", {"b": "c"}) => [{"a": "b"}, {"b", "c"}]
-			// We should wrap them to a single array first.
-			obj = json.CreateBinary([]interface{}{obj})
-		}
-
-		value, isnull, err := b.args[i+1].EvalJSON(b.ctx, row)
+		value, vNull, err := b.args[i+1].EvalJSON(b.ctx, row)
 		if err != nil {
 			return res, true, err
 		}
-
-		if isnull {
+		if vNull {
 			value = json.CreateBinary(nil)
 		}
-
-		obj = json.MergeBinary([]json.BinaryJSON{obj, value})
-		res, err = res.Modify([]json.PathExpression{pathExpr}, []json.BinaryJSON{obj}, json.ModifySet)
-		if err != nil {
-			// We checked pathExpr in the same way as res.Modify do.
-			// So err should always be nil, the function should never return here.
-			return res, true, err
+		res, isNull, err = b.appendJSONArray(res, s, value)
+		if isNull || err != nil {
+			return res, isNull, err
 		}
 	}
 	return res, false, nil
 }
 
+func (b *builtinJSONArrayAppendSig) appendJSONArray(res json.BinaryJSON, p string, v json.BinaryJSON) (json.BinaryJSON, bool, error) {
+	// We should do the following checks to get correct values in res.Extract
+	pathExpr, err := json.ParseJSONPathExpr(p)
+	if err != nil {
+		return res, true, json.ErrInvalidJSONPath.GenWithStackByArgs(p)
+	}
+	if pathExpr.ContainsAnyAsterisk() {
+		return res, true, json.ErrInvalidJSONPathWildcard.GenWithStackByArgs(p)
+	}
+
+	obj, exists := res.Extract([]json.PathExpression{pathExpr})
+	if !exists {
+		// If path not exists, just do nothing and no errors.
+		return res, false, nil
+	}
+
+	if obj.TypeCode != json.TypeCodeArray {
+		// res.Extract will return a json object instead of an array if there is an object at path pathExpr.
+		// JSON_ARRAY_APPEND({"a": "b"}, "$", {"b": "c"}) => [{"a": "b"}, {"b", "c"}]
+		// We should wrap them to a single array first.
+		obj = json.CreateBinary([]interface{}{obj})
+	}
+
+	obj = json.MergeBinary([]json.BinaryJSON{obj, v})
+	res, err = res.Modify([]json.PathExpression{pathExpr}, []json.BinaryJSON{obj}, json.ModifySet)
+	return res, false, err
+}
+
 type jsonArrayInsertFunctionClass struct {
 	baseFunctionClass
 }

diff --git a/expression/builtin_json_vec.go b/expression/builtin_json_vec.go
@@ -967,11 +967,87 @@ func (b *builtinJSONContainsPathSig) vecEvalInt(input *chunk.Chunk, result *chun
 }
 
 func (b *builtinJSONArrayAppendSig) vectorized() bool {
-	return false
+	return true
 }
 
 func (b *builtinJSONArrayAppendSig) vecEvalJSON(input *chunk.Chunk, result *chunk.Column) error {
-	return errors.Errorf("not implemented")
+	n := input.NumRows()
+	m := (len(b.args) - 1) / 2
+
+	jsonBufs, err := b.bufAllocator.get(types.ETJson, n)
+	if err != nil {
+		return err
+	}
+	defer b.bufAllocator.put(jsonBufs)
+	if err := b.args[0].VecEvalJSON(b.ctx, input, jsonBufs); err != nil {
+		return err
+	}
+
+	pathBufs := make([]*chunk.Column, 0, m)
+	valBufs := make([]*chunk.Column, 0, m)
+	defer func() {
+		for _, buf := range pathBufs {
+			b.bufAllocator.put(buf)
+		}
+		for _, buf := range valBufs {
+			b.bufAllocator.put(buf)
+		}
+	}()
+	for i := 1; i < len(b.args)-1; i += 2 {
+		pathBuf, err := b.bufAllocator.get(types.ETString, n)
+		if err != nil {
+			return err
+		}
+		pathBufs = append(pathBufs, pathBuf)
+		if err := b.args[i].VecEvalString(b.ctx, input, pathBuf); err != nil {
+			return err
+		}
+		valBuf, err := b.bufAllocator.get(types.ETJson, n)
+		if err != nil {
+			return err
+		}
+		if err := b.args[i+1].VecEvalJSON(b.ctx, input, valBuf); err != nil {
+			return err
+		}
+		valBufs = append(valBufs, valBuf)
+	}
+
+	result.ReserveJSON(n)
+	for i := 0; i < n; i++ {
+		if jsonBufs.IsNull(i) {
+			result.AppendNull()
+			continue
+		}
+		res := jsonBufs.GetJSON(i)
+		isNull := false
+		for j := 0; j < m; j++ {
+			if pathBufs[j].IsNull(i) {
+				isNull = true
+				break
+			}
+			s := pathBufs[j].GetString(i)
+			v, vNull := json.BinaryJSON{}, valBufs[j].IsNull(i)
+			if !vNull {
+				v = valBufs[j].GetJSON(i)
+			}
+			if vNull {
+				v = json.CreateBinary(nil)
+			}
+			res, isNull, err = b.appendJSONArray(res, s, v)
+			if err != nil {
+				return err
+			}
+			if isNull {
+				break
+			}
+		}
+		if isNull {
+			result.AppendNull()
+		} else {
+			result.AppendJSON(res)
+		}
+	}
+	return nil
 }
 
 func (b *builtinJSONUnquoteSig) vectorized() bool {

diff --git a/expression/builtin_json_vec_test.go b/expression/builtin_json_vec_test.go
@@ -26,7 +26,26 @@ var vecBuiltinJSONCases = map[string][]vecExprBenchCase{
 		{retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson}},
 		{retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson, types.ETString}, geners: []dataGenerator{&constJSONGener{"{\"a\": {\"c\": 3}, \"b\": 2}"}, &constStrGener{"$.a"}}},
 	},
-	ast.JSONArrayAppend: {},
+	ast.JSONArrayAppend: {
+		{retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson, types.ETString, types.ETJson},
+			geners: []dataGenerator{newNullWrappedGener(0.1, &constJSONGener{"{\"a\": {\"c\": 3}, \"b\": 2}"}),
+				newNullWrappedGener(0.1, &constStrGener{"$.a"}),
+				newNullWrappedGener(0.1, &constJSONGener{"{\"b\": 2}"})}},
+		{retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson, types.ETString, types.ETJson, types.ETString, types.ETJson},
+			geners: []dataGenerator{newNullWrappedGener(0.1, &constJSONGener{"{\"a\": {\"c\": 3}, \"b\": 2}"}),
+				newNullWrappedGener(0.1, &constStrGener{"$.a"}),
+				newNullWrappedGener(0.1, &constJSONGener{"{\"b\": 2}"}),
+				newNullWrappedGener(0.1, &constStrGener{"$.b"}),
+				newNullWrappedGener(0.1, &constJSONGener{"{\"x\": 3}"}),
+			}},
+		{retEvalType: types.ETJson, childrenTypes: []types.EvalType{types.ETJson, types.ETString, types.ETJson, types.ETString, types.ETJson},
+			geners: []dataGenerator{newNullWrappedGener(0.1, &constJSONGener{"{\"a\": {\"c\": 3}, \"b\": 2}"}),
+				newNullWrappedGener(0.1, &constStrGener{"$.a"}),
+				newNullWrappedGener(0.1, &constJSONGener{"{\"b\": 2}"}),
+				newNullWrappedGener(0.1, &constStrGener{"$.x"}), // not exists
+				newNullWrappedGener(0.1, &constJSONGener{"{\"x\": 3}"}),
+			}},
+	},
 	ast.JSONContainsPath: {
 		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETJson, types.ETString, types.ETString}, geners: []dataGenerator{&constJSONGener{"{\"a\": {\"c\": {\"d\": 4}}, \"b\": 2}"}, &constStrGener{"one"}, &constStrGener{"$.c"}}},
 		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETJson, types.ETString, types.ETString}, geners: []dataGenerator{&constJSONGener{"{\"a\": {\"c\": {\"d\": 4}}, \"b\": 2}"}, &constStrGener{"all"}, &constStrGener{"$.a"}, &constStrGener{"$.c"}}},