diff --git a/planner/core/point_get_plan.go b/planner/core/point_get_plan.go index e40d93edfc168..bb215c5e8df2b 100644 --- a/planner/core/point_get_plan.go +++ b/planner/core/point_get_plan.go @@ -462,9 +462,12 @@ func newBatchPointGetPlan( names []*types.FieldName, whereColNames []string, ) *BatchPointGetPlan { statsInfo := &property.StatsInfo{RowCount: float64(len(patternInExpr.List))} - partitionColName := getHashPartitionColumnName(ctx, tbl) - if tbl.GetPartitionInfo() != nil && partitionColName == nil { - return nil + var partitionColName *ast.ColumnName + if tbl.GetPartitionInfo() != nil { + partitionColName = getHashPartitionColumnName(ctx, tbl) + if partitionColName == nil { + return nil + } } if handleCol != nil { var handles = make([]kv.Handle, len(patternInExpr.List)) @@ -606,6 +609,10 @@ func tryWhereIn2BatchPointGet(ctx sessionctx.Context, selStmt *ast.SelectStmt) * if tbl == nil { return nil } + // Skip the optimization with partition selection. + if len(tblName.PartitionNames) > 0 { + return nil + } for _, col := range tbl.Columns { if col.IsGenerated() || col.State != model.StatePublic { @@ -734,6 +741,14 @@ func tryPointGetPlan(ctx sessionctx.Context, selStmt *ast.SelectStmt) *PointGetP if partitionInfo == nil { return nil } + // Take partition selection into consideration. + if len(tblName.PartitionNames) > 0 { + if !partitionNameInSet(partitionInfo.Name, tblName.PartitionNames) { + p := newPointGetPlan(ctx, tblName.Schema.O, schema, tbl, names) + p.IsTableDual = true + return p + } + } } handlePair, fieldType := findPKHandle(tbl, pairs) @@ -782,6 +797,16 @@ func tryPointGetPlan(ctx sessionctx.Context, selStmt *ast.SelectStmt) *PointGetP return nil } +func partitionNameInSet(name model.CIStr, pnames []model.CIStr) bool { + for _, pname := range pnames { + // Case insensitive, create table partition p0, query using P0 is OK. + if name.L == pname.L { + return true + } + } + return false +} + func newPointGetPlan(ctx sessionctx.Context, dbName string, schema *expression.Schema, tbl *model.TableInfo, names []*types.FieldName) *PointGetPlan { p := &PointGetPlan{ basePlan: newBasePlan(ctx, plancodec.TypePointGet, 0), diff --git a/planner/core/point_get_plan_test.go b/planner/core/point_get_plan_test.go index 401259bae257c..ed9451bdf6c76 100644 --- a/planner/core/point_get_plan_test.go +++ b/planner/core/point_get_plan_test.go @@ -429,3 +429,22 @@ func (s *testPointGetSuite) TestBatchPointGetPartition(c *C) { tk.MustExec("delete from t where (a,b) in ((1,1),(2,2),(3,3),(4,4))") tk.MustQuery("select * from t where (a, b) in ((1, 1), (2, 2), (3, 3), (4, 4))").Check(testkit.Rows()) } + +func (s *testPointGetSuite) TestIssue19141(c *C) { + // For issue 19141, fix partition selection on batch point get. + tk := testkit.NewTestKit(c, s.store) + tk.MustExec("use test") + tk.MustExec("create table t19141 (c_int int, primary key (c_int)) partition by hash ( c_int ) partitions 4") + tk.MustExec("insert into t19141 values (1), (2), (3), (4)") + tk.MustQuery("select * from t19141 partition (p0)").Check(testkit.Rows("4")) + tk.MustQuery("select * from t19141 partition (p0) where c_int = 1").Check(testkit.Rows()) + tk.MustExec("update t19141 partition (p0) set c_int = -c_int where c_int = 1") // TableDual after partition selection. + tk.MustQuery("select * from t19141 order by c_int").Check(testkit.Rows("1", "2", "3", "4")) + + // Bach point get + tk.MustQuery("select * from t19141 partition (p0, p2) where c_int in (1,2,3)").Check(testkit.Rows("2")) + tk.MustExec("update t19141 partition (p1) set c_int = -c_int where c_int in (2,3)") // No data changed + tk.MustQuery("select * from t19141 order by c_int").Check(testkit.Rows("1", "2", "3", "4")) + tk.MustExec("delete from t19141 partition (p0) where c_int in (2,3)") // No data changed + tk.MustQuery("select * from t19141 order by c_int").Check(testkit.Rows("1", "2", "3", "4")) +}