From 8215efd94dba209efa3bf0bc4ef80fa57b478b36 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Fri, 27 Dec 2024 15:53:24 +0800 Subject: [PATCH 1/2] This is an automated cherry-pick of #58499 Signed-off-by: ti-chi-bot --- .../core/casetest/partition/BUILD.bazel | 4 ++ .../partition/integration_partition_test.go | 20 ++++++++++ pkg/planner/core/physical_plans.go | 39 ++++++++++++------- 3 files changed, 50 insertions(+), 13 deletions(-) diff --git a/pkg/planner/core/casetest/partition/BUILD.bazel b/pkg/planner/core/casetest/partition/BUILD.bazel index bb7ecd6f0beda..93392a9f28c1d 100644 --- a/pkg/planner/core/casetest/partition/BUILD.bazel +++ b/pkg/planner/core/casetest/partition/BUILD.bazel @@ -10,7 +10,11 @@ go_test( ], data = glob(["testdata/**"]), flaky = True, +<<<<<<< HEAD shard_count = 6, +======= + shard_count = 10, +>>>>>>> b11cce72848 (planner/core: move all extra columns to the end in ExpandVirtualColumn (#58499)) 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 462db903509e5..60ea4baf72517 100644 --- a/pkg/planner/core/physical_plans.go +++ b/pkg/planner/core/physical_plans.go @@ -975,16 +975,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 { @@ -999,10 +1013,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 } From cc31dccfc1330a7f119cb55ce0c72f59c236326e Mon Sep 17 00:00:00 2001 From: Hangjie Mo Date: Tue, 11 Feb 2025 13:24:31 +0800 Subject: [PATCH 2/2] Update BUILD.bazel --- pkg/planner/core/casetest/partition/BUILD.bazel | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/pkg/planner/core/casetest/partition/BUILD.bazel b/pkg/planner/core/casetest/partition/BUILD.bazel index 93392a9f28c1d..0472ac87bda43 100644 --- a/pkg/planner/core/casetest/partition/BUILD.bazel +++ b/pkg/planner/core/casetest/partition/BUILD.bazel @@ -10,11 +10,7 @@ go_test( ], data = glob(["testdata/**"]), flaky = True, -<<<<<<< HEAD - shard_count = 6, -======= - shard_count = 10, ->>>>>>> b11cce72848 (planner/core: move all extra columns to the end in ExpandVirtualColumn (#58499)) + shard_count = 7, deps = [ "//pkg/config", "//pkg/planner/core/internal",