From c1befbba6107df88221841df9dc55293f383df0b Mon Sep 17 00:00:00 2001 From: jiyfhust Date: Mon, 4 Mar 2024 20:53:04 +0800 Subject: [PATCH] executor: fix `BatchPoint` leads to tidb panic when KeyPartition column is part of multi-column index (#51315) 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(-) diff --git a/pkg/executor/batch_point_get.go b/pkg/executor/batch_point_get.go index f9061fe04b4ca..939503caa4527 100644 --- a/pkg/executor/batch_point_get.go +++ b/pkg/executor/batch_point_get.go @@ -240,7 +240,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 } @@ -386,7 +386,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 } @@ -395,7 +395,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 d14615ea2c0fd..6dd282312c477 100644 --- a/pkg/planner/core/initialize.go +++ b/pkg/planner/core/initialize.go @@ -572,7 +572,7 @@ func (p *BatchPointGetPlan) Init(ctx PlanContext, stats *property.StatsInfo, sch 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 @@ -581,7 +581,7 @@ func (p *BatchPointGetPlan) Init(ctx PlanContext, stats *property.StatsInfo, sch } } 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 fca1573d2da9d..485ab59e5e255 100644 --- a/pkg/planner/core/util.go +++ b/pkg/planner/core/util.go @@ -413,7 +413,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 @@ -433,12 +433,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 index 9dbc8478bcce4..ba29094a46f3b 100644 --- a/tests/integrationtest/r/executor/batch_point_get.result +++ b/tests/integrationtest/r/executor/batch_point_get.result @@ -163,3 +163,15 @@ id c 1 a 11 b 21 c +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 index 58dc15a45e4de..6014239a27c01 100644 --- a/tests/integrationtest/t/executor/batch_point_get.test +++ b/tests/integrationtest/t/executor/batch_point_get.test @@ -107,3 +107,12 @@ explain format='brief' select * from t2 where id in (1, 1, 11); --sorted_result select * from t2 where id in (1, 11, 11, 21); +# 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;