planner: add maybe good heuristics for index selection

executor/set_test.go

@@ -539,6 +539,16 @@ func (s *testSerialSuite1) TestSetVar(c *C) {
 	tk.MustExec(`set tidb_opt_limit_push_down_threshold = 20`)
 	tk.MustQuery(`select @@global.tidb_opt_limit_push_down_threshold`).Check(testkit.Rows("100"))
 	tk.MustQuery(`select @@tidb_opt_limit_push_down_threshold`).Check(testkit.Rows("20"))
+
+	tk.MustQuery("select @@tidb_enable_maybe_good_heuristics").Check(testkit.Rows("0"))
+	tk.MustExec("set global tidb_enable_maybe_good_heuristics = 1")
+	tk.MustQuery("select @@global.tidb_enable_maybe_good_heuristics").Check(testkit.Rows("1"))
+	tk.MustExec("set global tidb_enable_maybe_good_heuristics = 0")
+	tk.MustQuery("select @@global.tidb_enable_maybe_good_heuristics").Check(testkit.Rows("0"))
+	tk.MustExec("set session tidb_enable_maybe_good_heuristics = 1")
+	tk.MustQuery("select @@session.tidb_enable_maybe_good_heuristics").Check(testkit.Rows("1"))
+	tk.MustExec("set session tidb_enable_maybe_good_heuristics = 0")
+	tk.MustQuery("select @@session.tidb_enable_maybe_good_heuristics").Check(testkit.Rows("0"))
 }
 
 func (s *testSuite5) TestTruncateIncorrectIntSessionVar(c *C) {

planner/core/integration_test.go

@@ -4178,6 +4178,31 @@ func (s *testIntegrationSuite) TestHeuristicIndexSelection(c *C) {
 	}
 }
 
+func (s *testIntegrationSuite) TestMaybeGoodHeuristics(c *C) {
+	tk := testkit.NewTestKit(c, s.store)
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t(a int primary key, b int, c int, index idx_b(b))")
+	tk.MustExec("set tidb_enable_maybe_good_heuristics = 1")
+
+	var input []string
+	var output []struct {
+		SQL      string
+		Plan     []string
+		Warnings []string
+	}
+	s.testData.GetTestCases(c, &input, &output)
+	for i, tt := range input {
+		s.testData.OnRecord(func() {
+			output[i].SQL = tt
+			output[i].Plan = s.testData.ConvertRowsToStrings(tk.MustQuery("explain format = 'verbose' " + tt).Rows())
+			output[i].Warnings = s.testData.ConvertRowsToStrings(tk.MustQuery("show warnings").Rows())
+		})
+		tk.MustQuery("explain format = 'verbose' " + tt).Check(testkit.Rows(output[i].Plan...))
+		tk.MustQuery("show warnings").Check(testkit.Rows(output[i].Warnings...))
+	}
+}
+
 func (s *testIntegrationSuite) TestOutputSkylinePruningInfo(c *C) {
 	tk := testkit.NewTestKit(c, s.store)
 	tk.MustExec("use test")

planner/core/logical_plan_builder.go

@@ -3407,6 +3407,27 @@ func (b *PlanBuilder) TableHints() *tableHintInfo {
 	return &(b.tableHintInfo[len(b.tableHintInfo)-1])
 }
 
+func checkOrderByPK(ds *DataSource, byItems []*util.ByItems) bool {
+	if ds.tableInfo.PKIsHandle && len(byItems) == 1 {
+		if col, isCol := byItems[0].Expr.(*expression.Column); isCol && col.Equal(nil, ds.getPKIsHandleCol()) {
+			return true
+		}
+		return false
+	}
+	if ds.tableInfo.IsCommonHandle && len(byItems) == len(ds.commonHandleCols) {
+		orderByPK := true
+		for i, byItem := range byItems {
+			if col, isCol := byItem.Expr.(*expression.Column); !isCol || ds.commonHandleLens[i] != types.UnspecifiedLength ||
+				!col.Equal(nil, ds.commonHandleCols[i]) || (i > 0 && byItem.Desc != byItems[i-1].Desc) {
+				orderByPK = false
+				break
+			}
+		}
+		return orderByPK
+	}
+	return false
+}
+
 func (b *PlanBuilder) buildSelect(ctx context.Context, sel *ast.SelectStmt) (p LogicalPlan, err error) {
 	b.pushSelectOffset(sel.QueryBlockOffset)
 	b.pushTableHints(sel.TableHints, sel.QueryBlockOffset)
@@ -3466,6 +3487,8 @@ func (b *PlanBuilder) buildSelect(ctx context.Context, sel *ast.SelectStmt) (p L
 	if err != nil {
 		return nil, err
 	}
+	// For filling DataSource.orderByPKLimitN
+	ds, isDataSource := p.(*DataSource)
 
 	originalFields := sel.Fields.Fields
 	sel.Fields.Fields, err = b.unfoldWildStar(p, sel.Fields.Fields)
@@ -3650,6 +3673,9 @@ func (b *PlanBuilder) buildSelect(ctx context.Context, sel *ast.SelectStmt) (p L
 		if err != nil {
 			return nil, err
 		}
+		if logicalSort, isSort := p.(*LogicalSort); isSort && sel.Limit != nil && isDataSource {
+			ds.orderByPKLimitN = checkOrderByPK(ds, logicalSort.ByItems)
+		}
 	}
 
 	if sel.Limit != nil {

planner/core/logical_plans.go

@@ -550,6 +550,8 @@ type DataSource struct {
 	// 1. use `inside insert`, `update`, `delete` or `select for update` statement
 	// 2. isolation level is RC
 	isForUpdateRead bool
+	// orderByPKLimitN is true iff there exists `order by pk limit n` pattern.
+	orderByPKLimitN bool
 }
 
 // ExtractCorrelatedCols implements LogicalPlan interface.

planner/core/stats.go

@@ -373,6 +373,36 @@ func (ds *DataSource) derivePathStatsAndTryHeuristics() error {
 				ds.ctx.GetSessionVars().StmtCtx.AppendNote(errors.New(sb.String()))
 			}
 		}
+	} else if ds.ctx.GetSessionVars().EnableMaybeGoodHeuristics && ds.orderByPKLimitN {
+		// maybe-good heuristics
+		// For query like `where index_col = ... order by pk limit n`, if the count of `index_col = ...` is small enough, we prefer the index.
+		for _, path := range ds.possibleAccessPaths {
+			const smallCountAfterAccess = 100
+			if path.OnlyPointRange(ds.ctx.GetSessionVars().StmtCtx) && path.CountAfterAccess < smallCountAfterAccess {
+				selected = path
+				break
+			}
+		}
+		if selected != nil {
+			ds.possibleAccessPaths[0] = selected
+			ds.possibleAccessPaths = ds.possibleAccessPaths[:1]
+			// TODO: Can we make a more careful check on whether the optimization depends on mutable constants?
+			ds.ctx.GetSessionVars().StmtCtx.OptimDependOnMutableConst = true
+			if ds.ctx.GetSessionVars().StmtCtx.InVerboseExplain {
+				var tableName, pathName string
+				if ds.TableAsName.O == "" {
+					tableName = ds.tableInfo.Name.O
+				} else {
+					tableName = ds.TableAsName.O
+				}
+				if selected.IsTablePath() {
+					pathName = "handle of " + tableName
+				} else {
+					pathName = "index " + selected.Index.Name.O + " of " + tableName
+				}
+				ds.ctx.GetSessionVars().StmtCtx.AppendNote(errors.New(pathName + " is selected since the path has point ranges and fetches limited number of rows under ORDER BY PK LIMIT N pattern"))
+			}
+		}
 	}
 	return nil
 }

planner/core/testdata/integration_suite_in.json

@@ -331,6 +331,14 @@
       "select a, b, c from t2 where a = 1 and b = 2 and c in (1, 2, 3, 4, 5)"
     ]
   },
+  {
+    "name": "TestMaybeGoodHeuristics",
+    "cases": [
+      "select * from t where b = 3 order by a limit 10",
+      "select * from t where b in (2, 3, 4) order by a limit 10",
+      "select * from t as t1 where t1.a > any(select t2.b from t as t2 where t2.b = 2 and t2.c > t1.a order by t2.a limit 10)"
+    ]
+  },
   {
     "name": "TestOutputSkylinePruningInfo",
     "cases": [

planner/core/testdata/integration_suite_out.json

@@ -1732,6 +1732,57 @@
       }
     ]
   },
+  {
+    "Name": "TestMaybeGoodHeuristics",
+    "Cases": [
+      {
+        "SQL": "select * from t where b = 3 order by a limit 10",
+        "Plan": [
+          "IndexLookUp_20 10.00 201.28 root  limit embedded(offset:0, count:10)",
+          "├─Limit_19(Build) 10.00 590.00 cop[tikv]  offset:0, count:10",
+          "│ └─IndexRangeScan_17 10.00 590.00 cop[tikv] table:t, index:idx_b(b) range:[3,3], keep order:true, stats:pseudo",
+          "└─TableRowIDScan_18(Probe) 10.00 590.00 cop[tikv] table:t keep order:false, stats:pseudo"
+        ],
+        "Warnings": [
+          "Note 1105 index idx_b of t is selected since the path has point ranges and fetches limited number of rows under ORDER BY PK LIMIT N pattern"
+        ]
+      },
+      {
+        "SQL": "select * from t where b in (2, 3, 4) order by a limit 10",
+        "Plan": [
+          "TopN_9 10.00 379.61 root  test.t.a, offset:0, count:10",
+          "└─IndexLookUp_16 10.00 279.95 root  ",
+          "  ├─TopN_15(Build) 10.00 0.00 cop[tikv]  test.t.a, offset:0, count:10",
+          "  │ └─IndexRangeScan_13 30.00 1770.00 cop[tikv] table:t, index:idx_b(b) range:[2,2], [3,3], [4,4], keep order:false, stats:pseudo",
+          "  └─TableRowIDScan_14(Probe) 10.00 1770.00 cop[tikv] table:t keep order:false, stats:pseudo"
+        ],
+        "Warnings": [
+          "Note 1105 index idx_b of t is selected since the path has point ranges and fetches limited number of rows under ORDER BY PK LIMIT N pattern"
+        ]
+      },
+      {
+        "SQL": "select * from t as t1 where t1.a > any(select t2.b from t as t2 where t2.b = t1.b order by t2.a limit 10)",
+        "Plan": [
+          "Projection_16 10000.00 2533281.69 root  test.t.a, test.t.b, test.t.c",
+          "└─Apply_18 10000.00 2527263.69 root  CARTESIAN inner join, other cond:or(gt(test.t.a, Column#8), if(ne(Column#9, 0), NULL, 0))",
+          "  ├─TableReader_20(Build) 10000.00 54251.33 root  data:TableFullScan_19",
+          "  │ └─TableFullScan_19 10000.00 570020.00 cop[tikv] table:t1 keep order:false, stats:pseudo",
+          "  └─Selection_21(Probe) 0.80 244.90 root  ne(Column#10, 0)",
+          "    └─HashAgg_22 1.00 241.90 root  funcs:min(Column#13)->Column#8, funcs:sum(Column#14)->Column#9, funcs:count(1)->Column#10",
+          "      └─Projection_46 8.00 0.00 root  test.t.b, cast(isnull(test.t.b), decimal(22,0) BINARY)->Column#14",
+          "        └─Limit_28 8.00 194.50 root  offset:0, count:10",
+          "          └─IndexLookUp_37 8.00 194.50 root  ",
+          "            ├─IndexRangeScan_34(Build) 10.00 590.00 cop[tikv] table:t2, index:idx_b(b) range:[2,2], keep order:true, stats:pseudo",
+          "            └─Selection_36(Probe) 8.00 0.00 cop[tikv]  gt(test.t.c, test.t.a)",
+          "              └─TableRowIDScan_35 10.00 590.00 cop[tikv] table:t2 keep order:false, stats:pseudo"
+        ],
+        "Warnings": [
+          "Note 1105 index idx_b of t2 is selected since the path has point ranges and fetches limited number of rows under ORDER BY PK LIMIT N pattern",
+          "Note 1105 [t1] remain after pruning paths for t1 given Prop{SortItems: [], TaskTp: rootTask}"
+        ]
+      }
+    ]
+  },
   {
     "Name": "TestOutputSkylinePruningInfo",
     "Cases": [

sessionctx/variable/session.go

@@ -877,6 +877,9 @@ type SessionVars struct {
 	// TemporaryTableData stores committed kv values for temporary table for current session.
 	TemporaryTableData kv.MemBuffer
 
+	// EnableMaybeGoodHeuristics indicates whether to apply maybe-good heuristics when the optimizer generates plans.
+	EnableMaybeGoodHeuristics bool
+
 	// MPPStoreLastFailTime records the lastest fail time that a TiFlash store failed.
 	MPPStoreLastFailTime map[string]time.Time
 
@@ -1103,6 +1106,7 @@ func NewSessionVars() *SessionVars {
 		CTEMaxRecursionDepth:        DefCTEMaxRecursionDepth,
 		TMPTableSize:                DefTMPTableSize,
 		EnableGlobalTemporaryTable:  DefTiDBEnableGlobalTemporaryTable,
+		EnableMaybeGoodHeuristics:   DefTiDBEnableMaybeGoodHeuristics,
 		MPPStoreLastFailTime:        make(map[string]time.Time),
 		MPPStoreFailTTL:             DefTiDBMPPStoreFailTTL,
 	}

sessionctx/variable/sysvar.go

@@ -1787,6 +1787,10 @@ var defaultSysVars = []*SysVar{
 		s.EnableStableResultMode = TiDBOptOn(val)
 		return nil
 	}},
+	{Scope: ScopeGlobal | ScopeSession, Name: TiDBEnableMaybeGoodHeuristics, Value: BoolToOnOff(DefTiDBEnableMaybeGoodHeuristics), Hidden: true, Type: TypeBool, SetSession: func(s *SessionVars, val string) error {
+		s.EnableMaybeGoodHeuristics = TiDBOptOn(val)
+		return nil
+	}},
 }
 
 // FeedbackProbability points to the FeedbackProbability in statistics package.

sessionctx/variable/tidb_vars.go

@@ -583,6 +583,9 @@ const (
 
 	// TiDBEnableOrderedResultMode indicates if stabilize query results.
 	TiDBEnableOrderedResultMode = "tidb_enable_ordered_result_mode"
+
+	// TiDBEnableMaybeGoodHeuristics indicates whether to apply maybe-good heuristics when the optimizer generates plans.
+	TiDBEnableMaybeGoodHeuristics = "tidb_enable_maybe_good_heuristics"
 )
 
 // TiDB vars that have only global scope
@@ -744,6 +747,7 @@ const (
 	DefTMPTableSize                       = 16777216
 	DefTiDBEnableLocalTxn                 = false
 	DefTiDBEnableOrderedResultMode        = false
+	DefTiDBEnableMaybeGoodHeuristics      = false
 )
 
 // Process global variables.