From cfa16a3fc3efbb52bfa89c70758ea498e6331b5f Mon Sep 17 00:00:00 2001
From: Yuanjia Zhang <zhangyuanjia@pingcap.com>
Date: Tue, 26 Nov 2024 10:51:03 +0800
Subject: [PATCH] planner: fix the wrong schema.Clone implementation (#57632)

close pingcap/tidb#57635
---
 pkg/expression/expression.go  |  2 +-
 pkg/expression/schema.go      | 28 +++++++++++++++++++++++-----
 pkg/expression/schema_test.go | 21 +++++++++++++++++++--
 3 files changed, 43 insertions(+), 8 deletions(-)

diff --git a/pkg/expression/expression.go b/pkg/expression/expression.go
index 0d0d4dd3b2..0eb654a0c7 100644
--- a/pkg/expression/expression.go
+++ b/pkg/expression/expression.go
@@ -1071,7 +1071,7 @@ func TableInfo2SchemaAndNames(ctx BuildContext, dbName pmodel.CIStr, tbl *model.
 		}
 	}
 	schema := NewSchema(cols...)
-	schema.SetUniqueKeys(keys)
+	schema.SetKeys(keys)
 	return schema, names, nil
 }
 
diff --git a/pkg/expression/schema.go b/pkg/expression/schema.go
index b85524db84..04eb189e55 100644
--- a/pkg/expression/schema.go
+++ b/pkg/expression/schema.go
@@ -58,11 +58,17 @@ func (s *Schema) String() string {
 	for _, col := range s.Columns {
 		colStrs = append(colStrs, col.String())
 	}
-	ukStrs := make([]string, 0, len(s.Keys))
+	strs := make([]string, 0, len(s.Keys))
 	for _, key := range s.Keys {
+		strs = append(strs, key.String())
+	}
+	ukStrs := make([]string, 0, len(s.Keys))
+	for _, key := range s.UniqueKeys {
 		ukStrs = append(ukStrs, key.String())
 	}
-	return "Column: [" + strings.Join(colStrs, ",") + "] Unique key: [" + strings.Join(ukStrs, ",") + "]"
+	return "Column: [" + strings.Join(colStrs, ",") +
+		"] Key: [" + strings.Join(strs, ",") +
+		"] Unique key: [" + strings.Join(ukStrs, ",") + "]"
 }
 
 // Clone copies the total schema.
@@ -79,7 +85,14 @@ func (s *Schema) Clone() *Schema {
 		keys = append(keys, key.Clone())
 	}
 	schema := NewSchema(cols...)
-	schema.SetUniqueKeys(keys)
+	schema.SetKeys(keys)
+	if s.UniqueKeys != nil {
+		uniqueKeys := make([]KeyInfo, 0, len(s.UniqueKeys))
+		for _, key := range s.UniqueKeys {
+			uniqueKeys = append(uniqueKeys, key.Clone())
+		}
+		schema.SetUniqueKeys(uniqueKeys)
+	}
 	return schema
 }
 
@@ -195,11 +208,16 @@ func (s *Schema) Append(col ...*Column) {
 	s.Columns = append(s.Columns, col...)
 }
 
-// SetUniqueKeys will set the value of Schema.Keys.
-func (s *Schema) SetUniqueKeys(keys []KeyInfo) {
+// SetKeys will set the value of Schema.Keys.
+func (s *Schema) SetKeys(keys []KeyInfo) {
 	s.Keys = keys
 }
 
+// SetUniqueKeys will set the value of Schema.UniqueKeys.
+func (s *Schema) SetUniqueKeys(keys []KeyInfo) {
+	s.UniqueKeys = keys
+}
+
 // ColumnsIndices will return a slice which contains the position of each column in schema.
 // If there is one column that doesn't match, nil will be returned.
 func (s *Schema) ColumnsIndices(cols []*Column) (ret []int) {
diff --git a/pkg/expression/schema_test.go b/pkg/expression/schema_test.go
index 73d1af98fe..6d9e45c7cc 100644
--- a/pkg/expression/schema_test.go
+++ b/pkg/expression/schema_test.go
@@ -48,12 +48,29 @@ func (s *schemaGenerator) generateSchema(colCount int) *Schema {
 	return NewSchema(cols...)
 }
 
+func TestSchemaClone(t *testing.T) {
+	s := &schemaGenerator{}
+	schema := s.generateSchema(5)
+	generateKeys4Schema(schema)
+
+	uniKeys := make([]KeyInfo, 0, len(schema.Columns)-1)
+	for i := 0; i < len(schema.Columns)-1; i++ {
+		uniKeys = append(uniKeys, []*Column{schema.Columns[i]})
+	}
+	schema.SetUniqueKeys(uniKeys)
+
+	clonedSchema := schema.Clone()
+	require.Equal(t, schema.String(), clonedSchema.String())
+	require.NotSame(t, schema.Keys, clonedSchema.Keys)
+	require.NotSame(t, schema.UniqueKeys, clonedSchema.UniqueKeys)
+}
+
 func TestSchemaString(t *testing.T) {
 	s := &schemaGenerator{}
 	schema := s.generateSchema(5)
-	require.Equal(t, "Column: [Column#1,Column#2,Column#3,Column#4,Column#5] Unique key: []", schema.String())
+	require.Equal(t, "Column: [Column#1,Column#2,Column#3,Column#4,Column#5] Key: [] Unique key: []", schema.String())
 	generateKeys4Schema(schema)
-	require.Equal(t, "Column: [Column#1,Column#2,Column#3,Column#4,Column#5] Unique key: [[Column#1],[Column#2],[Column#3],[Column#4]]", schema.String())
+	require.Equal(t, "Column: [Column#1,Column#2,Column#3,Column#4,Column#5] Key: [[Column#1],[Column#2],[Column#3],[Column#4]] Unique key: []", schema.String())
 }
 
 func TestSchemaRetrieveColumn(t *testing.T) {