planner, add hint nth_plan(x) to help user force a plan

planner/core/find_best_task_test.go

@@ -42,15 +42,16 @@ func (ds mockDataSource) Init(ctx sessionctx.Context) *mockDataSource {
 	return &ds
 }
 
-func (ds *mockDataSource) findBestTask(prop *property.PhysicalProperty) (task, error) {
+func (ds *mockDataSource) findBestTask(prop *property.PhysicalProperty, planCounter *PlanCounterTp) (task, int64, error) {
 	// It can satisfy any of the property!
 	// Just use a TableDual for convenience.
 	p := PhysicalTableDual{}.Init(ds.ctx, &property.StatsInfo{RowCount: 1}, 0)
 	task := &rootTask{
 		p:   p,
 		cst: 10000,
 	}
-	return task, nil
+	planCounter.Dec(1)
+	return task, 1, nil
 }
 
 // mockLogicalPlan4Test is a LogicalPlan which is used for unit test.
@@ -151,7 +152,7 @@ func (s *testFindBestTaskSuite) TestCostOverflow(c *C) {
 	mockPlan.SetChildren(mockDS)
 	// An empty property is enough for this test.
 	prop := property.NewPhysicalProperty(property.RootTaskType, nil, false, 0, false)
-	t, err := mockPlan.findBestTask(prop)
+	t, _, err := mockPlan.findBestTask(prop, &PlanCounterDisabled)
 	c.Assert(err, IsNil)
 	// The cost should be overflowed, but the task shouldn't be invalid.
 	c.Assert(t.invalid(), IsFalse)
@@ -178,7 +179,7 @@ func (s *testFindBestTaskSuite) TestEnforcedProperty(c *C) {
 		Enforced: false,
 	}
 	// should return invalid task because no physical plan can match this property.
-	task, err := mockPlan.findBestTask(prop0)
+	task, _, err := mockPlan.findBestTask(prop0, &PlanCounterDisabled)
 	c.Assert(err, IsNil)
 	c.Assert(task.invalid(), IsTrue)
 
@@ -187,7 +188,7 @@ func (s *testFindBestTaskSuite) TestEnforcedProperty(c *C) {
 		Enforced: true,
 	}
 	// should return the valid task when the property is enforced.
-	task, err = mockPlan.findBestTask(prop1)
+	task, _, err = mockPlan.findBestTask(prop1, &PlanCounterDisabled)
 	c.Assert(err, IsNil)
 	c.Assert(task.invalid(), IsFalse)
 }
@@ -210,7 +211,7 @@ func (s *testFindBestTaskSuite) TestHintCannotFitProperty(c *C) {
 		Items:    items,
 		Enforced: true,
 	}
-	task, err := mockPlan0.findBestTask(prop0)
+	task, _, err := mockPlan0.findBestTask(prop0, &PlanCounterDisabled)
 	c.Assert(err, IsNil)
 	c.Assert(task.invalid(), IsFalse)
 	_, enforcedSort := task.plan().(*PhysicalSort)
@@ -226,7 +227,7 @@ func (s *testFindBestTaskSuite) TestHintCannotFitProperty(c *C) {
 		Items:    items,
 		Enforced: false,
 	}
-	task, err = mockPlan0.findBestTask(prop1)
+	task, _, err = mockPlan0.findBestTask(prop1, &PlanCounterDisabled)
 	c.Assert(err, IsNil)
 	c.Assert(task.invalid(), IsFalse)
 	_, enforcedSort = task.plan().(*PhysicalSort)
@@ -247,7 +248,7 @@ func (s *testFindBestTaskSuite) TestHintCannotFitProperty(c *C) {
 		canGeneratePlan2: false,
 	}.Init(ctx)
 	mockPlan1.SetChildren(mockDS)
-	task, err = mockPlan1.findBestTask(prop2)
+	task, _, err = mockPlan1.findBestTask(prop2, &PlanCounterDisabled)
 	c.Assert(err, IsNil)
 	c.Assert(task.invalid(), IsFalse)
 	c.Assert(ctx.GetSessionVars().StmtCtx.WarningCount(), Equals, uint16(1))
@@ -263,7 +264,7 @@ func (s *testFindBestTaskSuite) TestHintCannotFitProperty(c *C) {
 		Items:    items,
 		Enforced: true,
 	}
-	task, err = mockPlan1.findBestTask(prop3)
+	task, _, err = mockPlan1.findBestTask(prop3, &PlanCounterDisabled)
 	c.Assert(err, IsNil)
 	c.Assert(task.invalid(), IsFalse)
 	c.Assert(ctx.GetSessionVars().StmtCtx.WarningCount(), Equals, uint16(1))

planner/core/logical_plan_test.go

@@ -1379,7 +1379,7 @@ func (s *testPlanSuite) optimize(ctx context.Context, sql string) (PhysicalPlan,
 	if err != nil {
 		return nil, nil, err
 	}
-	p, _, err = physicalOptimize(p.(LogicalPlan))
+	p, _, err = physicalOptimize(p.(LogicalPlan), &PlanCounterDisabled)
 	return p.(PhysicalPlan), stmt, err
 }
 

planner/core/optimizer.go

@@ -132,7 +132,11 @@ func DoOptimize(ctx context.Context, sctx sessionctx.Context, flag uint64, logic
 	if !AllowCartesianProduct.Load() && existsCartesianProduct(logic) {
 		return nil, 0, errors.Trace(ErrCartesianProductUnsupported)
 	}
-	physical, cost, err := physicalOptimize(logic)
+	planCounter := PlanCounterTp(sctx.GetSessionVars().StmtCtx.StmtHints.ForceNthPlan)
+	if planCounter == 0 {
+		planCounter = -1
+	}
+	physical, cost, err := physicalOptimize(logic, &planCounter)
 	if err != nil {
 		return nil, 0, err
 	}
@@ -169,7 +173,7 @@ func isLogicalRuleDisabled(r logicalOptRule) bool {
 	return disabled
 }
 
-func physicalOptimize(logic LogicalPlan) (PhysicalPlan, float64, error) {
+func physicalOptimize(logic LogicalPlan, planCounter *PlanCounterTp) (PhysicalPlan, float64, error) {
 	if _, err := logic.recursiveDeriveStats(); err != nil {
 		return nil, 0, err
 	}
@@ -181,10 +185,14 @@ func physicalOptimize(logic LogicalPlan) (PhysicalPlan, float64, error) {
 		ExpectedCnt: math.MaxFloat64,
 	}
 
-	t, err := logic.findBestTask(prop)
+	logic.SCtx().GetSessionVars().StmtCtx.TaskMapBakTS = 0
+	t, _, err := logic.findBestTask(prop, planCounter)
 	if err != nil {
 		return nil, 0, err
 	}
+	if *planCounter > 0 {
+		logic.SCtx().GetSessionVars().StmtCtx.AppendWarning(errors.Errorf("The parameter of nth_plan() is out of range."))
+	}
 	if t.invalid() {
 		return nil, 0, ErrInternal.GenWithStackByArgs("Can't find a proper physical plan for this query")
 	}

planner/core/physical_plan_test.go

@@ -1488,3 +1488,45 @@ func (s *testPlanSuite) TestHintFromDiffDatabase(c *C) {
 		c.Assert(core.ToString(p), Equals, output[i].Plan, comment)
 	}
 }
+
+func (s *testPlanSuite) TestNthPlanHintWithExplain(c *C) {
+	defer testleak.AfterTest(c)()
+	store, dom, err := newStoreWithBootstrap()
+	c.Assert(err, IsNil)
+	defer func() {
+		dom.Close()
+		store.Close()
+	}()
+	se, err := session.CreateSession4Test(store)
+	c.Assert(err, IsNil)
+	ctx := context.Background()
+	_, err = se.Execute(ctx, "use test")
+	c.Assert(err, IsNil)
+	_, err = se.Execute(ctx, `drop table if exists test.tt`)
+	c.Assert(err, IsNil)
+	_, err = se.Execute(ctx, `create table test.tt (a int,b int, index(a), index(b));`)
+	c.Assert(err, IsNil)
+
+	_, err = se.Execute(ctx, "insert into tt values (1, 1), (2, 2), (3, 4)")
+	c.Assert(err, IsNil)
+
+	var input []string
+	var output []struct {
+		SQL  string
+		Plan string
+	}
+	is := domain.GetDomain(se).InfoSchema()
+	s.testData.GetTestCases(c, &input, &output)
+	for i, tt := range input {
+		comment := Commentf("case:%v sql: %s", i, tt)
+		stmt, err := s.ParseOneStmt(tt, "", "")
+		c.Assert(err, IsNil, comment)
+		p, _, err := planner.Optimize(ctx, se, stmt, is)
+		c.Assert(err, IsNil, comment)
+		s.testData.OnRecord(func() {
+			output[i].SQL = tt
+			output[i].Plan = core.ToString(p)
+		})
+		c.Assert(core.ToString(p), Equals, output[i].Plan, comment)
+	}
+}

planner/core/plan.go

@@ -154,8 +154,12 @@ type LogicalPlan interface {
 	// findBestTask converts the logical plan to the physical plan. It's a new interface.
 	// It is called recursively from the parent to the children to create the result physical plan.
 	// Some logical plans will convert the children to the physical plans in different ways, and return the one
-	// with the lowest cost.
-	findBestTask(prop *property.PhysicalProperty) (task, error)
+	// With the lowest cost and how many plans are found in this function.
+	// planCounter is a counter for planner to force a plan.
+	// If planCounter > 0, the clock_th plan generated in this function will be returned.
+	// If planCounter = 0, the plan generated in this function will not be considered.
+	// If planCounter = -1, then we will not force plan.
+	findBestTask(prop *property.PhysicalProperty, planCounter *PlanCounterTp) (task, int64, error)
 
 	// BuildKeyInfo will collect the information of unique keys into schema.
 	// Because this method is also used in cascades planner, we cannot use
@@ -198,6 +202,9 @@ type LogicalPlan interface {
 
 	// SetChild sets the ith child for the plan.
 	SetChild(i int, child LogicalPlan)
+
+	// rollBackTaskMap roll back all taskMap's logs after TimeStamp TS.
+	rollBackTaskMap(TS uint64)
 }
 
 // PhysicalPlan is a tree of the physical operators.
@@ -245,10 +252,14 @@ type PhysicalPlan interface {
 type baseLogicalPlan struct {
 	basePlan
 
-	taskMap   map[string]task
-	self      LogicalPlan
-	maxOneRow bool
-	children  []LogicalPlan
+	taskMap map[string]task
+	// taskMapBak forms a backlog stack of taskMap, used to roll back the taskMap.
+	taskMapBak []string
+	// taskMapBakTS stores the timestamps of logs.
+	taskMapBakTS []uint64
+	self         LogicalPlan
+	maxOneRow    bool
+	children     []LogicalPlan
 }
 
 func (p *baseLogicalPlan) MaxOneRow() bool {
@@ -310,13 +321,53 @@ func (p *basePhysicalPlan) ExtractCorrelatedCols() []*expression.CorrelatedColum
 	return nil
 }
 
+// GetlogicalTS4TaskMap get the logical TimeStamp now to help rollback the TaskMap changes after that.
+func (p *baseLogicalPlan) GetlogicalTS4TaskMap() uint64 {
+	p.ctx.GetSessionVars().StmtCtx.TaskMapBakTS += 1
+	return p.ctx.GetSessionVars().StmtCtx.TaskMapBakTS
+}
+
+func (p *baseLogicalPlan) rollBackTaskMap(TS uint64) {
+	if !p.ctx.GetSessionVars().StmtCtx.StmtHints.TaskMapNeedBackUp() {
+		return
+	}
+	if len(p.taskMapBak) > 0 {
+		// Rollback all the logs with TimeStamp TS.
+		N := len(p.taskMapBak)
+		for i := 0; i < N; i++ {
+			cur := p.taskMapBak[i]
+			if p.taskMapBakTS[i] < TS {
+				continue
+			}
+
+			// Remove the i_th log.
+			p.taskMapBak = append(p.taskMapBak[:i], p.taskMapBak[i+1:]...)
+			p.taskMapBakTS = append(p.taskMapBakTS[:i], p.taskMapBakTS[i+1:]...)
+			i--
+			N--
+
+			// Roll back taskMap.
+			p.taskMap[cur] = nil
+		}
+	}
+	for _, child := range p.children {
+		child.rollBackTaskMap(TS)
+	}
+}
+
 func (p *baseLogicalPlan) getTask(prop *property.PhysicalProperty) task {
 	key := prop.HashCode()
 	return p.taskMap[string(key)]
 }
 
 func (p *baseLogicalPlan) storeTask(prop *property.PhysicalProperty, task task) {
 	key := prop.HashCode()
+	if p.ctx.GetSessionVars().StmtCtx.StmtHints.TaskMapNeedBackUp() {
+		// Empty string for useless change.
+		TS := p.GetlogicalTS4TaskMap()
+		p.taskMapBakTS = append(p.taskMapBakTS, TS)
+		p.taskMapBak = append(p.taskMapBak, string(key))
+	}
 	p.taskMap[string(key)] = task
 }
 
@@ -373,9 +424,11 @@ func newBasePlan(ctx sessionctx.Context, tp string, offset int) basePlan {
 
 func newBaseLogicalPlan(ctx sessionctx.Context, tp string, self LogicalPlan, offset int) baseLogicalPlan {
 	return baseLogicalPlan{
-		taskMap:  make(map[string]task),
-		basePlan: newBasePlan(ctx, tp, offset),
-		self:     self,
+		taskMap:      make(map[string]task),
+		taskMapBak:   make([]string, 0, 10),
+		taskMapBakTS: make([]uint64, 0, 10),
+		basePlan:     newBasePlan(ctx, tp, offset),
+		self:         self,
 	}
 }
 

planner/core/plan_test.go

@@ -232,3 +232,49 @@ func compareStringSlice(c *C, ss1, ss2 []string) {
 		c.Assert(s, Equals, ss2[i])
 	}
 }
+
+func (s *testPlanNormalize) TestNthPlanHint(c *C) {
+	tk := testkit.NewTestKit(c, s.store)
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists tt")
+	tk.MustExec("create table tt (a int,b int, index(a), index(b));")
+	tk.MustExec("insert into tt values (1, 1), (2, 2), (3, 4)")
+
+	tk.MustExec("explain select /*+nth_plan(4)*/ * from tt where a=1 and b=1;")
+	tk.MustQuery("show warnings").Check(testkit.Rows(
+		"Warning 1105 The parameter of nth_plan() is out of range."))
+
+	// Test hints for nth_plan(x).
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t (a int, b int, c int, index(a), index(b), index(a,b))")
+	tk.MustQuery("explain format='hint' select * from t where a=1 and b=1").Check(testkit.Rows(
+		"use_index(@`sel_1` `test`.`t` `a_2`)"))
+	tk.MustQuery("explain format='hint' select /*+ nth_plan(1) */ * from t where a=1 and b=1").Check(testkit.Rows(
+		"use_index(@`sel_1` `test`.`t` ), nth_plan(1)"))
+	tk.MustQuery("explain format='hint' select /*+ nth_plan(2) */ * from t where a=1 and b=1").Check(testkit.Rows(
+		"use_index(@`sel_1` `test`.`t` `a_2`), nth_plan(2)"))
+
+	tk.MustExec("explain format='hint' select /*+ nth_plan(3) */ * from t where a=1 and b=1")
+	tk.MustQuery("show warnings").Check(testkit.Rows(
+		"Warning 1105 The parameter of nth_plan() is out of range."))
+
+	// Test warning for multiply hints.
+	tk.MustQuery("explain format='hint' select /*+ nth_plan(1) nth_plan(2) */ * from t where a=1 and b=1").Check(testkit.Rows(
+		"use_index(@`sel_1` `test`.`t` `a_2`), nth_plan(1), nth_plan(2)"))
+	tk.MustQuery("show warnings").Check(testkit.Rows(
+		"Warning 1105 NTH_PLAN() is defined more than once, only the last definition takes effect: NTH_PLAN(2)",
+		"Warning 1105 NTH_PLAN() is defined more than once, only the last definition takes effect: NTH_PLAN(2)"))
+
+	// Test the correctness of generated plans.
+	tk.MustExec("insert into t values (1,1,1)")
+	tk.MustQuery("select  /*+ nth_plan(1) */ * from t where a=1 and b=1;").Check(testkit.Rows(
+		"1 1 1"))
+	tk.MustQuery("select  /*+ nth_plan(2) */ * from t where a=1 and b=1;").Check(testkit.Rows(
+		"1 1 1"))
+	tk.MustQuery("select  /*+ nth_plan(1) */ * from tt where a=1 and b=1;").Check(testkit.Rows(
+		"1 1"))
+	tk.MustQuery("select  /*+ nth_plan(2) */ * from tt where a=1 and b=1;").Check(testkit.Rows(
+		"1 1"))
+	tk.MustQuery("select  /*+ nth_plan(3) */ * from tt where a=1 and b=1;").Check(testkit.Rows(
+		"1 1"))
+}

planner/core/testdata/plan_suite_in.json

@@ -632,5 +632,14 @@
     "cases": [
       "select /*+ inl_hash_join(test.t1) */ * from test.t2 join test.t1 on test.t2.a = test.t1.a"
     ]
+  },
+  {
+    "name": "TestNthPlanHintWithExplain",
+    "cases": [
+      "select /*+nth_plan(1)*/ * from test.tt where a=1 and b=1",
+      "select /*+nth_plan(2)*/ * from test.tt where a=1 and b=1;",
+      "select /*+nth_plan(3)*/ * from test.tt where a=1 and b=1;",
+      "select /*+nth_plan(2)*/ * from test.tt where a=1 and b=1;"
+    ]
   }
 ]

planner/core/testdata/plan_suite_out.json

@@ -2075,5 +2075,26 @@
         "Plan": "IndexHashJoin{IndexReader(Index(t2.idx_a)[[-inf,+inf]])->IndexReader(Index(t1.idx_a)[[NULL,+inf]]->Sel([not(isnull(test.t1.a))]))}(test.t2.a,test.t1.a)"
       }
     ]
+  },
+  {
+    "Name": "TestNthPlanHintWithExplain",
+    "Cases": [
+      {
+        "SQL": "select /*+nth_plan(1)*/ * from test.tt where a=1 and b=1",
+        "Plan": "TableReader(Table(tt)->Sel([eq(test.tt.a, 1) eq(test.tt.b, 1)]))"
+      },
+      {
+        "SQL": "select /*+nth_plan(2)*/ * from test.tt where a=1 and b=1;",
+        "Plan": "IndexLookUp(Index(tt.a)[[1,1]], Table(tt)->Sel([eq(test.tt.b, 1)]))"
+      },
+      {
+        "SQL": "select /*+nth_plan(3)*/ * from test.tt where a=1 and b=1;",
+        "Plan": "IndexLookUp(Index(tt.b)[[1,1]], Table(tt)->Sel([eq(test.tt.a, 1)]))"
+      },
+      {
+        "SQL": "select /*+nth_plan(2)*/ * from test.tt where a=1 and b=1;",
+        "Plan": "IndexLookUp(Index(tt.a)[[1,1]], Table(tt)->Sel([eq(test.tt.b, 1)]))"
+      }
+    ]
   }
 ]