From d4ec23a5d1835a4f9dfa404dfc6ef8462a1ad88d Mon Sep 17 00:00:00 2001 From: Ti Chi Robot Date: Wed, 22 May 2024 18:53:47 +0800 Subject: [PATCH] executor: fix `BatchPoint` leads to tidb panic when KeyPartition column is part of multi-column index (#51315) (#53486) close pingcap/tidb#51313 --- pkg/executor/batch_point_get.go | 6 +++--- pkg/planner/core/initialize.go | 4 ++-- pkg/planner/core/util.go | 13 ++++++------- .../r/executor/batch_point_get.result | 12 ++++++++++++ .../integrationtest/t/executor/batch_point_get.test | 9 +++++++++ 5 files changed, 32 insertions(+), 12 deletions(-) create mode 100644 tests/integrationtest/r/executor/batch_point_get.result create mode 100644 tests/integrationtest/t/executor/batch_point_get.test diff --git a/pkg/executor/batch_point_get.go b/pkg/executor/batch_point_get.go index 06afa7285fa8f..5ac5d3d7fbfd7 100644 --- a/pkg/executor/batch_point_get.go +++ b/pkg/executor/batch_point_get.go @@ -229,7 +229,7 @@ func (e *BatchPointGetExec) initialize(ctx context.Context) error { if len(e.planPhysIDs) > 0 { physID = e.planPhysIDs[i] } else { - physID, err = core.GetPhysID(e.tblInfo, e.partExpr, e.partPos, idxVals[e.partPos]) + physID, err = core.GetPhysID(e.tblInfo, e.partExpr, idxVals[e.partPos]) if err != nil { continue } @@ -364,7 +364,7 @@ func (e *BatchPointGetExec) initialize(ctx context.Context) error { } else { if handle.IsInt() { d := types.NewIntDatum(handle.IntValue()) - tID, err = core.GetPhysID(e.tblInfo, e.partExpr, e.partPos, d) + tID, err = core.GetPhysID(e.tblInfo, e.partExpr, d) if err != nil { continue } @@ -373,7 +373,7 @@ func (e *BatchPointGetExec) initialize(ctx context.Context) error { if err1 != nil { return err1 } - tID, err = core.GetPhysID(e.tblInfo, e.partExpr, e.partPos, d) + tID, err = core.GetPhysID(e.tblInfo, e.partExpr, d) if err != nil { continue } diff --git a/pkg/planner/core/initialize.go b/pkg/planner/core/initialize.go index f48eb8296c59b..c685f6068e63d 100644 --- a/pkg/planner/core/initialize.go +++ b/pkg/planner/core/initialize.go @@ -569,7 +569,7 @@ func (p *BatchPointGetPlan) Init(ctx sessionctx.Context, stats *property.StatsIn break } } - pid, err := GetPhysID(p.TblInfo, p.PartitionExpr, p.PartitionColPos, d) + pid, err := GetPhysID(p.TblInfo, p.PartitionExpr, d) if err != nil { hasErr = true break @@ -578,7 +578,7 @@ func (p *BatchPointGetPlan) Init(ctx sessionctx.Context, stats *property.StatsIn } } else { for _, idxVals := range p.IndexValues { - pid, err := GetPhysID(p.TblInfo, p.PartitionExpr, p.PartitionColPos, idxVals[p.PartitionColPos]) + pid, err := GetPhysID(p.TblInfo, p.PartitionExpr, idxVals[p.PartitionColPos]) if err != nil { hasErr = true break diff --git a/pkg/planner/core/util.go b/pkg/planner/core/util.go index 1b0a5c7e804f8..af04375f6a0f4 100644 --- a/pkg/planner/core/util.go +++ b/pkg/planner/core/util.go @@ -414,7 +414,7 @@ func clonePhysicalPlan(plans []PhysicalPlan) ([]PhysicalPlan, error) { } // GetPhysID returns the physical table ID. -func GetPhysID(tblInfo *model.TableInfo, partitionExpr *tables.PartitionExpr, colPos int, d types.Datum) (int64, error) { +func GetPhysID(tblInfo *model.TableInfo, partitionExpr *tables.PartitionExpr, d types.Datum) (int64, error) { pi := tblInfo.GetPartitionInfo() if pi == nil { return tblInfo.ID, nil @@ -434,12 +434,11 @@ func GetPhysID(tblInfo *model.TableInfo, partitionExpr *tables.PartitionExpr, co len(pi.Columns) > 1 { return 0, errors.Errorf("unsupported partition type in BatchGet") } - newKeyPartExpr := tables.ForKeyPruning{ - KeyPartCols: []*expression.Column{{ - Index: colPos, - UniqueID: partitionExpr.KeyPartCols[0].UniqueID, - }}, - } + // We need to change the partition column index! + col := &expression.Column{} + *col = *partitionExpr.KeyPartCols[0] + col.Index = 0 + newKeyPartExpr := tables.ForKeyPruning{KeyPartCols: []*expression.Column{col}} partIdx, err := newKeyPartExpr.LocateKeyPartition(pi.Num, []types.Datum{d}) if err != nil { return 0, errors.Errorf("unsupported partition type in BatchGet") diff --git a/tests/integrationtest/r/executor/batch_point_get.result b/tests/integrationtest/r/executor/batch_point_get.result new file mode 100644 index 0000000000000..1f371d4b0136f --- /dev/null +++ b/tests/integrationtest/r/executor/batch_point_get.result @@ -0,0 +1,12 @@ +drop table if exists tkey; +create table tkey (col1 int not null, col2 varchar(32) not null, col3 int not null, unique(col1, col2)) partition by key(col2) partitions 4; +insert into tkey values(1, 'a', 1), (2, 'b', 2); +set session tidb_skip_missing_partition_stats=0; +set session tidb_opt_fix_control = ""; +explain format='brief' select col1, col2, col3 from tkey where col1 = 1 and col2 = 'a' or col1 = 3 and col2 = 'c'; +id estRows task access object operator info +Batch_Point_Get 2.00 root table:tkey, index:col1(col1, col2) keep order:false, desc:false +select col1, col2, col3 from tkey where col1 = 1 and col2 = 'a' or col1 = 3 and col2 = 'c'; +col1 col2 col3 +1 a 1 +drop table tkey; diff --git a/tests/integrationtest/t/executor/batch_point_get.test b/tests/integrationtest/t/executor/batch_point_get.test new file mode 100644 index 0000000000000..b2af4d3d66f78 --- /dev/null +++ b/tests/integrationtest/t/executor/batch_point_get.test @@ -0,0 +1,9 @@ +# test BatchPointGet panic issue(#51313) when KeyPartition column is part of multiColumn index. +drop table if exists tkey; +create table tkey (col1 int not null, col2 varchar(32) not null, col3 int not null, unique(col1, col2)) partition by key(col2) partitions 4; +insert into tkey values(1, 'a', 1), (2, 'b', 2); +set session tidb_skip_missing_partition_stats=0; +set session tidb_opt_fix_control = ""; +explain format='brief' select col1, col2, col3 from tkey where col1 = 1 and col2 = 'a' or col1 = 3 and col2 = 'c'; +select col1, col2, col3 from tkey where col1 = 1 and col2 = 'a' or col1 = 3 and col2 = 'c'; +drop table tkey;