expression: support member_of pushdown to tikv (#46308)

close #46307, close #46316

DEPS.bzl

@@ -5780,13 +5780,13 @@ def go_deps():
         name = "com_github_pingcap_tipb",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/pingcap/tipb",
-        sha256 = "74b710db49c16c0400c44d14e7d2e23257d094d7cf39c71e5715bb5ac59fd57a",
-        strip_prefix = "github.com/pingcap/[email protected]20230802082933-256c189cd860",
+        sha256 = "d6e086c68505edca6fba1a6264dadbab29abc85722d520355c485da901dfdb41",
+        strip_prefix = "github.com/pingcap/[email protected]20230822064221-711da6fede03",
         urls = [
-            "http://bazel-cache.pingcap.net:8080/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20230802082933-256c189cd860.zip",
-            "http://ats.apps.svc/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20230802082933-256c189cd860.zip",
-            "https://cache.hawkingrei.com/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20230802082933-256c189cd860.zip",
-            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20230802082933-256c189cd860.zip",
+            "http://bazel-cache.pingcap.net:8080/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20230822064221-711da6fede03.zip",
+            "http://ats.apps.svc/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20230822064221-711da6fede03.zip",
+            "https://cache.hawkingrei.com/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20230822064221-711da6fede03.zip",
+            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20230822064221-711da6fede03.zip",
         ],
     )
     go_repository(

expression/builtin_json.go

@@ -779,6 +779,7 @@ func (c *jsonMemberOfFunctionClass) getFunction(ctx sessionctx.Context, args []E
 	}
 	DisableParseJSONFlag4Expr(args[0])
 	sig := &builtinJSONMemberOfSig{bf}
+	sig.setPbCode(tipb.ScalarFuncSig_JsonMemberOfSig)
 	return sig, nil
 }
 

expression/distsql_builtin.go

@@ -728,6 +728,8 @@ func getSignatureByPB(ctx sessionctx.Context, sigCode tipb.ScalarFuncSig, tp *ti
 		f = &builtinJSONValidStringSig{base}
 	case tipb.ScalarFuncSig_JsonValidOthersSig:
 		f = &builtinJSONValidOthersSig{base}
+	case tipb.ScalarFuncSig_JsonMemberOfSig:
+		f = &builtinJSONMemberOfSig{base}
 	case tipb.ScalarFuncSig_DateFormatSig:
 		f = &builtinDateFormatSig{base}
 	// case tipb.ScalarFuncSig_DateLiteral:

expression/expression.go

@@ -1075,7 +1075,7 @@ func scalarExprSupportedByTiKV(sf *ScalarFunction) bool {
 		// json functions.
 		ast.JSONType, ast.JSONExtract, ast.JSONObject, ast.JSONArray, ast.JSONMerge, ast.JSONSet,
 		ast.JSONInsert /*ast.JSONReplace,*/, ast.JSONRemove, ast.JSONLength,
-		ast.JSONUnquote, ast.JSONContains, ast.JSONValid,
+		ast.JSONUnquote, ast.JSONContains, ast.JSONValid, ast.JSONMemberOf,
 
 		// date functions.
 		ast.Date, ast.Week /* ast.YearWeek, ast.ToSeconds */, ast.DateDiff,

go.mod

@@ -82,7 +82,7 @@ require (
 	github.com/pingcap/log v1.1.1-0.20230317032135-a0d097d16e22
 	github.com/pingcap/sysutil v1.0.1-0.20230407040306-fb007c5aff21
 	github.com/pingcap/tidb/parser v0.0.0-20211011031125-9b13dc409c5e
-	github.com/pingcap/tipb v0.0.0-20230802082933-256c189cd860
+	github.com/pingcap/tipb v0.0.0-20230822064221-711da6fede03
 	github.com/pkg/errors v0.9.1
 	github.com/prometheus/client_golang v1.16.0
 	github.com/prometheus/client_model v0.4.0

go.sum

@@ -828,8 +828,8 @@ github.com/pingcap/log v1.1.1-0.20230317032135-a0d097d16e22 h1:2SOzvGvE8beiC1Y4g
 github.com/pingcap/log v1.1.1-0.20230317032135-a0d097d16e22/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pingcap/sysutil v1.0.1-0.20230407040306-fb007c5aff21 h1:QV6jqlfOkh8hqvEAgwBZa+4bSgO0EeKC7s5c6Luam2I=
 github.com/pingcap/sysutil v1.0.1-0.20230407040306-fb007c5aff21/go.mod h1:QYnjfA95ZaMefyl1NO8oPtKeb8pYUdnDVhQgf+qdpjM=
-github.com/pingcap/tipb v0.0.0-20230802082933-256c189cd860 h1:2DXvgx8MGCJitU7/IaUC+2onDAuqB7B/KMyxxiX3Y5I=
-github.com/pingcap/tipb v0.0.0-20230802082933-256c189cd860/go.mod h1:A7mrd7WHBl1o63LE2bIBGEJMTNWXqhgmYiOvMLxozfs=
+github.com/pingcap/tipb v0.0.0-20230822064221-711da6fede03 h1:puhZJ86ZXvSCZLTshL6xBEhvF963MAoL+/Jydh2tvO8=
+github.com/pingcap/tipb v0.0.0-20230822064221-711da6fede03/go.mod h1:A7mrd7WHBl1o63LE2bIBGEJMTNWXqhgmYiOvMLxozfs=
 github.com/pkg/browser v0.0.0-20210115035449-ce105d075bb4 h1:Qj1ukM4GlMWXNdMBuXcXfz/Kw9s1qm0CLY32QxuSImI=
 github.com/pkg/browser v0.0.0-20210115035449-ce105d075bb4/go.mod h1:N6UoU20jOqggOuDwUaBQpluzLNDqif3kq9z2wpdYEfQ=
 github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=

planner/core/indexmerge_path.go

@@ -735,6 +735,12 @@ func (ds *DataSource) buildPartialPaths4MVIndex(accessFilters []expression.Expre
 		return nil, false, false, nil
 	}
 
+	for _, v := range virColVals {
+		if !isSafeTypeConversion4MVIndexRange(v.GetType(), virCol.GetType()) {
+			return nil, false, false, nil
+		}
+	}
+
 	for _, v := range virColVals {
 		// rewrite json functions to EQ to calculate range, `(1 member of j)` -> `j=1`.
 		eq, err := expression.NewFunction(ds.SCtx(), ast.EQ, types.NewFieldType(mysql.TypeTiny), virCol, v)
@@ -752,6 +758,14 @@ func (ds *DataSource) buildPartialPaths4MVIndex(accessFilters []expression.Expre
 	return partialPaths, isIntersection, true, nil
 }
 
+// isSafeTypeConversion4MVIndexRange checks whether it is safe to convert valType to mvIndexType when building ranges for MVIndexes.
+func isSafeTypeConversion4MVIndexRange(valType, mvIndexType *types.FieldType) (safe bool) {
+	// for safety, forbid type conversion when building ranges for MVIndexes.
+	// TODO: loose this restriction.
+	// for example, converting '1' to 1 to access INT MVIndex may cause some wrong result.
+	return valType.EvalType() == mvIndexType.EvalType()
+}
+
 // buildPartialPath4MVIndex builds a partial path on this MVIndex with these accessFilters.
 func (ds *DataSource) buildPartialPath4MVIndex(accessFilters []expression.Expression, idxCols []*expression.Column, mvIndex *model.IndexInfo) (*util.AccessPath, bool, error) {
 	partialPath := &util.AccessPath{Index: mvIndex}

planner/core/indexmerge_path_test.go

@@ -259,27 +259,26 @@ func TestMVIndexInvisible(t *testing.T) {
 
 	tk.MustExec(`create table t(a int, j json, index kj((cast(j as signed array))))`)
 	tk.MustQuery(`explain format='brief' select /*+ use_index(t, kj) */ * from t where (1 member of (j))`).Check(testkit.Rows(
-		`Selection 8000.00 root  json_memberof(cast(1, json BINARY), test.t.j)`,
-		`└─IndexMerge 10.00 root  type: union`,
-		"  ├─IndexRangeScan(Build) 10.00 cop[tikv] table:t, index:kj(cast(`j` as signed array)) range:[1,1], keep order:false, stats:pseudo",
-		`  └─TableRowIDScan(Probe) 10.00 cop[tikv] table:t keep order:false, stats:pseudo`))
+		"IndexMerge 10.00 root  type: union",
+		"├─IndexRangeScan(Build) 10.00 cop[tikv] table:t, index:kj(cast(`j` as signed array)) range:[1,1], keep order:false, stats:pseudo",
+		"└─TableRowIDScan(Probe) 10.00 cop[tikv] table:t keep order:false, stats:pseudo",
+	))
 
 	tk.MustExec(`ALTER TABLE t ALTER INDEX kj INVISIBLE`)
 	tk.MustQuery(`explain format='brief' select /*+ use_index(t, kj) */ * from t where (1 member of (j))`).Check(testkit.Rows(
-		"Selection 8000.00 root  json_memberof(cast(1, json BINARY), test.t.j)",
-		"└─TableReader 10000.00 root  data:TableFullScan",
+		"TableReader 8000.00 root  data:Selection",
+		"└─Selection 8000.00 cop[tikv]  json_memberof(cast(1, json BINARY), test.t.j)",
 		"  └─TableFullScan 10000.00 cop[tikv] table:t keep order:false, stats:pseudo"))
 	tk.MustQuery(`explain format='brief' select /*+ use_index_merge(t, kj) */ * from t where (1 member of (j))`).Check(testkit.Rows(
-		"Selection 8000.00 root  json_memberof(cast(1, json BINARY), test.t.j)",
-		"└─TableReader 10000.00 root  data:TableFullScan",
+		"TableReader 8000.00 root  data:Selection",
+		"└─Selection 8000.00 cop[tikv]  json_memberof(cast(1, json BINARY), test.t.j)",
 		"  └─TableFullScan 10000.00 cop[tikv] table:t keep order:false, stats:pseudo"))
 
 	tk.MustExec(`ALTER TABLE t ALTER INDEX kj VISIBLE`)
 	tk.MustQuery(`explain format='brief' select /*+ use_index(t, kj) */ * from t where (1 member of (j))`).Check(testkit.Rows(
-		`Selection 8000.00 root  json_memberof(cast(1, json BINARY), test.t.j)`,
-		`└─IndexMerge 10.00 root  type: union`,
-		"  ├─IndexRangeScan(Build) 10.00 cop[tikv] table:t, index:kj(cast(`j` as signed array)) range:[1,1], keep order:false, stats:pseudo",
-		`  └─TableRowIDScan(Probe) 10.00 cop[tikv] table:t keep order:false, stats:pseudo`))
+		`IndexMerge 10.00 root  type: union`,
+		"├─IndexRangeScan(Build) 10.00 cop[tikv] table:t, index:kj(cast(`j` as signed array)) range:[1,1], keep order:false, stats:pseudo",
+		`└─TableRowIDScan(Probe) 10.00 cop[tikv] table:t keep order:false, stats:pseudo`))
 }
 
 func TestMVIndexFullScan(t *testing.T) {
@@ -425,3 +424,17 @@ func randMVIndexValue(opts randMVIndexValOpts) string {
 	}
 	return ""
 }
+
+func TestIndexMergeJSONMemberOf2(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec(`create table t(
+a int, j0 json, j1 json,
+index j0_0((cast(j0->'$.path0' as signed array))));`)
+	tk.MustExec("insert into t values(1, '{\"path0\" : [1,2,3]}', null ); ")
+	tk.MustQuery("select /*+ no_index_merge() */ a from t where (1 member of (j0->'$.path0')); ").Check(testkit.Rows("1"))
+	tk.MustQuery("select /*+ no_index_merge() */ a from t where ('1' member of (j0->'$.path0')); ").Check(testkit.Rows())
+	tk.MustQuery("select /*+ use_index_merge(t, j0_0) */ a from t where (1 member of (j0->'$.path0')); ").Check(testkit.Rows("1"))
+	tk.MustQuery("select /*+ use_index_merge(t, j0_0) */ a from t where ('1' member of (j0->'$.path0')); ").Check(testkit.Rows())
+}