planner: add more test cases for Plan Clone

pkg/expression/builtin.go

@@ -409,11 +409,11 @@ func (b *baseBuiltinFunc) cloneFrom(from *baseBuiltinFunc) {
 	for _, arg := range from.args {
 		b.args = append(b.args, arg.Clone())
 	}
-	b.tp = from.tp
+	b.tp = from.tp.Clone()
 	b.pbCode = from.pbCode
 	b.bufAllocator = newLocalColumnPool()
 	b.childrenVectorizedOnce = new(sync.Once)
-	b.ctor = from.ctor
+	b.ctor = from.ctor.Clone()
 }
 
 func (*baseBuiltinFunc) Clone() builtinFunc {

pkg/expression/column.go

@@ -517,6 +517,13 @@ func (col *Column) EvalJSON(ctx EvalContext, row chunk.Row) (types.BinaryJSON, b
 // Clone implements Expression interface.
 func (col *Column) Clone() Expression {
 	newCol := *col
+	if col.RetType != nil {
+		newCol.RetType = col.RetType.Clone()
+	}
+	if col.hashcode != nil {
+		newCol.hashcode = make([]byte, len(col.hashcode))
+		copy(newCol.hashcode, col.hashcode)
+	}
 	return &newCol
 }
 

pkg/expression/constant.go

@@ -163,6 +163,17 @@ func (c *Constant) StringWithCtx(ctx ParamValues, redact string) string {
 // Clone implements Expression interface.
 func (c *Constant) Clone() Expression {
 	con := *c
+	con.RetType = c.RetType.Clone()
+	if c.ParamMarker != nil {
+		con.ParamMarker = &ParamMarker{order: c.ParamMarker.order}
+	}
+	if c.DeferredExpr != nil {
+		con.DeferredExpr = c.DeferredExpr.Clone()
+	}
+	if c.hashcode != nil {
+		con.hashcode = make([]byte, len(c.hashcode))
+		copy(con.hashcode, c.hashcode)
+	}
 	return &con
 }
 

pkg/expression/scalar_function.go

@@ -336,9 +336,16 @@ func ScalarFuncs2Exprs(funcs []*ScalarFunction) []Expression {
 func (sf *ScalarFunction) Clone() Expression {
 	c := &ScalarFunction{
 		FuncName: sf.FuncName,
-		RetType:  sf.RetType,
+		RetType:  sf.RetType.Clone(),
 		Function: sf.Function.Clone(),
-		hashcode: sf.hashcode,
+	}
+	if sf.hashcode != nil {
+		c.hashcode = make([]byte, len(sf.hashcode))
+		copy(c.hashcode, sf.hashcode)
+	}
+	if sf.canonicalhashcode != nil {
+		c.canonicalhashcode = make([]byte, len(sf.canonicalhashcode))
+		copy(c.canonicalhashcode, sf.canonicalhashcode)
 	}
 	c.SetCharsetAndCollation(sf.CharsetAndCollation())
 	c.SetCoercibility(sf.Coercibility())

pkg/expression/util_test.go

@@ -604,8 +604,18 @@ func (m *MockExpr) EvalJSON(ctx EvalContext, row chunk.Row) (val types.BinaryJSO
 	}
 	return types.BinaryJSON{}, m.i == nil, m.err
 }
-func (m *MockExpr) GetType(_ EvalContext) *types.FieldType            { return m.t }
-func (m *MockExpr) Clone() Expression                                 { return nil }
+func (m *MockExpr) GetType(_ EvalContext) *types.FieldType { return m.t }
+
+func (m *MockExpr) Clone() Expression {
+	cloned := new(MockExpr)
+	cloned.i = m.i
+	cloned.err = m.err
+	if m.t != nil {
+		cloned.t = m.t.Clone()
+	}
+	return cloned
+}
+
 func (m *MockExpr) Equal(ctx EvalContext, e Expression) bool          { return false }
 func (m *MockExpr) IsCorrelated() bool                                { return false }
 func (m *MockExpr) ConstLevel() ConstLevel                            { return ConstNone }

pkg/planner/core/integration_test.go

@@ -2107,8 +2107,8 @@ func TestIssue46556(t *testing.T) {
 		testkit.Rows(`HashJoin 0.00 root  inner join, equal:[eq(Column#5, test.t0.c0)]`,
 			`├─Projection(Build) 0.00 root  <nil>->Column#5`,
 			`│ └─TableDual 0.00 root  rows:0`,
-			`└─TableReader(Probe) 7992.00 root  data:Selection`,
-			`  └─Selection 7992.00 cop[tikv]  like(test.t0.c0, test.t0.c0, 92), not(isnull(test.t0.c0))`,
+			`└─TableReader(Probe) 9990.00 root  data:Selection`,
+			`  └─Selection 9990.00 cop[tikv]  not(isnull(test.t0.c0))`,
 			`    └─TableFullScan 10000.00 cop[tikv] table:t0 keep order:false, stats:pseudo`))
 }
 

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

@@ -31,7 +31,7 @@ import (
 // Plan Should be used as embedded struct in Plan implementations.
 type Plan struct {
 	ctx     context.PlanContext
-	stats   *property.StatsInfo
+	stats   *property.StatsInfo `plan-cache-clone:"shallow"`
 	tp      string
 	id      int
 	qbBlock int // Query Block offset

pkg/planner/core/physical_plans.go

@@ -159,6 +159,9 @@ const emptyPartitionInfoSize = int64(unsafe.Sizeof(PhysPlanPartInfo{}))
 
 // Clone clones the PhysPlanPartInfo.
 func (pi *PhysPlanPartInfo) Clone() *PhysPlanPartInfo {
+	if pi == nil {
+		return nil
+	}
 	cloned := new(PhysPlanPartInfo)
 	cloned.PruningConds = util.CloneExprs(pi.PruningConds)
 	cloned.PartitionNames = util.CloneCIStrs(pi.PartitionNames)
@@ -411,6 +414,9 @@ type PushedDownLimit struct {
 
 // Clone clones this pushed-down list.
 func (p *PushedDownLimit) Clone() *PushedDownLimit {
+	if p == nil {
+		return nil
+	}
 	cloned := new(PushedDownLimit)
 	*cloned = *p
 	return cloned

pkg/planner/core/plan_cache_rebuild_test.go

@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"math/rand"
 	"reflect"
+	"strings"
 	"testing"
 	"unsafe"
 
@@ -70,6 +71,24 @@ func TestPlanCacheClone(t *testing.T) {
 	testCachedPlanClone(t, tk1, tk2, `prepare st from 'select * from t use index(b) where a<? and b+?=10'`,
 		`set @a1=1, @b1=1, @a2=2, @b2=2`, `execute st using @a1,@b1`, `execute st using @a2,@b2`)
 
+	// IndexLookUp
+	testCachedPlanClone(t, tk1, tk2, `prepare st from 'select * from t use index(b) where b<=?'`,
+		`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>?'`,
+		`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>?'`,
+		`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`)
+	testCachedPlanClone(t, tk1, tk2, `prepare st from 'select * from t where a>=? order by b'`,
+		`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'`,
+		`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'`,
+		`set @a1=1, @a2=2`, `execute st using @a1`, `execute st using @a2`)
+
 	// TODO: PointGet doesn't support Clone
 	// PointPlan
 	//testCachedPlanClone(t, tk1, tk2, `prepare st from 'select * from t where a=?'`,
@@ -92,7 +111,9 @@ func testCachedPlanClone(t *testing.T, tk1, tk2 *testkit.TestKit, prep, set, exe
 	checked := false
 	ctx := context.WithValue(context.Background(), core.PlanCacheKeyTestClone{}, func(plan, cloned base.Plan) {
 		checked = true
-		// TODO: check cloned is deeply cloned from plan.
+		require.NoError(t, checkUnclearPlanCacheClone(plan, cloned,
+			".ctx",
+			"*collate"))
 	})
 	tk2.MustQueryWithContext(ctx, exec2)
 	require.True(t, checked)
@@ -169,11 +190,16 @@ func TestCheckPlanClone(t *testing.T) {
 
 // checkUnclearPlanCacheClone checks whether this cloned plan is safe for instance plan cache.
 // All fields in the plan should be deeply cloned except the fields with tag "plan-cache-shallow-clone:'true'".
-func checkUnclearPlanCacheClone(plan, cloned any) error {
-	return planCacheUnclearCloneCheck(reflect.ValueOf(plan), reflect.ValueOf(cloned), reflect.TypeOf(plan).String(), nil)
+func checkUnclearPlanCacheClone(plan, cloned any, whiteLists ...string) error {
+	return planCacheUnclearCloneCheck(reflect.ValueOf(plan), reflect.ValueOf(cloned), reflect.TypeOf(plan).String(), nil, whiteLists...)
 }
 
-func planCacheUnclearCloneCheck(v1, v2 reflect.Value, path string, visited map[visit]bool) error {
+func planCacheUnclearCloneCheck(v1, v2 reflect.Value, path string, visited map[visit]bool, whiteLists ...string) error {
+	for _, l := range whiteLists {
+		if strings.Contains(path, l) {
+			return nil
+		}
+	}
 	if !v1.IsValid() || !v2.IsValid() {
 		if v1.IsValid() != v2.IsValid() {
 			return errors.Errorf("invalid")
@@ -212,7 +238,7 @@ func planCacheUnclearCloneCheck(v1, v2 reflect.Value, path string, visited map[v
 	switch v1.Kind() {
 	case reflect.Array:
 		for i := 0; i < v1.Len(); i++ {
-			if err := planCacheUnclearCloneCheck(v1.Index(i), v2.Index(i), fmt.Sprintf("%v[%v]", path, i), visited); err != nil {
+			if err := planCacheUnclearCloneCheck(v1.Index(i), v2.Index(i), fmt.Sprintf("%v[%v]", path, i), visited, whiteLists...); err != nil {
 				return err
 			}
 		}
@@ -233,7 +259,7 @@ func planCacheUnclearCloneCheck(v1, v2 reflect.Value, path string, visited map[v
 			return errors.Errorf("same slice pointers, path %v", path)
 		}
 		for i := 0; i < v1.Len(); i++ {
-			if err := planCacheUnclearCloneCheck(v1.Index(i), v2.Index(i), fmt.Sprintf("%v[%v]", path, i), visited); err != nil {
+			if err := planCacheUnclearCloneCheck(v1.Index(i), v2.Index(i), fmt.Sprintf("%v[%v]", path, i), visited, whiteLists...); err != nil {
 				return err
 			}
 		}
@@ -244,23 +270,23 @@ func planCacheUnclearCloneCheck(v1, v2 reflect.Value, path string, visited map[v
 		if v1.IsNil() != v2.IsNil() {
 			return errors.Errorf("invalid interfaces, path %v", path)
 		}
-		return planCacheUnclearCloneCheck(v1.Elem(), v2.Elem(), fmt.Sprintf("%v(%v)", path, v1.Elem().Type().String()), visited)
+		return planCacheUnclearCloneCheck(v1.Elem(), v2.Elem(), fmt.Sprintf("%v(%v)", path, v1.Elem().Type().String()), visited, whiteLists...)
 	case reflect.Ptr:
 		if v1.IsNil() && v2.IsNil() {
 			return nil
 		}
 		if v1.Pointer() == v2.Pointer() {
 			return errors.Errorf("same pointer, path %v", path)
 		}
-		return planCacheUnclearCloneCheck(v1.Elem(), v2.Elem(), path, visited)
+		return planCacheUnclearCloneCheck(v1.Elem(), v2.Elem(), path, visited, whiteLists...)
 	case reflect.Struct:
 		for i, n := 0, v1.NumField(); i < n; i++ {
 			tag := v1.Type().Field(i).Tag.Get("plan-cache-clone")
 			if tag == "shallow" {
 				continue
 			}
 			fieldName := v1.Type().Field(i).Name
-			if err := planCacheUnclearCloneCheck(v1.Field(i), v2.Field(i), fmt.Sprintf("%v.%v", path, fieldName), visited); err != nil {
+			if err := planCacheUnclearCloneCheck(v1.Field(i), v2.Field(i), fmt.Sprintf("%v.%v", path, fieldName), visited, whiteLists...); err != nil {
 				return err
 			}
 		}
@@ -283,7 +309,7 @@ func planCacheUnclearCloneCheck(v1, v2 reflect.Value, path string, visited map[v
 			if !val1.IsValid() || !val2.IsValid() {
 				return errors.Errorf("invalid map value at %v", fmt.Sprintf("%v[%v]", path, k.Type().Name()))
 			}
-			if err := planCacheUnclearCloneCheck(val1, val2, fmt.Sprintf("%v[%v]", path, k.Type().Name()), visited); err != nil {
+			if err := planCacheUnclearCloneCheck(val1, val2, fmt.Sprintf("%v[%v]", path, k.Type().Name()), visited, whiteLists...); err != nil {
 				return err
 			}
 		}

pkg/planner/core/plan_cache_utils.go

@@ -418,8 +418,8 @@ func (v *PlanCacheValue) CloneForInstancePlanCache(ctx context.Context, newCtx b
 	if !ok {
 		return nil, false
 	}
-	clonedPlan, err := phyPlan.Clone(newCtx)
-	if err != nil {
+	clonedPlan, ok := phyPlan.CloneForPlanCache(newCtx)
+	if !ok {
 		return nil, false
 	}
 	if intest.InTest && ctx.Value(PlanCacheKeyTestClone{}) != nil {

pkg/planner/core/plan_clone_generator.go

@@ -105,11 +105,15 @@ func genPlanCloneForPlanCache(x any) ([]byte, error) {
 		case "[]property.SortItem":
 			c.write("cloned.%v = util.CloneSortItem(op.%v)", f.Name, 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":
 			c.write("cloned.%v = op.%v.Clone()", f.Name, f.Name)
 		case "*expression.Column":
+			c.write("if op.%v != nil {", f.Name)
 			c.write("cloned.%v = op.%v.Clone().(*expression.Column)", f.Name, f.Name)
+			c.write("}")
 		case "base.PhysicalPlan":
 			c.write("%v, ok := op.%v.CloneForPlanCache(newCtx)", f.Name, f.Name)
 			c.write("if !ok {return nil, false}")

pkg/planner/util/misc.go

@@ -34,6 +34,7 @@ func CloneFieldNames(names []*types.FieldName) []*types.FieldName {
 	}
 	cloned := make([]*types.FieldName, len(names))
 	for i, name := range names {
+		cloned[i] = new(types.FieldName)
 		*cloned[i] = *name
 	}
 	return cloned

pkg/util/collate/bin.go

@@ -44,6 +44,11 @@ func (*binCollator) Pattern() WildcardPattern {
 	return &binPattern{}
 }
 
+// Clone implements Collator interface.
+func (*binCollator) Clone() Collator {
+	return new(binCollator)
+}
+
 type derivedBinCollator struct {
 	binCollator
 }
@@ -75,6 +80,11 @@ func (*binPaddingCollator) Pattern() WildcardPattern {
 	return &derivedBinPattern{}
 }
 
+// Clone implements Collator interface.
+func (*binPaddingCollator) Clone() Collator {
+	return new(binPaddingCollator)
+}
+
 type derivedBinPattern struct {
 	patChars []rune
 	patTypes []byte

pkg/util/collate/collate.go

@@ -72,6 +72,8 @@ type Collator interface {
 	KeyWithoutTrimRightSpace(str string) []byte
 	// Pattern get a collation-aware WildcardPattern.
 	Pattern() WildcardPattern
+	// Clone returns a copy of the collator.
+	Clone() Collator
 }
 
 // WildcardPattern is the interface used for wildcard pattern match.

pkg/util/collate/gbk_bin.go

@@ -17,6 +17,7 @@ package collate
 import (
 	"bytes"
 
+	"github.com/pingcap/tidb/pkg/parser/charset"
 	"github.com/pingcap/tidb/pkg/util/hack"
 	"golang.org/x/text/encoding"
 )
@@ -26,6 +27,11 @@ type gbkBinCollator struct {
 	e *encoding.Encoder
 }
 
+// Clone implements Collator interface.
+func (*gbkBinCollator) Clone() Collator {
+	return &gbkBinCollator{charset.NewCustomGBKEncoder()}
+}
+
 // Compare implement Collator interface.
 func (g *gbkBinCollator) Compare(a, b string) int {
 	a = truncateTailingSpace(a)

pkg/util/collate/gbk_chinese_ci.go

@@ -64,6 +64,11 @@ func (*gbkChineseCICollator) Pattern() WildcardPattern {
 	return &gbkChineseCIPattern{}
 }
 
+// Clone implements Collator interface.
+func (*gbkChineseCICollator) Clone() Collator {
+	return new(gbkChineseCICollator)
+}
+
 type gbkChineseCIPattern struct {
 	patChars []rune
 	patTypes []byte

pkg/util/collate/general_ci.go

@@ -62,6 +62,11 @@ func (*generalCICollator) Pattern() WildcardPattern {
 	return &ciPattern{}
 }
 
+// Clone implements Collator interface.
+func (*generalCICollator) Clone() Collator {
+	return new(generalCICollator)
+}
+
 type ciPattern struct {
 	patChars []rune
 	patTypes []byte