partition: Revert support batch_point_get for partition table in dyna…

…mic mode (#45891) (#45902)

close #45889

executor/batch_point_get.go

@@ -51,7 +51,8 @@ type BatchPointGetExec struct {
 	partExpr    *tables.PartitionExpr
 	partPos     int
 	planPhysIDs []int64
-	partTblID   []int64
+	singlePart  bool
+	partTblID   int64
 	idxVals     [][]types.Datum
 	txn         kv.Transaction
 	lock        bool
@@ -234,11 +235,9 @@ func (e *BatchPointGetExec) initialize(ctx context.Context) error {
 				}
 			}
 
-			// If this BatchPointGetExec is built only for the specific table partitions, skip those filters not matching those partitions.
-			if len(e.partTblID) >= 1 {
-				if _, found := slices.BinarySearch(e.partTblID, physID); !found {
-					continue
-				}
+			// If this BatchPointGetExec is built only for the specific table partition, skip those filters not matching this partition.
+			if e.singlePart && e.partTblID != physID {
+				continue
 			}
 			idxKey, err1 := EncodeUniqueIndexKey(e.Ctx(), e.tblInfo, e.idxInfo, idxVals, physID)
 			if err1 != nil && !kv.ErrNotExist.Equal(err1) {
@@ -380,11 +379,9 @@ func (e *BatchPointGetExec) initialize(ctx context.Context) error {
 				}
 			}
 		}
-		// If this BatchPointGetExec is built only for the specific table partitions, skip those handles not matching those partitions.
-		if len(e.partTblID) >= 1 {
-			if _, found := slices.BinarySearch(e.partTblID, tID); !found {
-				continue
-			}
+		// If this BatchPointGetExec is built only for the specific table partition, skip those handles not matching this partition.
+		if e.singlePart && e.partTblID != tID {
+			continue
 		}
 		key := tablecodec.EncodeRowKeyWithHandle(tID, handle)
 		keys = append(keys, key)

executor/builder.go

@@ -5217,6 +5217,7 @@ func (b *executorBuilder) buildBatchPointGet(plan *plannercore.BatchPointGetPlan
 		partExpr:     plan.PartitionExpr,
 		partPos:      plan.PartitionColPos,
 		planPhysIDs:  plan.PartitionIDs,
+		singlePart:   plan.SinglePart,
 		partTblID:    plan.PartTblID,
 		columns:      plan.Columns,
 	}

planner/core/casetest/partition/integration_partition_test.go

@@ -292,6 +292,9 @@ func TestBatchPointGetTablePartition(t *testing.T) {
 	tk.MustExec("create table tlist3(a int, b int, primary key(a)) partition by list(a) (partition p0 values in (0, 1, 2), partition p1 values in (3, 4, 5))")
 	tk.MustExec("insert into tlist3 values(1,0),(2,0),(3,0),(4,0)")
 
+	tk.MustExec("create table issue45889(a int) partition by list(a) (partition p0 values in (0, 1), partition p1 values in (2, 3))")
+	tk.MustExec("insert into issue45889 values (0),(0),(1),(1),(2),(2),(3),(3)")
+
 	var input []string
 	var output []struct {
 		SQL         string

planner/core/casetest/partition/testdata/integration_partition_suite_in.json

@@ -263,7 +263,8 @@
       "select * from tlist3 where a in (1,2) and 1 = 1 order by a desc",
       "select * from tlist3 partition(p0) where a in (1,4)",
       "select * from tlist3 partition(p1) where a in (1,2)",
-      "select * from tlist3 partition(p0,p1) where a in (1,2)"
+      "select * from tlist3 partition(p0,p1) where a in (1,2)",
+      "select _tidb_rowid, a from issue45889 where _tidb_rowid in (7, 8)"
     ]
   },
   {

planner/core/casetest/partition/testdata/partition_pruner_out.json

@@ -740,8 +740,9 @@
         ],
         "Plan": [
           "HashJoin 2.20 root  inner join, equal:[eq(test_partition.t4.id, test_partition.t5.id)]",
-          "├─Selection(Build) 2.00 root  not(isnull(test_partition.t4.id))",
-          "│ └─Batch_Point_Get 2.00 root table:t4 handle:[1 3], keep order:false, desc:false",
+          "├─TableReader(Build) 2.00 root partition:p0 data:Selection",
+          "│ └─Selection 2.00 cop[tikv]  not(isnull(test_partition.t4.id))",
+          "│   └─TableRangeScan 2.00 cop[tikv] table:t4 range:[1,1], [3,3], keep order:false",
           "└─IndexLookUp(Probe) 3.64 root partition:p0,p1 ",
           "  ├─IndexRangeScan(Build) 4.00 cop[tikv] table:t5, index:a(a, b) range:[1 1,1 1], [1 6,1 6], [6 1,6 1], [6 6,6 6], keep order:false",
           "  └─Selection(Probe) 3.64 cop[tikv]  not(isnull(test_partition.t5.id))",
@@ -755,10 +756,12 @@
         ],
         "Plan": [
           "HashJoin 2.00 root  inner join, equal:[eq(test_partition.t4.id, test_partition.t4.id)]",
-          "├─Selection(Build) 2.00 root  not(isnull(test_partition.t4.id))",
-          "│ └─Batch_Point_Get 2.00 root table:t4 handle:[1 4], keep order:false, desc:false",
-          "└─Selection(Probe) 4.00 root  not(isnull(test_partition.t4.id))",
-          "  └─Batch_Point_Get 4.00 root table:t4 handle:[1 3 9 100], keep order:false, desc:false"
+          "├─TableReader(Build) 2.00 root partition:p0 data:Selection",
+          "│ └─Selection 2.00 cop[tikv]  not(isnull(test_partition.t4.id))",
+          "│   └─TableRangeScan 2.00 cop[tikv] table:t1 range:[1,1], [4,4], keep order:false",
+          "└─TableReader(Probe) 4.00 root partition:p0,p1 data:Selection",
+          "  └─Selection 4.00 cop[tikv]  not(isnull(test_partition.t4.id))",
+          "    └─TableRangeScan 4.00 cop[tikv] table:t2 range:[1,1], [3,3], [9,9], [100,100], keep order:false"
         ]
       },
       {
@@ -1373,8 +1376,9 @@
         ],
         "Plan": [
           "HashJoin 2.17 root  inner join, equal:[eq(default_partition.t4.id, default_partition.t5.id)]",
-          "├─Selection(Build) 2.00 root  not(isnull(default_partition.t4.id))",
-          "│ └─Batch_Point_Get 2.00 root table:t4 handle:[1 3], keep order:false, desc:false",
+          "├─TableReader(Build) 2.00 root partition:p0 data:Selection",
+          "│ └─Selection 2.00 cop[tikv]  not(isnull(default_partition.t4.id))",
+          "│   └─TableRangeScan 2.00 cop[tikv] table:t4 range:[1,1], [3,3], keep order:false",
           "└─IndexLookUp(Probe) 3.69 root partition:p0,p1 ",
           "  ├─IndexRangeScan(Build) 4.00 cop[tikv] table:t5, index:a(a, b) range:[1 1,1 1], [1 6,1 6], [6 1,6 1], [6 6,6 6], keep order:false",
           "  └─Selection(Probe) 3.69 cop[tikv]  not(isnull(default_partition.t5.id))",
@@ -1388,10 +1392,12 @@
         ],
         "Plan": [
           "HashJoin 2.00 root  inner join, equal:[eq(default_partition.t4.id, default_partition.t4.id)]",
-          "├─Selection(Build) 2.00 root  not(isnull(default_partition.t4.id))",
-          "│ └─Batch_Point_Get 2.00 root table:t4 handle:[1 4], keep order:false, desc:false",
-          "└─Selection(Probe) 4.00 root  not(isnull(default_partition.t4.id))",
-          "  └─Batch_Point_Get 4.00 root table:t4 handle:[1 3 9 100], keep order:false, desc:false"
+          "├─TableReader(Build) 2.00 root partition:p0 data:Selection",
+          "│ └─Selection 2.00 cop[tikv]  not(isnull(default_partition.t4.id))",
+          "│   └─TableRangeScan 2.00 cop[tikv] table:t1 range:[1,1], [4,4], keep order:false",
+          "└─TableReader(Probe) 4.00 root partition:p0,p1 data:Selection",
+          "  └─Selection 4.00 cop[tikv]  not(isnull(default_partition.t4.id))",
+          "    └─TableRangeScan 4.00 cop[tikv] table:t2 range:[1,1], [3,3], [9,9], [100,100], keep order:false"
         ]
       },
       {

planner/core/find_best_task.go

@@ -1136,12 +1136,19 @@ func (ds *DataSource) findBestTask(prop *property.PhysicalProperty, planCounter
 				}
 			}
 		}
-		var partColName *model.CIStr
+		var hashPartColName *model.CIStr
 		if tblInfo := ds.table.Meta(); canConvertPointGet && tblInfo.GetPartitionInfo() != nil {
+			// partition table with dynamic prune not support batchPointGet
+			if canConvertPointGet && len(path.Ranges) > 1 && ds.SCtx().GetSessionVars().StmtCtx.UseDynamicPartitionPrune() {
+				canConvertPointGet = false
+			}
 			if canConvertPointGet && len(path.Ranges) > 1 {
-				// not support some complex situation, like `by HASH( col DIV 80 )` etc.
-				partColName = getPartitionColumnName(getPartitionExpr(ds.SCtx(), tblInfo), tblInfo)
-				if partColName == nil {
+				// We can only build batch point get for hash partitions on a simple column now. This is
+				// decided by the current implementation of `BatchPointGetExec::initialize()`, specifically,
+				// the `getPhysID()` function. Once we optimize that part, we can come back and enable
+				// BatchPointGet plan for more cases.
+				hashPartColName = getHashOrKeyPartitionColumnName(ds.SCtx(), tblInfo)
+				if hashPartColName == nil {
 					canConvertPointGet = false
 				}
 			}
@@ -1170,7 +1177,7 @@ func (ds *DataSource) findBestTask(prop *property.PhysicalProperty, planCounter
 				if len(path.Ranges) == 1 {
 					pointGetTask = ds.convertToPointGet(prop, candidate)
 				} else {
-					pointGetTask = ds.convertToBatchPointGet(prop, candidate, partColName)
+					pointGetTask = ds.convertToBatchPointGet(prop, candidate, hashPartColName)
 				}
 
 				// Batch/PointGet plans may be over-optimized, like `a>=1(?) and a<=1(?)` --> `a=1` --> PointGet(a=1).
@@ -2526,7 +2533,7 @@ func (ds *DataSource) convertToPointGet(prop *property.PhysicalProperty, candida
 	return rTsk
 }
 
-func (ds *DataSource) convertToBatchPointGet(prop *property.PhysicalProperty, candidate *candidatePath, partColName *model.CIStr) (task task) {
+func (ds *DataSource) convertToBatchPointGet(prop *property.PhysicalProperty, candidate *candidatePath, hashPartColName *model.CIStr) (task task) {
 	if !prop.IsSortItemEmpty() && !candidate.isMatchProp {
 		return invalidTask
 	}
@@ -2543,31 +2550,11 @@ func (ds *DataSource) convertToBatchPointGet(prop *property.PhysicalProperty, ca
 		TblInfo:          ds.TableInfo(),
 		KeepOrder:        !prop.IsSortItemEmpty(),
 		Columns:          ds.Columns,
+		SinglePart:       ds.isPartition,
+		PartTblID:        ds.physicalTableID,
 		PartitionExpr:    getPartitionExpr(ds.SCtx(), ds.TableInfo()),
 	}
-	if ds.isPartition {
-		// static prune
-		batchPointGetPlan.PartTblID = make([]int64, 1)
-		batchPointGetPlan.PartTblID[0] = ds.physicalTableID
-	} else if ds.tableInfo.GetPartitionInfo() != nil {
-		// dynamic prune
-		idxs, err := PartitionPruning(ds.SCtx(), ds.table.GetPartitionedTable(), ds.allConds, ds.partitionNames, ds.TblCols, ds.names)
-		if err != nil || len(idxs) == 0 {
-			return invalidTask
-		}
-		if idxs[0] != FullRange {
-			batchPointGetPlan.PartTblID = make([]int64, len(idxs))
-			for i, idx := range idxs {
-				batchPointGetPlan.PartTblID[i] = ds.tableInfo.GetPartitionInfo().Definitions[idx].ID
-			}
-			slices.Sort(batchPointGetPlan.PartTblID)
-		}
-	}
 	if batchPointGetPlan.KeepOrder {
-		// TODO: support keepOrder for partition table with dynamic pruning
-		if ds.TableInfo().GetPartitionInfo() != nil && ds.SCtx().GetSessionVars().StmtCtx.UseDynamicPruneMode {
-			return invalidTask
-		}
 		batchPointGetPlan.Desc = prop.SortItems[0].Desc
 	}
 	rTsk := &rootTask{}
@@ -2589,7 +2576,7 @@ func (ds *DataSource) convertToBatchPointGet(prop *property.PhysicalProperty, ca
 		batchPointGetPlan.IndexInfo = candidate.path.Index
 		batchPointGetPlan.IdxCols = candidate.path.IdxCols
 		batchPointGetPlan.IdxColLens = candidate.path.IdxColLens
-		batchPointGetPlan.PartitionColPos = getColumnPosInIndex(candidate.path.Index, partColName)
+		batchPointGetPlan.PartitionColPos = getColumnPosInIndex(candidate.path.Index, hashPartColName)
 		for _, ran := range candidate.path.Ranges {
 			batchPointGetPlan.IndexValues = append(batchPointGetPlan.IndexValues, ran.LowVal)
 		}

planner/core/point_get_plan.go

@@ -329,8 +329,13 @@ type BatchPointGetPlan struct {
 	Columns          []*model.ColumnInfo
 	cost             float64
 
-	// PartTblID is the table IDs for the specific table partitions.
-	PartTblID []int64
+	// SinglePart indicates whether this BatchPointGetPlan is just for a single partition, instead of the whole partition table.
+	// If the BatchPointGetPlan is built in fast path, this value is false; if the plan is generated in physical optimization for a partition,
+	// this value would be true. This value would decide the behavior of BatchPointGetExec, i.e, whether to compute the table ID of the partition
+	// on the fly.
+	SinglePart bool
+	// PartTblID is the table ID for the specific table partition.
+	PartTblID int64
 
 	// required by cost model
 	planCostInit bool
@@ -1887,53 +1892,39 @@ func getPartitionColumnPos(idx *model.IndexInfo, partitionExpr *tables.Partition
 		return 0, nil
 	}
 
-	partitionColName := getPartitionColumnName(partitionExpr, tbl)
-	if partitionColName == nil {
-		return 0, errors.Errorf("unsupported partition type in BatchGet")
-	}
-
-	return getColumnPosInIndex(idx, partitionColName), nil
-}
-
-func getPartitionColumnName(partitionExpr *tables.PartitionExpr, tbl *model.TableInfo) *model.CIStr {
-	if partitionExpr == nil {
-		return nil
-	}
-
-	pi := tbl.GetPartitionInfo()
 	var partitionColName model.CIStr
 	switch pi.Type {
 	case model.PartitionTypeHash:
 		col, ok := partitionExpr.OrigExpr.(*ast.ColumnNameExpr)
 		if !ok {
-			return nil
+			return 0, errors.Errorf("unsupported partition type in BatchGet")
 		}
 		partitionColName = col.Name.Name
 	case model.PartitionTypeKey:
 		if len(partitionExpr.KeyPartCols) != 1 {
-			return nil
+			return 0, errors.Errorf("unsupported partition type in BatchGet")
 		}
 		colInfo := findColNameByColID(tbl.Columns, partitionExpr.KeyPartCols[0])
 		partitionColName = colInfo.Name
 	case model.PartitionTypeRange:
 		// left range columns partition for future development
 		col, ok := partitionExpr.Expr.(*expression.Column)
 		if !(ok && len(pi.Columns) == 0) {
-			return nil
+			return 0, errors.Errorf("unsupported partition type in BatchGet")
 		}
 		colInfo := findColNameByColID(tbl.Columns, col)
 		partitionColName = colInfo.Name
 	case model.PartitionTypeList:
 		// left list columns partition for future development
 		locateExpr, ok := partitionExpr.ForListPruning.LocateExpr.(*expression.Column)
 		if !(ok && partitionExpr.ForListPruning.ColPrunes == nil) {
-			return nil
+			return 0, errors.Errorf("unsupported partition type in BatchGet")
 		}
 		colInfo := findColNameByColID(tbl.Columns, locateExpr)
 		partitionColName = colInfo.Name
 	}
 
-	return &partitionColName
+	return getColumnPosInIndex(idx, &partitionColName), nil
 }
 
 // getColumnPosInIndex gets the column's position in the index.
@@ -1966,6 +1957,36 @@ func getPartitionExpr(ctx sessionctx.Context, tbl *model.TableInfo) *tables.Part
 	return partTable.PartitionExpr()
 }
 
+func getHashOrKeyPartitionColumnName(ctx sessionctx.Context, tbl *model.TableInfo) *model.CIStr {
+	pi := tbl.GetPartitionInfo()
+	if pi == nil {
+		return nil
+	}
+	if pi.Type != model.PartitionTypeHash && pi.Type != model.PartitionTypeKey {
+		return nil
+	}
+	is := ctx.GetInfoSchema().(infoschema.InfoSchema)
+	table, ok := is.TableByID(tbl.ID)
+	if !ok {
+		return nil
+	}
+	// PartitionExpr don't need columns and names for hash partition.
+	partitionExpr := table.(partitionTable).PartitionExpr()
+	if pi.Type == model.PartitionTypeKey {
+		// used to judge whether the key partition contains only one field
+		if len(pi.Columns) != 1 {
+			return nil
+		}
+		return &pi.Columns[0]
+	}
+	expr := partitionExpr.OrigExpr
+	col, ok := expr.(*ast.ColumnNameExpr)
+	if !ok {
+		return nil
+	}
+	return &col.Name.Name
+}
+
 func findColNameByColID(cols []*model.ColumnInfo, col *expression.Column) *model.ColumnInfo {
 	for _, c := range cols {
 		if c.ID == col.ID {