planner: support plan cache for cluster index (#18716) (#18819)

executor/prepared_test.go

@@ -15,6 +15,7 @@ package executor_test
 
 import (
 	"fmt"
+	"strings"
 
 	. "github.com/pingcap/check"
 	"github.com/pingcap/parser/auth"
@@ -23,6 +24,7 @@ import (
 	plannercore "github.com/pingcap/tidb/planner/core"
 	"github.com/pingcap/tidb/util"
 	"github.com/pingcap/tidb/util/testkit"
+	"github.com/pingcap/tidb/util/testleak"
 )
 
 func (s *testSuite1) TestPreparedNameResolver(c *C) {
@@ -102,3 +104,87 @@ func (s *testSuite1) TestPrepareStmtAfterIsolationReadChange(c *C) {
 	c.Assert(tk.Se.GetSessionVars().PreparedStmts[1].(*plannercore.CachedPrepareStmt).NormalizedSQL, Equals, "select * from t")
 	c.Assert(tk.Se.GetSessionVars().PreparedStmts[1].(*plannercore.CachedPrepareStmt).NormalizedPlan, Equals, "")
 }
+
+func (s *testSuite9) TestPlanCacheOnPointGet(c *C) {
+	defer testleak.AfterTest(c)()
+	store, dom, err := newStoreWithBootstrap()
+	c.Assert(err, IsNil)
+	tk := testkit.NewTestKit(c, store)
+	defer func() {
+		dom.Close()
+		store.Close()
+	}()
+	orgEnable := plannercore.PreparedPlanCacheEnabled()
+	defer func() {
+		plannercore.SetPreparedPlanCache(orgEnable)
+	}()
+	plannercore.SetPreparedPlanCache(true)
+	tk.MustExec("use test")
+
+	// For point get
+	tk.MustExec("drop table if exists t1")
+	tk.MustExec("create table t1(a varchar(20), b varchar(20), c varchar(20), primary key(a, b))")
+	tk.MustExec("insert into t1 values('1','1','111'),('2','2','222'),('3','3','333')")
+	tk.MustExec(`prepare stmt2 from "select * from t1 where t1.a = ? and t1.b = ?"`)
+	tk.MustExec("set @v1 = 1")
+	tk.MustExec("set @v2 = 1")
+	tk.MustQuery("execute stmt2 using @v1,@v2").Check(testkit.Rows("1 1 111"))
+	tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("0"))
+	tk.MustExec("set @v1 = 2")
+	tk.MustExec("set @v2 = 2")
+	tk.MustQuery("execute stmt2 using @v1,@v2").Check(testkit.Rows("2 2 222"))
+	tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("1"))
+	tk.MustExec("set @v1 = 3")
+	tk.MustExec("set @v2 = 3")
+	tk.MustQuery("execute stmt2 using @v1,@v2").Check(testkit.Rows("3 3 333"))
+	tkProcess := tk.Se.ShowProcess()
+	ps := []*util.ProcessInfo{tkProcess}
+	tk.Se.SetSessionManager(&mockSessionManager1{PS: ps})
+	rows := tk.MustQuery(fmt.Sprintf("explain for connection %d", tkProcess.ID)).Rows()
+	c.Assert(strings.Index(rows[len(rows)-1][0].(string), `Point_Get`), Equals, 0)
+
+	// For CBO point get and batch point get
+	// case 1:
+	tk.MustExec(`drop table if exists ta, tb`)
+	tk.MustExec(`create table ta (a int primary key, b int)`)
+	tk.MustExec(`insert ta values (1, 1), (2, 2)`)
+	tk.MustExec(`create table tb (a int primary key, b int)`)
+	tk.MustExec(`insert tb values (1, 1), (2, 2)`)
+	tk.MustExec(`prepare stmt1 from "select * from ta, tb where ta.a = tb.a and ta.a = ?"`)
+	tk.MustExec(`set @v1 = 1, @v2 = 2`)
+	tk.MustQuery(`execute stmt1 using @v1`).Check(testkit.Rows("1 1 1 1"))
+	tk.MustQuery(`execute stmt1 using @v2`).Check(testkit.Rows("2 2 2 2"))
+	tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("1"))
+
+	// case 2:
+	tk.MustExec(`drop table if exists ta, tb`)
+	tk.MustExec(`create table ta (a varchar(10) primary key, b int not null)`)
+	tk.MustExec(`insert ta values ('a', 1), ('b', 2)`)
+	tk.MustExec(`create table tb (b int primary key, c int)`)
+	tk.MustExec(`insert tb values (1, 1), (2, 2)`)
+	tk.MustExec(`prepare stmt1 from "select * from ta, tb where ta.b = tb.b and ta.a = ?"`)
+	tk.MustExec(`set @v1 = 'a', @v2 = 'b'`)
+	tk.MustQuery(`execute stmt1 using @v1`).Check(testkit.Rows("a 1 1 1"))
+	tk.MustQuery(`execute stmt1 using @v2`).Check(testkit.Rows("b 2 2 2"))
+	tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("1"))
+	tk.MustQuery(`execute stmt1 using @v2`).Check(testkit.Rows("b 2 2 2"))
+	tkProcess = tk.Se.ShowProcess()
+	ps = []*util.ProcessInfo{tkProcess}
+	tk.Se.SetSessionManager(&mockSessionManager1{PS: ps})
+	rows = tk.MustQuery(fmt.Sprintf("explain for connection %d", tkProcess.ID)).Rows()
+	c.Assert(strings.Index(rows[1][0].(string), `Point_Get`), Equals, 6)
+
+	// case 3:
+	tk.MustExec(`drop table if exists ta, tb`)
+	tk.MustExec(`create table ta (a varchar(10), b varchar(10), c int, primary key (a, b))`)
+	tk.MustExec(`insert ta values ('a', 'a', 1), ('b', 'b', 2), ('c', 'c', 3)`)
+	tk.MustExec(`create table tb (b int primary key, c int)`)
+	tk.MustExec(`insert tb values (1, 1), (2, 2), (3,3)`)
+	tk.MustExec(`prepare stmt1 from "select * from ta, tb where ta.c = tb.b and ta.a = ? and ta.b = ?"`)
+	tk.MustExec(`set @v1 = 'a', @v2 = 'b', @v3 = 'c'`)
+	tk.MustQuery(`execute stmt1 using @v1, @v1`).Check(testkit.Rows("a a 1 1 1"))
+	tk.MustQuery(`execute stmt1 using @v2, @v2`).Check(testkit.Rows("b b 2 2 2"))
+	tk.MustExec(`prepare stmt2 from "select * from ta, tb where ta.c = tb.b and (ta.a, ta.b) in ((?, ?), (?, ?))"`)
+	tk.MustQuery(`execute stmt2 using @v1, @v1, @v2, @v2`).Check(testkit.Rows("a a 1 1 1", "b b 2 2 2"))
+	tk.MustQuery(`execute stmt2 using @v2, @v2, @v3, @v3`).Check(testkit.Rows("b b 2 2 2", "c c 3 3 3"))
+}

planner/core/common_plans.go

@@ -513,6 +513,30 @@ func (e *Execute) rebuildRange(p Plan) error {
 		}
 		return nil
 	case *BatchPointGetPlan:
+		if x.Path != nil {
+			if x.Path.IsTablePath {
+				x.Path.Ranges, err = ranger.BuildTableRange(x.Path.AccessConds, sc, x.Path.PkCol.RetType)
+				// For col = NULL case, the length of the final ranges could be empty.
+				if err != nil || len(x.Path.Ranges) != 1 {
+					return errors.Errorf("Rebuilding range for PointGet failed")
+				}
+				x.Handles = make([]int64, len(x.Path.Ranges))
+				for i, ran := range x.Path.Ranges {
+					x.Handles[i] = ran.LowVal[0].GetInt64()
+				}
+				return nil
+			}
+			res, err := ranger.DetachCondAndBuildRangeForIndex(p.SCtx(), x.Path.AccessConds, x.Path.IdxCols, x.Path.IdxColLens)
+			// For col = NULL case, the length of the final ranges could be empty.
+			if err != nil || len(res.Ranges) != 1 {
+				return errors.Errorf("Rebuilding range for BatchPointGet failed")
+			}
+			x.IndexValues = make([][]types.Datum, 0, len(res.Ranges))
+			for _, ran := range res.Ranges {
+				x.IndexValues = append(x.IndexValues, ran.LowVal)
+			}
+			return nil
+		}
 		for i, param := range x.HandleParams {
 			if param != nil {
 				x.Handles[i], err = param.Datum.ToInt64(sc)

planner/core/find_best_task.go

@@ -1295,6 +1295,7 @@ func (ds *DataSource) convertToBatchPointGet(prop *property.PhysicalProperty, ca
 		TblInfo:   ds.TableInfo(),
 		KeepOrder: !prop.IsEmpty(),
 		Columns:   ds.Columns,
+		Path:      candidate.path,
 	}.Init(ds.ctx, ds.stats.ScaleByExpectCnt(float64(len(candidate.path.Ranges))), ds.schema.Clone(), ds.names, ds.blockOffset)
 	if batchPointGetPlan.KeepOrder {
 		batchPointGetPlan.Desc = prop.Items[0].Desc

planner/core/point_get_plan.go

@@ -231,6 +231,8 @@ type BatchPointGetPlan struct {
 	Lock             bool
 	LockWaitTime     int64
 	Columns          []*model.ColumnInfo
+
+	Path *util.AccessPath
 }
 
 // attach2Task makes the current physical plan as the father of task's physicalPlan and updates the cost of