From 157161bbc2a7ecd6f4cb1fbbb11f2f036e11229d Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Mon, 23 Dec 2024 10:44:52 +0800 Subject: [PATCH 1/5] Handle extra columns in ExpandVirtualColumn --- pkg/planner/core/physical_plans.go | 33 ++++++++++++++++++++---------- 1 file changed, 22 insertions(+), 11 deletions(-) diff --git a/pkg/planner/core/physical_plans.go b/pkg/planner/core/physical_plans.go index 1887ea8d39991..82739166a2941 100644 --- a/pkg/planner/core/physical_plans.go +++ b/pkg/planner/core/physical_plans.go @@ -1061,14 +1061,26 @@ 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] + + oldNumColumns := len(schema.Columns) + numExtraColumns := 0 + for i := oldNumColumns - 1; i >= 0; i-- { + cid := schema.Columns[i].ID + // All columns with negative cid should be placed at last. + if cid < 0 { + 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 +1095,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 92e886789f5d93e4daef46a0e9effe3fc79aa234 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Tue, 24 Dec 2024 14:58:09 +0800 Subject: [PATCH 2/5] Add test --- pkg/planner/core/integration_test.go | 20 ++++++++++++++++++++ pkg/planner/core/physical_plans.go | 4 ++-- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/pkg/planner/core/integration_test.go b/pkg/planner/core/integration_test.go index 38837cc33b38e..cd82049e9a9aa 100644 --- a/pkg/planner/core/integration_test.go +++ b/pkg/planner/core/integration_test.go @@ -2318,3 +2318,23 @@ func TestNestedVirtualGeneratedColumnUpdate(t *testing.T) { tk.MustExec("UPDATE test1 SET col7 = '{\"col10\":\"DDDDD\",\"col9\":[\"abcdefg\"]}';\n") tk.MustExec("DELETE FROM test1 WHERE col1 < 0;\n") } + +// 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 82739166a2941..3cdc80454a105 100644 --- a/pkg/planner/core/physical_plans.go +++ b/pkg/planner/core/physical_plans.go @@ -1059,8 +1059,8 @@ 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) + copyColumn := make([]*model.ColumnInfo, 0, len(columns)) + copyColumn = append(copyColumn, columns...) oldNumColumns := len(schema.Columns) numExtraColumns := 0 From 288e2bd05cb1d9ff8ccf551a0dc5c38b40e4bbd5 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Tue, 24 Dec 2024 16:42:08 +0800 Subject: [PATCH 3/5] Address comment --- .../partition/integration_partition_test.go | 20 +++++++++++++++++++ pkg/planner/core/integration_test.go | 20 ------------------- 2 files changed, 20 insertions(+), 20 deletions(-) diff --git a/pkg/planner/core/casetest/partition/integration_partition_test.go b/pkg/planner/core/casetest/partition/integration_partition_test.go index 8436d1788301b..4c5633f53472a 100644 --- a/pkg/planner/core/casetest/partition/integration_partition_test.go +++ b/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)`) +} diff --git a/pkg/planner/core/integration_test.go b/pkg/planner/core/integration_test.go index cd82049e9a9aa..38837cc33b38e 100644 --- a/pkg/planner/core/integration_test.go +++ b/pkg/planner/core/integration_test.go @@ -2318,23 +2318,3 @@ func TestNestedVirtualGeneratedColumnUpdate(t *testing.T) { tk.MustExec("UPDATE test1 SET col7 = '{\"col10\":\"DDDDD\",\"col9\":[\"abcdefg\"]}';\n") tk.MustExec("DELETE FROM test1 WHERE col1 < 0;\n") } - -// 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)`) -} From f2f1028ae3cf19a67c5ebb75439d57222499093d Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Fri, 27 Dec 2024 13:25:49 +0800 Subject: [PATCH 4/5] Fix error --- pkg/planner/core/casetest/partition/BUILD.bazel | 2 +- pkg/planner/core/physical_plans.go | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pkg/planner/core/casetest/partition/BUILD.bazel b/pkg/planner/core/casetest/partition/BUILD.bazel index b372c30815a03..571ae2469b6d0 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/util/coretestsdk", diff --git a/pkg/planner/core/physical_plans.go b/pkg/planner/core/physical_plans.go index 3cdc80454a105..05f48d31be1fd 100644 --- a/pkg/planner/core/physical_plans.go +++ b/pkg/planner/core/physical_plans.go @@ -1067,7 +1067,7 @@ func ExpandVirtualColumn(columns []*model.ColumnInfo, schema *expression.Schema, for i := oldNumColumns - 1; i >= 0; i-- { cid := schema.Columns[i].ID // All columns with negative cid should be placed at last. - if cid < 0 { + if cid >= 0 { break } numExtraColumns++ From 0f2fc5d94ff9dc81a1cefc17a5b9fa211d73a598 Mon Sep 17 00:00:00 2001 From: Ruihao Chen Date: Fri, 27 Dec 2024 14:54:57 +0800 Subject: [PATCH 5/5] Fix --- pkg/planner/core/physical_plans.go | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/pkg/planner/core/physical_plans.go b/pkg/planner/core/physical_plans.go index 05f48d31be1fd..26363bc24b827 100644 --- a/pkg/planner/core/physical_plans.go +++ b/pkg/planner/core/physical_plans.go @@ -1066,8 +1066,10 @@ func ExpandVirtualColumn(columns []*model.ColumnInfo, schema *expression.Schema, numExtraColumns := 0 for i := oldNumColumns - 1; i >= 0; i-- { cid := schema.Columns[i].ID - // All columns with negative cid should be placed at last. - if cid >= 0 { + // 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++