planner: use code-gen to generate CloneForPlanCache method for Point/…

…BatchPoint/Limit (pingcap#55096)

ref pingcap#54057

pkg/kv/key.go

@@ -312,6 +312,9 @@ func NewCommonHandle(encoded []byte) (*CommonHandle, error) {
 
 // Copy implements the Handle interface.
 func (ch *CommonHandle) Copy() Handle {
+	if ch == nil {
+		return nil
+	}
 	encoded := make([]byte, len(ch.encoded))
 	copy(encoded, ch.encoded)
 	colEndOffsets := make([]uint16, len(ch.colEndOffsets))

pkg/planner/core/operator/baseimpl/plan.go

@@ -138,6 +138,14 @@ func (p *Plan) BuildPlanTrace() *tracing.PlanTrace {
 	return planTrace
 }
 
+// CloneWithNewCtx clones the plan with new context.
+func (p *Plan) CloneWithNewCtx(newCtx base.PlanContext) *Plan {
+	cloned := new(Plan)
+	*cloned = *p
+	cloned.ctx = newCtx
+	return cloned
+}
+
 // CloneForPlanCache clones the plan for Plan Cache.
 func (*Plan) CloneForPlanCache(base.PlanContext) (cloned base.Plan, ok bool) {
 	return nil, false

pkg/planner/core/plan_cache_rebuild_test.go

@@ -138,6 +138,14 @@ func TestPlanCacheClone(t *testing.T) {
 	//testCachedPlanClone(t, tk1, tk2, `prepare st from 'select /*+ inl_join(t1, t2, t3) */ * from t t1, t t2, t t3 where t1.b=t2.b and t2.b<t3.b and t1.a<?'`,
 	//	`set @a1=1, @a2=2`, `execute st using @a1`, `execute st using @a2`)
 
+	// Limit
+	testCachedPlanClone(t, tk1, tk2, `prepare st from 'select * from t where a<? limit 1'`,
+		`set @a1=1, @a2=2`, `execute st using @a1`, `execute st using @a2`)
+	testCachedPlanClone(t, tk1, tk2, `prepare st from 'select b from t use index(b) where b<=? limit 10'`,
+		`set @a1=1, @a2=2`, `execute st using @a1`, `execute st using @a2`)
+	testCachedPlanClone(t, tk1, tk2, `prepare st from 'select * from t use index(b) where b<=? limit 100'`,
+		`set @a1=1, @a2=2`, `execute st using @a1`, `execute st using @a2`)
+
 	// Sort
 	testCachedPlanClone(t, tk1, tk2, `prepare st from 'select * from t where a<? order by a'`,
 		`set @a1=1, @a2=2`, `execute st using @a1`, `execute st using @a2`)
@@ -148,16 +156,27 @@ func TestPlanCacheClone(t *testing.T) {
 	testCachedPlanClone(t, tk1, tk2, `prepare st from 'select * from t use index(b) where b<=? order by a+4'`,
 		`set @a1=1, @a2=2`, `execute st using @a1`, `execute st using @a2`)
 
-	// TODO: PointGet doesn't support Clone
+	// TopN
+	testCachedPlanClone(t, tk1, tk2, `prepare st from 'select * from t where a<? order by a limit 5'`,
+		`set @a1=1, @a2=2`, `execute st using @a1`, `execute st using @a2`)
+	testCachedPlanClone(t, tk1, tk2, `prepare st from 'select * from t where a>=? order by b limit 5'`,
+		`set @a1=1, @a2=2`, `execute st using @a1`, `execute st using @a2`)
+	testCachedPlanClone(t, tk1, tk2, `prepare st from 'select * from t use index(primary) where a<? and b<? order by a+b limit 5'`,
+		`set @a1=1, @b1=1, @a2=2, @b2=2`, `execute st using @a1,@b1`, `execute st using @a2,@b2`)
+	testCachedPlanClone(t, tk1, tk2, `prepare st from 'select * from t use index(b) where b<=? order by a+4 limit 5'`,
+		`set @a1=1, @a2=2`, `execute st using @a1`, `execute st using @a2`)
+
 	// PointPlan
-	//testCachedPlanClone(t, tk1, tk2, `prepare st from 'select * from t where a=?'`,
-	//	`set @a1=1, @a2=2`, `execute st using @a1`, `execute st using @a2`)
-	//testCachedPlanClone(t, tk1, tk2, `prepare st from 'select * from t where d=?'`,
-	//	`set @a1=1, @a2=2`, `execute st using @a1`, `execute st using @a2`)
-	//testCachedPlanClone(t, tk1, tk2, `prepare st from 'select * from t where a in (?,?)'`,
-	//	`set @a1=1,@b1=1, @a2=2,@b2=2`, `execute st using @a1,@b1`, `execute st using @a2,@b2`)
-	//testCachedPlanClone(t, tk1, tk2, `prepare st from 'select * from t where d in (?,?)'`,
-	//	`set @a1=1,@b1=1, @a2=2,@b2=2`, `execute st using @a1,@b1`, `execute st using @a2,@b2`)
+	testCachedPlanClone(t, tk1, tk2, `prepare st from 'select * from t where a=?'`,
+		`set @a1=1, @a2=2`, `execute st using @a1`, `execute st using @a2`)
+	testCachedPlanClone(t, tk1, tk2, `prepare st from 'select * from t where d=?'`,
+		`set @a1=1, @a2=2`, `execute st using @a1`, `execute st using @a2`)
+
+	// BatchPointPlan
+	testCachedPlanClone(t, tk1, tk2, `prepare st from 'select * from t where a in (?,?)'`,
+		`set @a1=1,@b1=1, @a2=2,@b2=2`, `execute st using @a1,@b1`, `execute st using @a2,@b2`)
+	testCachedPlanClone(t, tk1, tk2, `prepare st from 'select * from t where d in (?,?)'`,
+		`set @a1=1,@b1=1, @a2=2,@b2=2`, `execute st using @a1,@b1`, `execute st using @a2,@b2`)
 }
 
 func testCachedPlanClone(t *testing.T, tk1, tk2 *testkit.TestKit, prep, set, exec1, exec2 string) {

pkg/planner/core/plan_clone_generator.go

@@ -36,6 +36,7 @@ func genPlanCloneForPlanCacheCode() ([]byte, error) {
 	var structures = []any{PhysicalTableScan{}, PhysicalIndexScan{}, PhysicalSelection{}, PhysicalProjection{},
 		PhysicalSort{}, PhysicalTopN{}, PhysicalStreamAgg{}, PhysicalHashAgg{},
 		PhysicalHashJoin{}, PhysicalMergeJoin{}, PhysicalTableReader{}, PhysicalIndexReader{},
+		PointGetPlan{}, BatchPointGetPlan{}, PhysicalLimit{},
 		PhysicalIndexLookUpReader{}, PhysicalIndexMergeReader{}}
 	c := new(codeGen)
 	c.write(codeGenPrefix)
@@ -92,27 +93,30 @@ func genPlanCloneForPlanCache(x any) ([]byte, error) {
 			c.write(`basePlan, baseOK := op.%v.cloneForPlanCacheWithSelf(newCtx, cloned)
 							if !baseOK {return nil, false}
 							cloned.%v = *basePlan`, fieldName, fieldName)
-		case "[]expression.Expression":
-			c.write("cloned.%v = util.CloneExprs(op.%v)", f.Name, f.Name)
-		case "[]*ranger.Range":
-			c.write("cloned.%v = util.CloneRanges(op.%v)", f.Name, f.Name)
-		case "[]*util.ByItems":
-			c.write("cloned.%v = util.CloneByItems(op.%v)", f.Name, f.Name)
-		case "[]*expression.Column":
-			c.write("cloned.%v = util.CloneCols(op.%v)", f.Name, f.Name)
-		case "[]*expression.ScalarFunction":
-			c.write("cloned.%v = util.CloneScalarFunctions(op.%v)", f.Name, f.Name)
-		case "[]property.SortItem":
-			c.write("cloned.%v = util.CloneSortItem(op.%v)", f.Name, f.Name)
+		case "baseimpl.Plan", "core.baseSchemaProducer":
+			c.write("cloned.%v = *op.%v.CloneWithNewCtx(newCtx)", f.Name, f.Name)
+		case "[]expression.Expression", "[]*ranger.Range", "[]*util.ByItems", "[]*expression.Column", "[]model.CIStr",
+			"[]*expression.Constant", "[]*expression.ScalarFunction", "[]property.SortItem", "[]types.Datum", "[]kv.Handle":
+			structureName := strings.Split(f.Type.String(), ".")[1] + "s"
+			c.write("cloned.%v = util.Clone%v(op.%v)", f.Name, structureName, f.Name)
+		case "[][]*expression.Constant", "[][]types.Datum":
+			structureName := strings.Split(f.Type.String(), ".")[1]
+			c.write("cloned.%v = util.Clone%v2D(op.%v)", f.Name, structureName, f.Name)
+		case "context.PlanContext":
+			c.write("cloned.%v = newCtx", f.Name)
 		case "util.HandleCols":
 			c.write("if op.%v != nil {", f.Name)
 			c.write("cloned.%v = op.%v.Clone(newCtx.GetSessionVars().StmtCtx)", f.Name, f.Name)
 			c.write("}")
-		case "*core.PhysPlanPartInfo", "*core.PushedDownLimit":
+		case "*core.PhysPlanPartInfo", "*core.PushedDownLimit", "*expression.Schema":
 			c.write("cloned.%v = op.%v.Clone()", f.Name, f.Name)
-		case "*expression.Column":
+		case "kv.Handle":
 			c.write("if op.%v != nil {", f.Name)
-			c.write("cloned.%v = op.%v.Clone().(*expression.Column)", f.Name, f.Name)
+			c.write("cloned.%v = op.%v.Copy()", f.Name, f.Name)
+			c.write("}")
+		case "*expression.Column", "*expression.Constant":
+			c.write("if op.%v != nil {", f.Name)
+			c.write("cloned.%v = op.%v.Clone().(%v)", f.Name, f.Name, f.Type.String())
 			c.write("}")
 		case "base.PhysicalPlan":
 			c.write("%v, ok := op.%v.CloneForPlanCache(newCtx)", f.Name, f.Name)
@@ -122,6 +126,11 @@ func genPlanCloneForPlanCache(x any) ([]byte, error) {
 			c.write("%v, ok := clonePhysicalPlansForPlanCache(newCtx, op.%v)", f.Name, f.Name)
 			c.write("if !ok {return nil, false}")
 			c.write("cloned.%v = %v", f.Name, f.Name)
+		case "*int":
+			c.write("if op.%v != nil {", f.Name)
+			c.write("cloned.%v = new(int)", f.Name)
+			c.write("*cloned.%v = *op.%v", f.Name, f.Name)
+			c.write("}")
 		default:
 			return nil, fmt.Errorf("can't generate Clone method for type %v in %v", f.Type.String(), vType.String())
 		}