planner/core: move all extra columns to the end in ExpandVirtualColumn

pkg/planner/core/casetest/partition/BUILD.bazel

@@ -10,7 +10,7 @@ go_test(
     ],
     data = glob(["testdata/**"]),
     flaky = True,
-    shard_count = 9,
+    shard_count = 10,
     deps = [
         "//pkg/config",
         "//pkg/planner/util/coretestsdk",

pkg/planner/core/casetest/partition/integration_partition_test.go

@@ -247,3 +247,23 @@ func TestBatchPointGetPartitionForAccessObject(t *testing.T) {
 		tk.MustQuery(tt).Check(testkit.Rows(output[i].Plan...))
 	}
 }
+
+// Issue 58475
+func TestGeneratedColumnWithPartition(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+
+	tk.MustExec(`
+		CREATE TABLE tp (
+			id int,
+			c1 int,
+			c2 int GENERATED ALWAYS AS (c1) VIRTUAL,
+			KEY idx (id)
+		) PARTITION BY RANGE (id)
+		(PARTITION p0 VALUES LESS THAN (0),
+		PARTITION p1 VALUES LESS THAN (10000))
+	`)
+	tk.MustExec(`INSERT INTO tp (id, c1) VALUES (0, 1)`)
+	tk.MustExec(`select /*+ FORCE_INDEX(tp, idx) */id from tp where c2 = 2 group by id having id in (0)`)
+}

pkg/planner/core/physical_plans.go

@@ -1059,16 +1059,30 @@ func (ts *PhysicalTableScan) ResolveCorrelatedColumns() ([]*ranger.Range, error)
 // ExpandVirtualColumn expands the virtual column's dependent columns to ts's schema and column.
 func ExpandVirtualColumn(columns []*model.ColumnInfo, schema *expression.Schema,
 	colsInfo []*model.ColumnInfo) []*model.ColumnInfo {
-	copyColumn := make([]*model.ColumnInfo, len(columns))
-	copy(copyColumn, columns)
-	var extraColumn *expression.Column
-	var extraColumnModel *model.ColumnInfo
-	if schema.Columns[len(schema.Columns)-1].ID == model.ExtraHandleID {
-		extraColumn = schema.Columns[len(schema.Columns)-1]
-		extraColumnModel = copyColumn[len(copyColumn)-1]
-		schema.Columns = schema.Columns[:len(schema.Columns)-1]
-		copyColumn = copyColumn[:len(copyColumn)-1]
+	copyColumn := make([]*model.ColumnInfo, 0, len(columns))
+	copyColumn = append(copyColumn, columns...)
+
+	oldNumColumns := len(schema.Columns)
+	numExtraColumns := 0
+	for i := oldNumColumns - 1; i >= 0; i-- {
+		cid := schema.Columns[i].ID
+		// Move extra columns to the end.
+		// ExtraRowChecksumID is ignored here since it's treated as an ordinary column.
+		// https://github.com/pingcap/tidb/blob/3c407312a986327bc4876920e70fdd6841b8365f/pkg/util/rowcodec/decoder.go#L206-L222
+		if cid != model.ExtraHandleID && cid != model.ExtraPhysTblID {
+			break
+		}
+		numExtraColumns++
 	}
+
+	extraColumns := make([]*expression.Column, numExtraColumns)
+	copy(extraColumns, schema.Columns[oldNumColumns-numExtraColumns:])
+	schema.Columns = schema.Columns[:oldNumColumns-numExtraColumns]
+
+	extraColumnModels := make([]*model.ColumnInfo, numExtraColumns)
+	copy(extraColumnModels, copyColumn[len(copyColumn)-numExtraColumns:])
+	copyColumn = copyColumn[:len(copyColumn)-numExtraColumns]
+
 	schemaColumns := schema.Columns
 	for _, col := range schemaColumns {
 		if col.VirtualExpr == nil {
@@ -1083,10 +1097,9 @@ func ExpandVirtualColumn(columns []*model.ColumnInfo, schema *expression.Schema,
 			}
 		}
 	}
-	if extraColumn != nil {
-		schema.Columns = append(schema.Columns, extraColumn)
-		copyColumn = append(copyColumn, extraColumnModel) // nozero
-	}
+
+	schema.Columns = append(schema.Columns, extraColumns...)
+	copyColumn = append(copyColumn, extraColumnModels...)
 	return copyColumn
 }