plan: do not cache plan for query on range partition table (#15697)

planner/core/common_plans.go

@@ -253,7 +253,8 @@ func (e *Execute) getPhysicalPlan(ctx context.Context, sctx sessionctx.Context,
 		return nil, err
 	}
 	_, isTableDual := p.(*PhysicalTableDual)
-	if !isTableDual && prepared.UseCache {
+	isPartition := e.isPartition(p)
+	if !isTableDual && prepared.UseCache && !isPartition {
 		sctx.PreparedPlanCache().Put(cacheKey, NewPSTMTPlanCacheValue(p))
 	}
 	return p, err
@@ -329,6 +330,36 @@ func (e *Execute) rebuildRange(p Plan) error {
 	return nil
 }
 
+func checkPartitionInfo(pi *model.PartitionInfo) bool {
+	if pi != nil {
+		return true
+	}
+	return false
+}
+
+// Prepare plan cache is not support query plan on range partition table.
+func (e *Execute) isPartition(p Plan) bool {
+	isRange := false
+	switch x := p.(type) {
+	case *PhysicalTableReader:
+		ts := x.TablePlans[0].(*PhysicalTableScan)
+		return checkPartitionInfo(ts.Table.Partition)
+	case *PhysicalIndexLookUpReader:
+		is := x.IndexPlans[0].(*PhysicalIndexScan)
+		return checkPartitionInfo(is.Table.Partition)
+	case *PhysicalIndexReader:
+		is := x.IndexPlans[0].(*PhysicalIndexScan)
+		return checkPartitionInfo(is.Table.Partition)
+	case PhysicalPlan:
+		for _, child := range x.Children() {
+			if e.isPartition(child) {
+				isRange = true
+			}
+		}
+	}
+	return isRange
+}
+
 func (e *Execute) buildRangeForIndexScan(sctx sessionctx.Context, is *PhysicalIndexScan) ([]*ranger.Range, error) {
 	idxCols, colLengths := expression.IndexInfo2Cols(is.schema.Columns, is.Index)
 	if len(idxCols) == 0 {

planner/core/prepare_test.go

@@ -370,3 +370,92 @@ func (s *testPrepareSuite) TestPrepareForGroupByItems(c *C) {
 	tk.MustExec("set @a=2.0;")
 	tk.MustQuery("execute s1 using @a;").Check(testkit.Rows("3"))
 }
+
+func (s *testPrepareSuite) TestPrepareCacheForPartition(c *C) {
+	defer testleak.AfterTest(c)()
+	store, dom, err := newStoreWithBootstrap()
+	c.Assert(err, IsNil)
+	tk := testkit.NewTestKit(c, store)
+	orgEnable := core.PreparedPlanCacheEnabled()
+	orgCapacity := core.PreparedPlanCacheCapacity
+	orgMemGuardRatio := core.PreparedPlanCacheMemoryGuardRatio
+	orgMaxMemory := core.PreparedPlanCacheMaxMemory
+	defer func() {
+		dom.Close()
+		store.Close()
+		core.SetPreparedPlanCache(orgEnable)
+		core.PreparedPlanCacheCapacity = orgCapacity
+		core.PreparedPlanCacheMemoryGuardRatio = orgMemGuardRatio
+		core.PreparedPlanCacheMaxMemory = orgMaxMemory
+	}()
+	core.SetPreparedPlanCache(true)
+	core.PreparedPlanCacheCapacity = 100
+	core.PreparedPlanCacheMemoryGuardRatio = 0.1
+	// PreparedPlanCacheMaxMemory is set to MAX_UINT64 to make sure the cache
+	// behavior would not be effected by the uncertain memory utilization.
+	core.PreparedPlanCacheMaxMemory.Store(math.MaxUint64)
+
+	tk.MustExec("use test")
+	// Test for PointGet and IndexRead.
+	tk.MustExec("drop table if exists t_index_read")
+	tk.MustExec("create table t_index_read (id int, k int, c varchar(10), primary key (id, k)) partition by hash(id+k) partitions 10")
+	tk.MustExec("insert into t_index_read values (1, 2, 'abc'), (3, 4, 'def'), (5, 6, 'xyz')")
+	tk.MustExec("prepare stmt1 from 'select c from t_index_read where id = ? and k = ?;'")
+	tk.MustExec("set @id=1, @k=2")
+	// When executing one statement at the first time, we don't use cache, so we need to execute it at least twice to test the cache.
+	tk.MustQuery("execute stmt1 using @id, @k").Check(testkit.Rows("abc"))
+	tk.MustQuery("execute stmt1 using @id, @k").Check(testkit.Rows("abc"))
+	tk.MustExec("set @id=5, @k=6")
+	tk.MustQuery("execute stmt1 using @id, @k").Check(testkit.Rows("xyz"))
+	tk.MustExec("prepare stmt2 from 'select c from t_index_read where id = ? and k = ? and 1 = 1;'")
+	tk.MustExec("set @id=1, @k=2")
+	tk.MustQuery("execute stmt2 using @id, @k").Check(testkit.Rows("abc"))
+	tk.MustQuery("execute stmt2 using @id, @k").Check(testkit.Rows("abc"))
+	tk.MustExec("set @id=5, @k=6")
+	tk.MustQuery("execute stmt2 using @id, @k").Check(testkit.Rows("xyz"))
+	// Test for TableScan.
+	tk.MustExec("drop table if exists t_table_read")
+	tk.MustExec("create table t_table_read (id int, k int, c varchar(10), primary key(id)) partition by hash(id) partitions 10")
+	tk.MustExec("insert into t_table_read values (1, 2, 'abc'), (3, 4, 'def'), (5, 6, 'xyz')")
+	tk.MustExec("prepare stmt3 from 'select c from t_index_read where id = ?;'")
+	tk.MustExec("set @id=1")
+	// When executing one statement at the first time, we don't use cache, so we need to execute it at least twice to test the cache.
+	tk.MustQuery("execute stmt3 using @id").Check(testkit.Rows("abc"))
+	tk.MustQuery("execute stmt3 using @id").Check(testkit.Rows("abc"))
+	tk.MustExec("set @id=5")
+	tk.MustQuery("execute stmt3 using @id").Check(testkit.Rows("xyz"))
+	tk.MustExec("prepare stmt4 from 'select c from t_index_read where id = ? and k = ?'")
+	tk.MustExec("set @id=1, @k=2")
+	tk.MustQuery("execute stmt4 using @id, @k").Check(testkit.Rows("abc"))
+	tk.MustQuery("execute stmt4 using @id, @k").Check(testkit.Rows("abc"))
+	tk.MustExec("set @id=5, @k=6")
+	tk.MustQuery("execute stmt4 using @id, @k").Check(testkit.Rows("xyz"))
+	// Query on range partition tables should not raise error.
+	tk.MustExec("create table t_range_index (id int, k int, c varchar(10), primary key(id)) partition by range(id) ( PARTITION p0 VALUES LESS THAN (4), PARTITION p1 VALUES LESS THAN (14),PARTITION p2 VALUES LESS THAN (20) )")
+	tk.MustExec("insert into t_range_index values (1, 2, 'abc'), (5, 4, 'def'), (13, 6, 'xyz'), (17, 6, 'hij')")
+	tk.MustExec("prepare stmt5 from 'select c from t_range_index where id = ?'")
+	tk.MustExec("set @id=1")
+	tk.MustQuery("execute stmt5 using @id").Check(testkit.Rows("abc"))
+	tk.MustQuery("execute stmt5 using @id").Check(testkit.Rows("abc"))
+	tk.MustExec("set @id=5")
+	tk.MustQuery("execute stmt5 using @id").Check(testkit.Rows("def"))
+	tk.MustQuery("execute stmt5 using @id").Check(testkit.Rows("def"))
+	tk.MustExec("set @id=13")
+	tk.MustQuery("execute stmt5 using @id").Check(testkit.Rows("xyz"))
+	tk.MustExec("set @id=17")
+	tk.MustQuery("execute stmt5 using @id").Check(testkit.Rows("hij"))
+
+	tk.MustExec("create table t_range_table (id int, k int, c varchar(10)) partition by range(id) ( PARTITION p0 VALUES LESS THAN (4), PARTITION p1 VALUES LESS THAN (14),PARTITION p2 VALUES LESS THAN (20) )")
+	tk.MustExec("insert into t_range_table values (1, 2, 'abc'), (5, 4, 'def'), (13, 6, 'xyz'), (17, 6, 'hij')")
+	tk.MustExec("prepare stmt6 from 'select c from t_range_table where id = ?'")
+	tk.MustExec("set @id=1")
+	tk.MustQuery("execute stmt6 using @id").Check(testkit.Rows("abc"))
+	tk.MustQuery("execute stmt6 using @id").Check(testkit.Rows("abc"))
+	tk.MustExec("set @id=5")
+	tk.MustQuery("execute stmt6 using @id").Check(testkit.Rows("def"))
+	tk.MustQuery("execute stmt6 using @id").Check(testkit.Rows("def"))
+	tk.MustExec("set @id=13")
+	tk.MustQuery("execute stmt6 using @id").Check(testkit.Rows("xyz"))
+	tk.MustExec("set @id=17")
+	tk.MustQuery("execute stmt6 using @id").Check(testkit.Rows("hij"))
+}