planner: support more types to use IndexMerge to access MVIndex (#40343)

ref #40191
pingcap · Jan 6, 2023 · 99100d4 · 99100d4
1 parent 3291b6d
commit 99100d4
Show file tree

Hide file tree

Showing 6 changed files with 331 additions and 39 deletions.
diff --git a/expression/builtin_cast.go b/expression/builtin_cast.go
@@ -487,8 +487,17 @@ func (b *castJSONAsArrayFunctionSig) evalJSON(row chunk.Row) (res types.BinaryJS
 	return types.CreateBinaryJSON(arrayVals), false, nil
 }
 
-func convertJSON2Tp(eval types.EvalType) func(*stmtctx.StatementContext, types.BinaryJSON, *types.FieldType) (any, error) {
-	switch eval {
+// ConvertJSON2Tp returns a function that can convert JSON to the specified type.
+func ConvertJSON2Tp(v types.BinaryJSON, targetType *types.FieldType) (any, error) {
+	convertFunc := convertJSON2Tp(targetType.EvalType())
+	if convertFunc == nil {
+		return nil, ErrInvalidJSONForFuncIndex
+	}
+	return convertFunc(fakeSctx, v, targetType)
+}
+
+func convertJSON2Tp(evalType types.EvalType) func(*stmtctx.StatementContext, types.BinaryJSON, *types.FieldType) (any, error) {
+	switch evalType {
 	case types.ETString:
 		return func(sc *stmtctx.StatementContext, item types.BinaryJSON, tp *types.FieldType) (any, error) {
 			if item.TypeCode != types.JSONTypeCodeString {

diff --git a/expression/util.go b/expression/util.go
@@ -1145,6 +1145,33 @@ func IsMutableEffectsExpr(expr Expression) bool {
 	return false
 }
 
+// IsInmutableExpr checks whether this expression only consists of foldable functions and inmutable constants.
+// This expression can be evaluated by using `expr.Eval(chunk.Row{})` directly if it's inmutable.
+func IsInmutableExpr(expr Expression) bool {
+	switch x := expr.(type) {
+	case *ScalarFunction:
+		if _, ok := unFoldableFunctions[x.FuncName.L]; ok {
+			return false
+		}
+		if _, ok := mutableEffectsFunctions[x.FuncName.L]; ok {
+			return false
+		}
+		for _, arg := range x.GetArgs() {
+			if !IsInmutableExpr(arg) {
+				return false
+			}
+		}
+		return true
+	case *Constant:
+		if x.DeferredExpr != nil || x.ParamMarker != nil {
+			return false
+		}
+		return true
+	default:
+		return false
+	}
+}
+
 // RemoveDupExprs removes identical exprs. Not that if expr contains functions which
 // are mutable or have side effects, we cannot remove it even if it has duplicates;
 // if the plan is going to be cached, we cannot remove expressions containing `?` neither.

diff --git a/planner/core/indexmerge_path.go b/planner/core/indexmerge_path.go
@@ -563,7 +563,8 @@ func (ds *DataSource) generateIndexMergeJSONMVIndexPath(normalPathCnt int, filte
 				indexMergeIsIntersection = true
 				jsonPath = sf.GetArgs()[0]
 				var ok bool
-				vals, ok = jsonArrayExpr2Exprs(ds.ctx, sf.GetArgs()[1])
+				//virCol.RetType
+				vals, ok = jsonArrayExpr2Exprs(ds.ctx, sf.GetArgs()[1], virCol.GetType())
 				if !ok {
 					continue
 				}
@@ -578,7 +579,7 @@ func (ds *DataSource) generateIndexMergeJSONMVIndexPath(normalPathCnt int, filte
 				}
 				jsonPath = sf.GetArgs()[jsonPathIdx]
 				var ok bool
-				vals, ok = jsonArrayExpr2Exprs(ds.ctx, sf.GetArgs()[1-jsonPathIdx])
+				vals, ok = jsonArrayExpr2Exprs(ds.ctx, sf.GetArgs()[1-jsonPathIdx], virCol.GetType())
 				if !ok {
 					continue
 				}
@@ -593,21 +594,6 @@ func (ds *DataSource) generateIndexMergeJSONMVIndexPath(normalPathCnt int, filte
 			if !jsonPath.Equal(ds.ctx, targetJSONPath) {
 				continue // not on the same JSON col
 			}
-			// only support INT now
-			// TODO: support more types
-			if jsonPath.GetType().EvalType() == types.ETInt {
-				continue
-			}
-			allInt := true
-			// TODO: support using IndexLookUp to handle single-value cases.
-			for _, v := range vals {
-				if v.GetType().EvalType() != types.ETInt {
-					allInt = false
-				}
-			}
-			if !allInt {
-				continue
-			}
 
 			// Step 2.3. Generate a IndexMerge Path of this filter on the current MVIndex.
 			var partialPaths []*util.AccessPath
@@ -652,16 +638,8 @@ func (ds *DataSource) generateIndexMergeJSONMVIndexPath(normalPathCnt int, filte
 }
 
 // jsonArrayExpr2Exprs converts a JsonArray expression to expression list: cast('[1, 2, 3]' as JSON) --> []expr{1, 2, 3}
-func jsonArrayExpr2Exprs(sctx sessionctx.Context, jsonArrayExpr expression.Expression) ([]expression.Expression, bool) {
-	// only support cast(const as JSON)
-	arrayExpr, wrappedByJSONCast := unwrapJSONCast(jsonArrayExpr)
-	if !wrappedByJSONCast {
-		return nil, false
-	}
-	if _, isConst := arrayExpr.(*expression.Constant); !isConst {
-		return nil, false
-	}
-	if expression.IsMutableEffectsExpr(arrayExpr) {
+func jsonArrayExpr2Exprs(sctx sessionctx.Context, jsonArrayExpr expression.Expression, targetType *types.FieldType) ([]expression.Expression, bool) {
+	if !expression.IsInmutableExpr(jsonArrayExpr) || jsonArrayExpr.GetType().EvalType() != types.ETJson {
 		return nil, false
 	}
 
@@ -670,15 +648,15 @@ func jsonArrayExpr2Exprs(sctx sessionctx.Context, jsonArrayExpr expression.Expre
 		return nil, false
 	}
 	if jsonArray.TypeCode != types.JSONTypeCodeArray {
-		single, ok := jsonValue2Expr(jsonArray) // '1' -> []expr{1}
+		single, ok := jsonValue2Expr(jsonArray, targetType) // '1' -> []expr{1}
 		if ok {
 			return []expression.Expression{single}, true
 		}
 		return nil, false
 	}
 	var exprs []expression.Expression
 	for i := 0; i < jsonArray.GetElemCount(); i++ { // '[1, 2, 3]' -> []expr{1, 2, 3}
-		expr, ok := jsonValue2Expr(jsonArray.ArrayGetElem(i))
+		expr, ok := jsonValue2Expr(jsonArray.ArrayGetElem(i), targetType)
 		if !ok {
 			return nil, false
 		}
@@ -687,16 +665,14 @@ func jsonArrayExpr2Exprs(sctx sessionctx.Context, jsonArrayExpr expression.Expre
 	return exprs, true
 }
 
-func jsonValue2Expr(v types.BinaryJSON) (expression.Expression, bool) {
-	if v.TypeCode != types.JSONTypeCodeInt64 {
-		// only support INT now
-		// TODO: support more types
+func jsonValue2Expr(v types.BinaryJSON, targetType *types.FieldType) (expression.Expression, bool) {
+	datum, err := expression.ConvertJSON2Tp(v, targetType)
+	if err != nil {
 		return nil, false
 	}
-	val := v.GetInt64()
 	return &expression.Constant{
-		Value:   types.NewDatum(val),
-		RetType: types.NewFieldType(mysql.TypeLonglong),
+		Value:   types.NewDatum(datum),
+		RetType: targetType,
 	}, true
 }
 

diff --git a/planner/core/indexmerge_path_test.go b/planner/core/indexmerge_path_test.go
@@ -30,6 +30,10 @@ func TestIndexMergeJSONMemberOf(t *testing.T) {
 a int, j0 json, j1 json,
 index j0_0((cast(j0->'$.path0' as signed array))),
 index j0_1((cast(j0->'$.path1' as signed array))),
+index j0_double((cast(j0->'$.path_double' as double array))),
+index j0_decimal((cast(j0->'$.path_decimal' as decimal(10, 2) array))),
+index j0_string((cast(j0->'$.path_string' as char(10) array))),
+index j0_date((cast(j0->'$.path_date' as date array))),
 index j1((cast(j1 as signed array))))`)
 
 	var input []string

diff --git a/planner/core/testdata/index_merge_suite_in.json b/planner/core/testdata/index_merge_suite_in.json
@@ -18,7 +18,31 @@
       "select /*+ use_index_merge(t, j0_0) */ * from t where json_overlaps('1', (j0->'$.path0'))",
       "select /*+ use_index_merge(t, j0_0) */ * from t where json_contains((j0->'$.path0'), '1') and a<10",
       "select /*+ use_index_merge(t, j0_0) */ * from t where json_overlaps((j0->'$.path0'), '1') and a<10",
-      "select /*+ use_index_merge(t, j0_0) */ * from t where json_overlaps('1', (j0->'$.path0')) and a<10"
+      "select /*+ use_index_merge(t, j0_0) */ * from t where json_overlaps('1', (j0->'$.path0')) and a<10",
+      "select /*+ use_index_merge(t, j0_double) */ * from t where (1.5 member of (j0->'$.path_double'))",
+      "select /*+ use_index_merge(t, j0_double) */ * from t where (1.5 member of (j0->'$.path_double')) and a<10",
+      "select /*+ use_index_merge(t, j0_double) */ * from t where json_contains((j0->'$.path_double'), '[1.1, 1.2, 1.3]')",
+      "select /*+ use_index_merge(t, j0_double) */ * from t where json_contains((j0->'$.path_double'), '[1.1, 1.2, 1.3]') and a<10",
+      "select /*+ use_index_merge(t, j0_double) */ * from t where json_overlaps((j0->'$.path_double'), '[1.1, 1.2, 1.3]')",
+      "select /*+ use_index_merge(t, j0_double) */ * from t where json_overlaps((j0->'$.path_double'), '[1.1, 1.2, 1.3]') and a<10",
+      "select /*+ use_index_merge(t, j0_decimal) */ * from t where (1.5 member of (j0->'$.path_decimal'))",
+      "select /*+ use_index_merge(t, j0_decimal) */ * from t where (1.5 member of (j0->'$.path_decimal')) and a<10",
+      "select /*+ use_index_merge(t, j0_decimal) */ * from t where json_contains((j0->'$.path_decimal'), '[1.1, 1.2, 1.3]')",
+      "select /*+ use_index_merge(t, j0_decimal) */ * from t where json_contains((j0->'$.path_decimal'), '[1.1, 1.2, 1.3]') and a<10",
+      "select /*+ use_index_merge(t, j0_decimal) */ * from t where json_overlaps((j0->'$.path_decimal'), '[1.1, 1.2, 1.3]')",
+      "select /*+ use_index_merge(t, j0_decimal) */ * from t where json_overlaps((j0->'$.path_decimal'), '[1.1, 1.2, 1.3]') and a<10",
+      "select /*+ use_index_merge(t, j0_string) */ * from t where (\"a\" member of (j0->'$.path_string'))",
+      "select /*+ use_index_merge(t, j0_string) */ * from t where (\"a\" member of (j0->'$.path_string')) and a<10",
+      "select /*+ use_index_merge(t, j0_string) */ * from t where json_contains((j0->'$.path_string'), '[\"a\", \"b\", \"c\"]')",
+      "select /*+ use_index_merge(t, j0_string) */ * from t where json_contains((j0->'$.path_string'), '[\"a\", \"b\", \"c\"]') and a<10",
+      "select /*+ use_index_merge(t, j0_string) */ * from t where json_overlaps((j0->'$.path_string'), '[\"a\", \"b\", \"c\"]')",
+      "select /*+ use_index_merge(t, j0_string) */ * from t where json_overlaps((j0->'$.path_string'), '[\"a\", \"b\", \"c\"]') and a<10",
+      "select /*+ use_index_merge(t, j0_date) */ * from t where (\"2023-01-01\" member of (j0->'$.path_date'))",
+      "select /*+ use_index_merge(t, j0_date) */ * from t where (\"2023-01-01\" member of (j0->'$.path_date')) and a<10",
+      "select /*+ use_index_merge(t, j0_date) */ * from t where json_contains((j0->'$.path_date'), json_array(cast('2023-01-01' as date), cast('2023-01-02' as date), cast('2023-01-03' as date)))",
+      "select /*+ use_index_merge(t, j0_date) */ * from t where json_contains((j0->'$.path_date'), json_array(cast('2023-01-01' as date), cast('2023-01-02' as date), cast('2023-01-03' as date))) and a<10",
+      "select /*+ use_index_merge(t, j0_date) */ * from t where json_overlaps((j0->'$.path_date'), json_array(cast('2023-01-01' as date), cast('2023-01-02' as date), cast('2023-01-03' as date)))",
+      "select /*+ use_index_merge(t, j0_date) */ * from t where json_overlaps((j0->'$.path_date'), json_array(cast('2023-01-01' as date), cast('2023-01-02' as date), cast('2023-01-03' as date))) and a<10"
     ]
   },
   {