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

…mn is part of multi-column index (#51315) (#53486)

close #51313

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
 				}

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

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")

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;

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;