executor: fix BatchPoint leads to tidb panic when KeyPartition colu…

…mn is part of multi-column index (#51315) close #51313
pingcap · Mar 4, 2024 · c1befbb · c1befbb
1 parent b01942d
commit c1befbb
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 12 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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;