expression: implement vectorized evaluation for `builtinJSONContainsS…

…ig` (pingcap#12924)
cqlxcnp · Oct 29, 2019 · f3958c1 · f3958c1
1 parent f1f80b4
commit f3958c1
Show file tree

Hide file tree

Showing 2 changed files with 84 additions and 3 deletions.
diff --git a/expression/builtin_json_vec.go b/expression/builtin_json_vec.go
@@ -105,11 +105,88 @@ func (b *builtinJSONArraySig) vecEvalJSON(input *chunk.Chunk, result *chunk.Colu
 }
 
 func (b *builtinJSONContainsSig) vectorized() bool {
-	return false
+	return true
 }
 
 func (b *builtinJSONContainsSig) vecEvalInt(input *chunk.Chunk, result *chunk.Column) error {
-	return errors.Errorf("not implemented")
+	nr := input.NumRows()
+
+	objCol, err := b.bufAllocator.get(types.ETJson, nr)
+	if err != nil {
+		return err
+	}
+	defer b.bufAllocator.put(objCol)
+
+	if err := b.args[0].VecEvalJSON(b.ctx, input, objCol); err != nil {
+		return err
+	}
+
+	targetCol, err := b.bufAllocator.get(types.ETJson, nr)
+	if err != nil {
+		return err
+	}
+	defer b.bufAllocator.put(targetCol)
+
+	if err := b.args[1].VecEvalJSON(b.ctx, input, targetCol); err != nil {
+		return err
+	}
+
+	result.ResizeInt64(nr, false)
+	resI64s := result.Int64s()
+
+	if len(b.args) == 3 {
+		pathCol, err := b.bufAllocator.get(types.ETString, nr)
+		if err != nil {
+			return err
+		}
+		defer b.bufAllocator.put(pathCol)
+
+		if err := b.args[2].VecEvalString(b.ctx, input, pathCol); err != nil {
+			return err
+		}
+
+		result.MergeNulls(objCol, targetCol, pathCol)
+
+		var pathExpr json.PathExpression
+		for i := 0; i < nr; i++ {
+			if result.IsNull(i) {
+				continue
+			}
+			pathExpr, err = json.ParseJSONPathExpr(pathCol.GetString(i))
+			if err != nil {
+				return err
+			}
+			if pathExpr.ContainsAnyAsterisk() {
+				return json.ErrInvalidJSONPathWildcard
+			}
+
+			obj, exists := objCol.GetJSON(i).Extract([]json.PathExpression{pathExpr})
+			if !exists {
+				result.SetNull(i, true)
+				continue
+			}
+
+			if json.ContainsBinary(obj, targetCol.GetJSON(i)) {
+				resI64s[i] = 1
+			} else {
+				resI64s[i] = 0
+			}
+		}
+	} else {
+		result.MergeNulls(objCol, targetCol)
+		for i := 0; i < nr; i++ {
+			if result.IsNull(i) {
+				continue
+			}
+			if json.ContainsBinary(objCol.GetJSON(i), targetCol.GetJSON(i)) {
+				resI64s[i] = 1
+			} else {
+				resI64s[i] = 0
+			}
+		}
+	}
+
+	return nil
 }
 
 func (b *builtinJSONQuoteSig) vectorized() bool {

diff --git a/expression/builtin_json_vec_test.go b/expression/builtin_json_vec_test.go
@@ -32,7 +32,11 @@ var vecBuiltinJSONCases = map[string][]vecExprBenchCase{
 	ast.JSONType:         {{retEvalType: types.ETString, childrenTypes: []types.EvalType{types.ETJson}}},
 	ast.JSONArray:        {},
 	ast.JSONArrayInsert:  {},
-	ast.JSONContains:     {},
+	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"}}},
+		{retEvalType: types.ETInt, childrenTypes: []types.EvalType{types.ETJson, types.ETJson, types.ETString}, geners: []dataGenerator{nil, nil, &constStrGener{"$.key"}}},
+	},
 	ast.JSONObject: {
 		{
 			retEvalType: types.ETJson,