planner: support basic usage of partial index (#65051)

pkg/expression/util.go

@@ -782,6 +782,18 @@ var symmetricOp = map[opcode.Op]opcode.Op{
 	opcode.NullEQ: opcode.NullEQ,
 }
 
+// CompareOpMap records all comparison operators.
+var CompareOpMap = map[string]struct{}{
+	ast.LT:     {},
+	ast.GE:     {},
+	ast.GT:     {},
+	ast.LE:     {},
+	ast.EQ:     {},
+	ast.NE:     {},
+	ast.NullEQ: {},
+	ast.In:     {},
+}
+
 func pushNotAcrossArgs(ctx BuildContext, exprs []Expression, not bool) ([]Expression, bool) {
 	newExprs := make([]Expression, 0, len(exprs))
 	flag := false

pkg/planner/core/casetest/index/BUILD.bazel

@@ -9,18 +9,32 @@ go_test(
     ],
     data = glob(["testdata/**"]),
     flaky = True,
+<<<<<<< HEAD
     shard_count = 8,
     deps = [
         "//pkg/domain",
         "//pkg/domain/infosync",
         "//pkg/parser/model",
+=======
+    shard_count = 11,
+    deps = [
+        "//pkg/domain",
+        "//pkg/domain/infosync",
+        "//pkg/parser/ast",
+        "//pkg/planner/util",
+        "//pkg/session/sessmgr",
+>>>>>>> 959bf330874 (planner: support basic usage of partial index (#65051))
         "//pkg/store/mockstore",
         "//pkg/testkit",
         "//pkg/testkit/testdata",
         "//pkg/testkit/testfailpoint",
         "//pkg/testkit/testmain",
         "//pkg/testkit/testsetup",
+<<<<<<< HEAD
         "//pkg/util",
+=======
+        "@com_github_pingcap_failpoint//:failpoint",
+>>>>>>> 959bf330874 (planner: support basic usage of partial index (#65051))
         "@com_github_stretchr_testify//require",
         "@org_uber_go_goleak//:goleak",
     ],

pkg/planner/core/casetest/index/index_test.go

@@ -20,9 +20,16 @@ import (
 	"testing"
 	"time"
 
+	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/pkg/domain"
 	"github.com/pingcap/tidb/pkg/domain/infosync"
+<<<<<<< HEAD
 	"github.com/pingcap/tidb/pkg/parser/model"
+=======
+	"github.com/pingcap/tidb/pkg/parser/ast"
+	"github.com/pingcap/tidb/pkg/planner/util"
+	"github.com/pingcap/tidb/pkg/session/sessmgr"
+>>>>>>> 959bf330874 (planner: support basic usage of partial index (#65051))
 	"github.com/pingcap/tidb/pkg/store/mockstore"
 	"github.com/pingcap/tidb/pkg/testkit"
 	"github.com/pingcap/tidb/pkg/testkit/testdata"
@@ -346,3 +353,95 @@ func TestAnalyzeVectorIndex(t *testing.T) {
 		"Warning 1105 analyzing vector index is not supported, skip idx2",
 		"Warning 1681 ANALYZE with tidb_analyze_version=1 is deprecated and will be removed in a future release."))
 }
+
+func TestPartialIndexWithPlanCache(t *testing.T) {
+	testkit.RunTestUnderCascades(t, func(t *testing.T, tk *testkit.TestKit, cascades, caller string) {
+		tk.MustExec(`set tidb_enable_prepared_plan_cache=1`)
+		tk.MustExec("use test")
+		tk.MustExec("set @@tidb_enable_collect_execution_info=0;")
+		tk.MustExec("drop table if exists t")
+		tk.MustExec("create table t(a int, b int, index idx1(a) where a is not null, index idx2(b) where b > 10)")
+
+		tk.MustExec("prepare stmt from 'select * from t where a = ?'")
+		tk.MustExec("set @a = 123")
+
+		// IS NOT NULL pre condition can use plan cache.
+		tk.MustExec("execute stmt using @a")
+		tk.MustExec("execute stmt using @a")
+		tkProcess := tk.Session().ShowProcess()
+		ps := []*sessmgr.ProcessInfo{tkProcess}
+		tk.Session().SetSessionManager(&testkit.MockSessionManager{PS: ps})
+		tk.MustQuery(fmt.Sprintf("explain for connection %d", tkProcess.ID)).CheckContain("idx1")
+		tk.MustExec("execute stmt using @a")
+		tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("1"))
+
+		// Normal pre condition can not use plan cache.
+		tk.MustExec("prepare stmt from 'select * from t where b = ?'")
+		tk.MustExec("set @a = 20")
+		tk.MustExec("execute stmt using @a")
+		tk.MustExec("execute stmt using @a")
+		tkProcess = tk.Session().ShowProcess()
+		ps[0] = tkProcess
+		tk.Session().SetSessionManager(&testkit.MockSessionManager{PS: ps})
+		tk.MustQuery(fmt.Sprintf("explain for connection %d", tkProcess.ID)).CheckContain("idx2")
+		tk.MustExec("execute stmt using @a")
+		tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("0"))
+	})
+}
+
+func TestPartialIndexWithIndexPrune(t *testing.T) {
+	testkit.RunTestUnderCascades(t, func(t *testing.T, tk *testkit.TestKit, cascades, caller string) {
+		tk.MustExec("use test")
+		tk.MustExec("set @@tidb_enable_collect_execution_info=0;")
+		tk.MustExec("drop table if exists t")
+		tk.MustExec("create table t(a int, b int, index idx1(a) where a is not null, index idx2(b) where b > 10)")
+		tk.MustQuery("explain select * from t use index(idx1) where a > 1").CheckContain("idx1")
+
+		// Set the prune behavior to prune all non interesting ones.
+		tk.MustExec("set @@tidb_opt_index_prune_threshold=0")
+
+		// The failpoint will check whether all partial indexes are pruned.
+		fpName := "github.com/pingcap/tidb/pkg/planner/core/rule/InjectCheckForIndexPrune"
+		require.NoError(t, failpoint.EnableCall(fpName, func(paths []*util.AccessPath) {
+			for _, path := range paths {
+				if path != nil && path.Index != nil && path.Index.ConditionExprString != "" {
+					require.True(t, false, "Partial index should be pruned")
+				}
+			}
+		}))
+		tk.MustQuery("select * from t")
+
+		// idx1 is pruned because a is not referenced as interesting one.
+		// idx2 is kept though its constraint is not matched.
+		require.NoError(t, failpoint.EnableCall(fpName, func(paths []*util.AccessPath) {
+			idx2Found := false
+			for _, path := range paths {
+				if path != nil && path.Index != nil && path.Index.Name.L == "idx1" {
+					require.True(t, false, "Partial index idx1 should be pruned")
+				}
+				if path != nil && path.Index != nil && path.Index.Name.L == "idx2" {
+					idx2Found = true
+				}
+			}
+			require.True(t, idx2Found, "Partial index idx2 should not be pruned")
+		}))
+		tk.MustQuery("explain select * from t order by b").CheckNotContain("idx2")
+
+		// idx2 is pruned because b is not referenced as interesting one.
+		// idx1 is kept though its constraint is not matched.
+		require.NoError(t, failpoint.EnableCall(fpName, func(paths []*util.AccessPath) {
+			idx1Found := false
+			for _, path := range paths {
+				if path != nil && path.Index != nil && path.Index.Name.L == "idx2" {
+					require.True(t, false, "Partial index idx2 should be pruned")
+				}
+				if path != nil && path.Index != nil && path.Index.Name.L == "idx1" {
+					idx1Found = true
+				}
+			}
+			require.True(t, idx1Found, "Partial index idx1 should not be pruned")
+		}))
+		tk.MustQuery("explain select * from t where a is null").CheckNotContain("idx1")
+		require.NoError(t, failpoint.Disable(fpName))
+	})
+}

pkg/planner/core/exhaust_physical_plans.go

@@ -796,6 +796,165 @@ childLoop:
 	return wrapper
 }
 
+<<<<<<< HEAD
+=======
+// buildDataSource2IndexScanByIndexJoinProp builds an IndexScan as the inner child for an
+// IndexJoin based on IndexJoinProp included in prop if possible.
+//
+// buildDataSource2IndexScanByIndexJoinProp differs with buildIndexJoinInner2IndexScan in that
+// the first one is try to build a single table scan as the inner child of an index join then return
+// this inner task(raw table scan) bottom-up, which will be attached with other inner parents of an
+// index join in attach2Task when bottom-up of enumerating the physical plans;
+//
+// while the second is try to build a table scan as the inner child of an index join, then build
+// entire inner subtree of a index join out as innerTask instantly according those validated and
+// zipped inner patterns with calling constructInnerIndexScanTask. That's not done yet, it also
+// tries to enumerate kinds of index join operators based on the finished innerTask and un-decided
+// outer child which will be physical-ed in the future.
+func buildDataSource2IndexScanByIndexJoinProp(
+	ds *logicalop.DataSource,
+	prop *property.PhysicalProperty) base.Task {
+	indexValid := func(path *util.AccessPath) bool {
+		if path.IsTablePath() {
+			return false
+		}
+		// if path is index path. index path currently include two kind of, one is normal, and the other is mv index.
+		// for mv index like mvi(a, json, b), if driving condition is a=1, and we build a prefix scan with range [1,1]
+		// on mvi, it will return many index rows which breaks handle-unique attribute here.
+		//
+		// the basic rule is that: mv index can be and can only be accessed by indexMerge operator. (embedded handle duplication)
+		if !path.IsIndexJoinUnapplicable() {
+			return true // not a MVIndex path, it can successfully be index join probe side.
+		}
+		return false
+	}
+	indexJoinResult, keyOff2IdxOff := getBestIndexJoinPathResultByProp(ds, prop.IndexJoinProp, indexValid)
+	if indexJoinResult == nil {
+		return base.InvalidTask
+	}
+	rangeInfo, maxOneRow := indexJoinPathGetRangeInfoAndMaxOneRow(ds.SCtx(), prop.IndexJoinProp.OuterJoinKeys, indexJoinResult)
+	var innerTask base.Task
+	if !prop.IsSortItemEmpty() && matchProperty(ds, indexJoinResult.chosenPath, prop) == property.PropMatched {
+		innerTask = constructDS2IndexScanTask(ds, indexJoinResult.chosenPath, indexJoinResult.chosenRanges.Range(), indexJoinResult.chosenRemained, indexJoinResult.idxOff2KeyOff, rangeInfo, true, prop.SortItems[0].Desc, prop.IndexJoinProp.AvgInnerRowCnt, maxOneRow)
+	} else {
+		innerTask = constructDS2IndexScanTask(ds, indexJoinResult.chosenPath, indexJoinResult.chosenRanges.Range(), indexJoinResult.chosenRemained, indexJoinResult.idxOff2KeyOff, rangeInfo, false, false, prop.IndexJoinProp.AvgInnerRowCnt, maxOneRow)
+	}
+	// since there is a possibility that inner task can't be built and the returned value is nil, we just return base.InvalidTask.
+	if innerTask == nil {
+		return base.InvalidTask
+	}
+	// prepare the index path chosen information and wrap them as IndexJoinInfo and fill back to CopTask.
+	// here we don't need to construct physical index join here anymore, because we will encapsulate it bottom-up.
+	// chosenPath and lastColManager of indexJoinResult should be returned to the caller (seen by index join to keep
+	// index join aware of indexColLens and compareFilters).
+	completeIndexJoinFeedBackInfo(innerTask.(*physicalop.CopTask), indexJoinResult, indexJoinResult.chosenRanges, keyOff2IdxOff)
+	return innerTask
+}
+
+// buildDataSource2TableScanByIndexJoinProp builds a TableScan as the inner child for an
+// IndexJoin if possible.
+// If the inner side of an index join is a TableScan, only one tuple will be
+// fetched from the inner side for every tuple from the outer side. This will be
+// promised to be no worse than building IndexScan as the inner child.
+func buildDataSource2TableScanByIndexJoinProp(
+	ds *logicalop.DataSource,
+	prop *property.PhysicalProperty) base.Task {
+	var tblPath *util.AccessPath
+	for _, path := range ds.PossibleAccessPaths {
+		if path.IsTablePath() && path.StoreType == kv.TiKV { // old logic
+			tblPath = path
+			break
+		}
+	}
+	if tblPath == nil {
+		return base.InvalidTask
+	}
+	var keyOff2IdxOff []int
+	var ranges ranger.MutableRanges = ranger.Ranges{}
+	var innerTask base.Task
+	var indexJoinResult *indexJoinPathResult
+	if ds.TableInfo.IsCommonHandle {
+		// for the leaf datasource, we use old logic to get the indexJoinResult, which contain the chosen path and ranges.
+		indexJoinResult, keyOff2IdxOff = getBestIndexJoinPathResultByProp(ds, prop.IndexJoinProp, func(path *util.AccessPath) bool { return path.IsCommonHandlePath })
+		// if there is no chosen info, it means the leaf datasource couldn't even leverage this indexJoinProp, return InvalidTask.
+		if indexJoinResult == nil {
+			return base.InvalidTask
+		}
+		// prepare the range info with outer join keys, it shows like: [xxx] decided by:
+		rangeInfo, maxOneRow := indexJoinPathGetRangeInfoAndMaxOneRow(ds.SCtx(), prop.IndexJoinProp.OuterJoinKeys, indexJoinResult)
+		// construct the inner task with chosen path and ranges, note: it only for this leaf datasource.
+		// like the normal way, we need to check whether the chosen path is matched with the prop, if so, we will set the `keepOrder` to true.
+		if matchProperty(ds, indexJoinResult.chosenPath, prop) == property.PropMatched {
+			innerTask = constructDS2TableScanTask(ds, indexJoinResult.chosenRanges.Range(), rangeInfo, true, !prop.IsSortItemEmpty() && prop.SortItems[0].Desc, prop.IndexJoinProp.AvgInnerRowCnt, maxOneRow)
+		} else {
+			innerTask = constructDS2TableScanTask(ds, indexJoinResult.chosenRanges.Range(), rangeInfo, false, false, prop.IndexJoinProp.AvgInnerRowCnt, maxOneRow)
+		}
+		ranges = indexJoinResult.chosenRanges
+	} else {
+		var (
+			ok               bool
+			chosenPath       *util.AccessPath
+			newOuterJoinKeys []*expression.Column
+			// note: pk col doesn't have mutableRanges, the global var(ranges) which will be handled as empty range in constructIndexJoin.
+			localRanges ranger.Ranges
+		)
+		keyOff2IdxOff, newOuterJoinKeys, localRanges, chosenPath, ok = getIndexJoinIntPKPathInfo(ds, prop.IndexJoinProp.InnerJoinKeys, prop.IndexJoinProp.OuterJoinKeys, func(path *util.AccessPath) bool { return path.IsIntHandlePath })
+		if !ok {
+			return base.InvalidTask
+		}
+		// For IntHandle (integer primary key), it's always a unique match.
+		maxOneRow := true
+		rangeInfo := indexJoinIntPKRangeInfo(ds.SCtx().GetExprCtx().GetEvalCtx(), newOuterJoinKeys)
+		if !prop.IsSortItemEmpty() && matchProperty(ds, chosenPath, prop) == property.PropMatched {
+			innerTask = constructDS2TableScanTask(ds, localRanges, rangeInfo, true, prop.SortItems[0].Desc, prop.IndexJoinProp.AvgInnerRowCnt, maxOneRow)
+		} else {
+			innerTask = constructDS2TableScanTask(ds, localRanges, rangeInfo, false, false, prop.IndexJoinProp.AvgInnerRowCnt, maxOneRow)
+		}
+	}
+	// since there is a possibility that inner task can't be built and the returned value is nil, we just return base.InvalidTask.
+	if innerTask == nil {
+		return base.InvalidTask
+	}
+	// prepare the index path chosen information and wrap them as IndexJoinInfo and fill back to CopTask.
+	// here we don't need to construct physical index join here anymore, because we will encapsulate it bottom-up.
+	// chosenPath and lastColManager of indexJoinResult should be returned to the caller (seen by index join to keep
+	// index join aware of indexColLens and compareFilters).
+	completeIndexJoinFeedBackInfo(innerTask.(*physicalop.CopTask), indexJoinResult, ranges, keyOff2IdxOff)
+	return innerTask
+}
+
+// completeIndexJoinFeedBackInfo completes the IndexJoinInfo for the innerTask.
+// indexJoin
+//
+//	+--- outer child
+//	+--- inner child (say: projection ------------> unionScan -------------> ds)
+//	        <-------RootTask(IndexJoinInfo) <--RootTask(IndexJoinInfo) <--copTask(IndexJoinInfo)
+//
+// when we build the underlying datasource as table-scan, we will return wrap it and
+// return as a CopTask, inside which the index join contains some index path chosen
+// information which will be used in indexJoin execution runtime: ref IndexJoinInfo
+// declaration for more information.
+// the indexJoinInfo will be filled back to the innerTask, passed upward to RootTask
+// once this copTask is converted to RootTask type, and finally end up usage in the
+// indexJoin's attach2Task with calling completePhysicalIndexJoin.
+func completeIndexJoinFeedBackInfo(innerTask *physicalop.CopTask, indexJoinResult *indexJoinPathResult, ranges ranger.MutableRanges, keyOff2IdxOff []int) {
+	info := innerTask.IndexJoinInfo
+	if info == nil {
+		info = &physicalop.IndexJoinInfo{}
+	}
+	if indexJoinResult != nil {
+		if indexJoinResult.chosenPath != nil {
+			info.IdxColLens = indexJoinResult.chosenPath.IdxColLens
+		}
+		info.CompareFilters = indexJoinResult.lastColManager
+	}
+	info.Ranges = ranges
+	info.KeyOff2IdxOff = keyOff2IdxOff
+	// fill it back to the bottom-up Task.
+	innerTask.IndexJoinInfo = info
+}
+
+>>>>>>> 959bf330874 (planner: support basic usage of partial index (#65051))
 // buildIndexJoinInner2TableScan builds a TableScan as the inner child for an
 // IndexJoin if possible.
 // If the inner side of a index join is a TableScan, only one tuple will be
@@ -915,7 +1074,7 @@ func buildIndexJoinInner2IndexScan(
 		// on mvi, it will return many index rows which breaks handle-unique attribute here.
 		//
 		// the basic rule is that: mv index can be and can only be accessed by indexMerge operator. (embedded handle duplication)
-		if !isMVIndexPath(path) {
+		if !path.IsIndexJoinUnapplicable() {
 			return true // not a MVIndex path, it can successfully be index join probe side.
 		}
 		return false

pkg/planner/core/indexmerge_path.go

@@ -135,11 +135,21 @@ func generateIndexMergePath(ds *logicalop.DataSource) error {
 
 func generateNormalIndexPartialPaths4DNF(
 	ds *logicalop.DataSource,
+<<<<<<< HEAD
 	dnfItems []expression.Expression,
 	candidatePaths []*util.AccessPath,
 ) (paths []*util.AccessPath, needSelection bool, usedMap []bool) {
 	paths = make([]*util.AccessPath, 0, len(dnfItems))
 	usedMap = make([]bool, len(dnfItems))
+=======
+	item expression.Expression,
+	candidatePath *util.AccessPath,
+) (paths *util.AccessPath, needSelection bool) {
+	// Reject partial index first.
+	if candidatePath.Index != nil && candidatePath.Index.HasCondition() {
+		return nil, false
+	}
+>>>>>>> 959bf330874 (planner: support basic usage of partial index (#65051))
 	pushDownCtx := util.GetPushDownCtx(ds.SCtx())
 	for offset, item := range dnfItems {
 		cnfItems := expression.SplitCNFItems(item)

pkg/planner/core/operator/logicalop/BUILD.bazel

@@ -53,6 +53,7 @@ go_library(
         "//pkg/planner/core/constraint",
         "//pkg/planner/core/cost",
         "//pkg/planner/core/operator/baseimpl",
+        "//pkg/planner/core/partidx",
         "//pkg/planner/core/resolve",
         "//pkg/planner/core/rule/util",
         "//pkg/planner/funcdep",