diff --git a/pkg/planner/core/casetest/partition/BUILD.bazel b/pkg/planner/core/casetest/partition/BUILD.bazel index 510e6aac3d7ec..0ec1371dfabe4 100644 --- a/pkg/planner/core/casetest/partition/BUILD.bazel +++ b/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/core/internal", diff --git a/pkg/planner/core/casetest/partition/integration_partition_test.go b/pkg/planner/core/casetest/partition/integration_partition_test.go index 811ea9a4e6ca1..1808cc942b0e9 100644 --- a/pkg/planner/core/casetest/partition/integration_partition_test.go +++ b/pkg/planner/core/casetest/partition/integration_partition_test.go @@ -246,3 +246,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)`) +} diff --git a/pkg/planner/core/physical_plans.go b/pkg/planner/core/physical_plans.go index 3ae807b4c015a..405f973641a6b 100644 --- a/pkg/planner/core/physical_plans.go +++ b/pkg/planner/core/physical_plans.go @@ -994,16 +994,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 { @@ -1018,10 +1032,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 }