Skip to content

Commit

Permalink
planner: add batch_point_get access object (#35230)
Browse files Browse the repository at this point in the history
close #32815
  • Loading branch information
ymkzpx authored Jun 21, 2022
1 parent f89466d commit 6cefaa9
Show file tree
Hide file tree
Showing 4 changed files with 192 additions and 29 deletions.
8 changes: 4 additions & 4 deletions executor/partition_table_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3468,8 +3468,8 @@ func TestPartitionTableExplain(t *testing.T) {
"PartitionUnion 2.00 root ",
"├─Batch_Point_Get 1.00 root table:t handle:[1 2], keep order:false, desc:false",
"└─Batch_Point_Get 1.00 root table:t handle:[1 2], keep order:false, desc:false"))
tk.MustQuery(`explain format = 'brief' select * from t where a IN (2,3,4)`).Check(testkit.Rows("Batch_Point_Get 3.00 root table:t handle:[2 3 4], keep order:false, desc:false"))
tk.MustQuery(`explain format = 'brief' select * from t where a IN (2,3)`).Check(testkit.Rows("Batch_Point_Get 2.00 root table:t handle:[2 3], keep order:false, desc:false"))
tk.MustQuery(`explain format = 'brief' select * from t where a IN (2,3,4)`).Check(testkit.Rows("Batch_Point_Get 3.00 root table:t, partition:P0,p1,P2 handle:[2 3 4], keep order:false, desc:false"))
tk.MustQuery(`explain format = 'brief' select * from t where a IN (2,3)`).Check(testkit.Rows("Batch_Point_Get 2.00 root table:t, partition:P0,P2 handle:[2 3], keep order:false, desc:false"))
// above ^^ is for completeness, the below vv is enough for Issue32719
tk.MustQuery(`explain format = 'brief' select * from t where b = 1`).Check(testkit.Rows(
"PartitionUnion 1.00 root ",
Expand Down Expand Up @@ -3553,8 +3553,8 @@ func TestPartitionTableExplain(t *testing.T) {
tk.MustQuery(`explain format = 'brief' select * from t where a = 1 OR a = 2`).Check(testkit.Rows(
"TableReader 2.00 root partition:p1,P2 data:TableRangeScan",
"└─TableRangeScan 2.00 cop[tikv] table:t range:[1,1], [2,2], keep order:false"))
tk.MustQuery(`explain format = 'brief' select * from t where a IN (2,3,4)`).Check(testkit.Rows("Batch_Point_Get 3.00 root table:t handle:[2 3 4], keep order:false, desc:false"))
tk.MustQuery(`explain format = 'brief' select * from t where a IN (2,3)`).Check(testkit.Rows("Batch_Point_Get 2.00 root table:t handle:[2 3], keep order:false, desc:false"))
tk.MustQuery(`explain format = 'brief' select * from t where a IN (2,3,4)`).Check(testkit.Rows("Batch_Point_Get 3.00 root table:t, partition:P0,p1,P2 handle:[2 3 4], keep order:false, desc:false"))
tk.MustQuery(`explain format = 'brief' select * from t where a IN (2,3)`).Check(testkit.Rows("Batch_Point_Get 2.00 root table:t, partition:P0,P2 handle:[2 3], keep order:false, desc:false"))
tk.MustQuery(`explain format = 'brief' select * from t where b = 1`).Check(testkit.Rows(
"IndexReader 1.00 root partition:all index:IndexRangeScan",
"└─IndexRangeScan 1.00 cop[tikv] table:t, index:b(b) range:[1,1], keep order:false"))
Expand Down
16 changes: 16 additions & 0 deletions planner/core/partition_pruner_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -704,6 +704,22 @@ func TestHashPartitionPruning(t *testing.T) {
tk.MustQuery("SELECT col1, COL3 FROM t WHERE COL1 IN (0,14158354938390,0) AND COL3 IN (3522101843073676459,-2846203247576845955,838395691793635638);").Check(testkit.Rows("0 3522101843073676459"))
}

func TestIssue32815(t *testing.T) {
store, clean := testkit.CreateMockStore(t)
defer clean()

tk := testkit.NewTestKit(t, store)
tk.MustExec("set @@tidb_partition_prune_mode='dynamic'")
tk.MustExec("USE test;")
tk.MustExec("DROP TABLE IF EXISTS t;")
tk.MustExec("create table t (a int primary key, b int, key (b)) partition by hash(a) (partition P0, partition p1, partition P2)")
tk.MustExec("insert into t values (1, 1),(2, 2),(3, 3)")
tk.MustQuery("explain select * from t where a IN (1, 2)").Check(testkit.Rows(
"Batch_Point_Get_1 2.00 root table:t, partition:p1,P2 handle:[1 2], keep order:false, desc:false"))
tk.MustQuery("explain select * from t where a IN (1, 2, 1)").Check(testkit.Rows(
"Batch_Point_Get_1 3.00 root table:t, partition:p1,P2 handle:[1 2 1], keep order:false, desc:false"))
}

func TestIssue32007(t *testing.T) {
store, clean := testkit.CreateMockStore(t)
defer clean()
Expand Down
123 changes: 104 additions & 19 deletions planner/core/point_get_plan.go
Original file line number Diff line number Diff line change
Expand Up @@ -270,6 +270,7 @@ type BatchPointGetPlan struct {
dbName string
TblInfo *model.TableInfo
IndexInfo *model.IndexInfo
PartitionInfos []*model.PartitionDefinition
Handles []kv.Handle
HandleType *types.FieldType
HandleParams []*expression.Constant // record all Parameters for Plan-Cache
Expand Down Expand Up @@ -345,11 +346,25 @@ func (p *BatchPointGetPlan) ExplainNormalizedInfo() string {
}

// AccessObject implements physicalScan interface.
func (p *BatchPointGetPlan) AccessObject(_ bool) string {
func (p *BatchPointGetPlan) AccessObject(normalized bool) string {
var buffer strings.Builder
tblName := p.TblInfo.Name.O
buffer.WriteString("table:")
buffer.WriteString(tblName)
if p.PartitionInfos != nil {
if normalized {
buffer.WriteString(", partition:?")
} else {
for i, partitionInfo := range p.PartitionInfos {
if i == 0 {
buffer.WriteString(", partition:")
} else {
buffer.WriteString(",")
}
buffer.WriteString(partitionInfo.Name.O)
}
}
}
if p.IndexInfo != nil {
if p.IndexInfo.Primary && p.TblInfo.IsCommonHandle {
buffer.WriteString(", clustered index:" + p.IndexInfo.Name.O + "(")
Expand Down Expand Up @@ -565,9 +580,13 @@ func newBatchPointGetPlan(
return nil
}
}

if handleCol != nil {
// condition key of where is primary key
var handles = make([]kv.Handle, len(patternInExpr.List))
var handleParams = make([]*expression.Constant, len(patternInExpr.List))
var pos2PartitionDefinition = make(map[int]*model.PartitionDefinition)
partitionInfos := make([]*model.PartitionDefinition, 0, len(patternInExpr.List))
for i, item := range patternInExpr.List {
// SELECT * FROM t WHERE (key) in ((1), (2))
if p, ok := item.(*ast.ParenthesesExpr); ok {
Expand Down Expand Up @@ -600,13 +619,39 @@ func newBatchPointGetPlan(
}
handles[i] = kv.IntHandle(intDatum.GetInt64())
handleParams[i] = con
pairs := []nameValuePair{{colName: handleCol.Name.L, colFieldType: item.GetType(), value: *intDatum, con: con}}
if tbl.GetPartitionInfo() != nil {
tmpPartitionDefinition, _, pos, isTableDual := getPartitionInfo(ctx, tbl, pairs)
if isTableDual {
return nil
}
if tmpPartitionDefinition != nil {
pos2PartitionDefinition[pos] = tmpPartitionDefinition
}
}
}

posArr := make([]int, len(pos2PartitionDefinition))
i := 0
for pos := range pos2PartitionDefinition {
posArr[i] = pos
i++
}
sort.Ints(posArr)
for _, pos := range posArr {
partitionInfos = append(partitionInfos, pos2PartitionDefinition[pos])
}
if len(partitionInfos) == 0 {
partitionInfos = nil
}

return BatchPointGetPlan{
TblInfo: tbl,
Handles: handles,
HandleParams: handleParams,
HandleType: &handleCol.FieldType,
PartitionExpr: partitionExpr,
TblInfo: tbl,
Handles: handles,
HandleParams: handleParams,
HandleType: &handleCol.FieldType,
PartitionExpr: partitionExpr,
PartitionInfos: partitionInfos,
}.Init(ctx, statsInfo, schema, names, 0)
}

Expand Down Expand Up @@ -661,28 +706,34 @@ func newBatchPointGetPlan(

indexValues := make([][]types.Datum, len(patternInExpr.List))
indexValueParams := make([][]*expression.Constant, len(patternInExpr.List))
partitionInfos := make([]*model.PartitionDefinition, 0, len(patternInExpr.List))
var pos2PartitionDefinition = make(map[int]*model.PartitionDefinition)

var indexTypes []*types.FieldType
for i, item := range patternInExpr.List {
// SELECT * FROM t WHERE (key) in ((1), (2))
// SELECT * FROM t WHERE (key) in ((1), (2)) or SELECT * FROM t WHERE (key1, key2) in ((1, 1), (2, 2))
if p, ok := item.(*ast.ParenthesesExpr); ok {
item = p.Expr
}
var values []types.Datum
var valuesParams []*expression.Constant
var pairs []nameValuePair
switch x := item.(type) {
case *ast.RowExpr:
// The `len(values) == len(valuesParams)` should be satisfied in this mode
if len(x.Values) != len(whereColNames) {
return nil
}
values = make([]types.Datum, len(x.Values))
pairs = make([]nameValuePair, 0, len(x.Values))
valuesParams = make([]*expression.Constant, len(x.Values))
initTypes := false
if indexTypes == nil { // only init once
indexTypes = make([]*types.FieldType, len(x.Values))
initTypes = true
}
for index, inner := range x.Values {
// permutations is used to match column and value.
permIndex := permutations[index]
switch innerX := inner.(type) {
case *driver.ValueExpr:
Expand All @@ -691,6 +742,7 @@ func newBatchPointGetPlan(
return nil
}
values[permIndex] = innerX.Datum
pairs = append(pairs, nameValuePair{colName: whereColNames[index], value: innerX.Datum})
case *driver.ParamMarkerExpr:
con, err := expression.ParamMarkerExpression(ctx, innerX, true)
if err != nil {
Expand All @@ -709,6 +761,7 @@ func newBatchPointGetPlan(
if initTypes {
indexTypes[permIndex] = &colInfos[index].FieldType
}
pairs = append(pairs, nameValuePair{colName: whereColNames[index], value: innerX.Datum})
default:
return nil
}
Expand All @@ -724,6 +777,8 @@ func newBatchPointGetPlan(
return nil
}
values = []types.Datum{*dval}
valuesParams = []*expression.Constant{nil}
pairs = append(pairs, nameValuePair{colName: whereColNames[0], value: *dval})
case *driver.ParamMarkerExpr:
if len(whereColNames) != 1 {
return nil
Expand All @@ -745,12 +800,39 @@ func newBatchPointGetPlan(
if indexTypes == nil { // only init once
indexTypes = []*types.FieldType{&colInfos[0].FieldType}
}
pairs = append(pairs, nameValuePair{colName: whereColNames[0], value: *dval})

default:
return nil
}
indexValues[i] = values
indexValueParams[i] = valuesParams
if tbl.GetPartitionInfo() != nil {
tmpPartitionDefinition, _, pos, isTableDual := getPartitionInfo(ctx, tbl, pairs)
if isTableDual {
return nil
}
if tmpPartitionDefinition != nil {
pos2PartitionDefinition[pos] = tmpPartitionDefinition
}
}

}

posArr := make([]int, len(pos2PartitionDefinition))
i := 0
for pos := range pos2PartitionDefinition {
posArr[i] = pos
i++
}
sort.Ints(posArr)
for _, pos := range posArr {
partitionInfos = append(partitionInfos, pos2PartitionDefinition[pos])
}
if len(partitionInfos) == 0 {
partitionInfos = nil
}

return BatchPointGetPlan{
TblInfo: tbl,
IndexInfo: matchIdxInfo,
Expand All @@ -759,6 +841,7 @@ func newBatchPointGetPlan(
IndexColTypes: indexTypes,
PartitionColPos: pos,
PartitionExpr: partitionExpr,
PartitionInfos: partitionInfos,
}.Init(ctx, statsInfo, schema, names, 0)
}

Expand All @@ -768,6 +851,8 @@ func tryWhereIn2BatchPointGet(ctx sessionctx.Context, selStmt *ast.SelectStmt) *
len(selStmt.WindowSpecs) > 0 {
return nil
}
// `expr1 in (1, 2) and expr2 in (1, 2)` isn't PatternInExpr, so it can't use tryWhereIn2BatchPointGet.
// (expr1, expr2) in ((1, 1), (2, 2)) can hit it.
in, ok := selStmt.Where.(*ast.PatternInExpr)
if !ok || in.Not || len(in.List) < 1 {
return nil
Expand Down Expand Up @@ -907,7 +992,7 @@ func tryPointGetPlan(ctx sessionctx.Context, selStmt *ast.SelectStmt, check bool
var partitionInfo *model.PartitionDefinition
var pos int
if pi != nil {
partitionInfo, pos, isTableDual = getPartitionInfo(ctx, tbl, pairs)
partitionInfo, pos, _, isTableDual = getPartitionInfo(ctx, tbl, pairs)
if isTableDual {
p := newPointGetPlan(ctx, tblName.Schema.O, schema, tbl, names)
p.IsTableDual = true
Expand Down Expand Up @@ -1583,35 +1668,35 @@ func buildHandleCols(ctx sessionctx.Context, tbl *model.TableInfo, schema *expre
return &IntHandleCols{col: handleCol}
}

func getPartitionInfo(ctx sessionctx.Context, tbl *model.TableInfo, pairs []nameValuePair) (*model.PartitionDefinition, int, bool) {
func getPartitionInfo(ctx sessionctx.Context, tbl *model.TableInfo, pairs []nameValuePair) (*model.PartitionDefinition, int, int, bool) {
partitionExpr := getPartitionExpr(ctx, tbl)
if partitionExpr == nil {
return nil, 0, false
return nil, 0, 0, false
}

pi := tbl.GetPartitionInfo()
if pi == nil {
return nil, 0, false
return nil, 0, 0, false
}

switch pi.Type {
case model.PartitionTypeHash:
expr := partitionExpr.OrigExpr
col, ok := expr.(*ast.ColumnNameExpr)
if !ok {
return nil, 0, false
return nil, 0, 0, false
}

partitionColName := col.Name
if partitionColName == nil {
return nil, 0, false
return nil, 0, 0, false
}

for i, pair := range pairs {
if partitionColName.Name.L == pair.colName {
val := pair.value.GetInt64()
pos := mathutil.Abs(val % int64(pi.Num))
return &pi.Definitions[pos], i, false
return &pi.Definitions[pos], i, int(pos), false
}
}
case model.PartitionTypeRange:
Expand All @@ -1629,9 +1714,9 @@ func getPartitionInfo(ctx sessionctx.Context, tbl *model.TableInfo, pairs []name
return ranges.Compare(i, val, unsigned) > 0
})
if pos >= 0 && pos < length {
return &pi.Definitions[pos], i, false
return &pi.Definitions[pos], i, pos, false
}
return nil, 0, true
return nil, 0, 0, true
}
}
}
Expand All @@ -1648,15 +1733,15 @@ func getPartitionInfo(ctx sessionctx.Context, tbl *model.TableInfo, pairs []name
isNull := false
pos := partitionExpr.ForListPruning.LocatePartition(val, isNull)
if pos >= 0 {
return &pi.Definitions[pos], i, false
return &pi.Definitions[pos], i, pos, false
}
return nil, 0, true
return nil, 0, 0, true
}
}
}
}
}
return nil, 0, false
return nil, 0, 0, false
}

func findPartitionIdx(idxInfo *model.IndexInfo, pos int, pairs []nameValuePair) int {
Expand Down
Loading

0 comments on commit 6cefaa9

Please sign in to comment.