From 10cf1b308b1e06bd6b3a1b8a1932d5648d202781 Mon Sep 17 00:00:00 2001 From: Mattias Jonsson Date: Tue, 20 Dec 2022 00:17:23 +0000 Subject: [PATCH] Less intrusive refactoring (due to only a bug fix). Refactoring should still be done I think, but in a separate PR. --- executor/builder.go | 77 +++++++++++++++++++-------------------- table/table.go | 1 - table/tables/partition.go | 23 ------------ 3 files changed, 37 insertions(+), 64 deletions(-) diff --git a/executor/builder.go b/executor/builder.go index 1fecec7753d62..1f215aaa40931 100644 --- a/executor/builder.go +++ b/executor/builder.go @@ -3511,14 +3511,39 @@ func buildIndexRangeForEachPartition(ctx sessionctx.Context, usedPartitions []ta return nextRange, nil } -func keyColumnsIncludeAllPartitionColumns(keyColumnIDs []int64, pt table.PartitionedTable) bool { - partColIDs := pt.GetPartitionColumnIDs() - for _, id := range keyColumnIDs { - if _, ok := partColIDs[id]; !ok { - return false +func getPartitionKeyColOffsets(keyColIDs []int64, pt table.PartitionedTable) []int { + keyColOffsets := make([]int, len(keyColIDs)) + for i, colID := range keyColIDs { + offset := -1 + for j, col := range pt.Cols() { + if colID == col.ID { + offset = j + break + } + } + if offset == -1 { + return nil } + keyColOffsets[i] = offset } - return true + + pe, err := pt.(interface { + PartitionExpr() (*tables.PartitionExpr, error) + }).PartitionExpr() + if err != nil { + return nil + } + + offsetMap := make(map[int]struct{}) + for _, offset := range keyColOffsets { + offsetMap[offset] = struct{}{} + } + for _, offset := range pe.ColumnOffset { + if _, ok := offsetMap[offset]; !ok { + return nil + } + } + return keyColOffsets } func (builder *dataReaderBuilder) prunePartitionForInnerExecutor(tbl table.Table, schema *expression.Schema, partitionInfo *plannercore.PartitionInfo, @@ -3533,15 +3558,6 @@ func (builder *dataReaderBuilder) prunePartitionForInnerExecutor(tbl table.Table return nil, false, nil, err } - // check whether can runtime prune. - type partitionExpr interface { - PartitionExpr() (*tables.PartitionExpr, error) - } - pe, err := tbl.(partitionExpr).PartitionExpr() - if err != nil { - return nil, false, nil, err - } - // recalculate key column offsets if len(lookUpContent) == 0 { return nil, false, nil, nil @@ -3549,29 +3565,9 @@ func (builder *dataReaderBuilder) prunePartitionForInnerExecutor(tbl table.Table if lookUpContent[0].keyColIDs == nil { return nil, false, nil, plannercore.ErrInternal.GenWithStack("cannot get column IDs when dynamic pruning") } - keyColOffsets := make([]int, len(lookUpContent[0].keyColIDs)) - for i, colID := range lookUpContent[0].keyColIDs { - offset := -1 - for j, col := range partitionTbl.Cols() { - if colID == col.ID { - offset = j - break - } - } - if offset == -1 { - return nil, false, nil, plannercore.ErrInternal.GenWithStack("invalid column offset when dynamic pruning") - } - keyColOffsets[i] = offset - } - - offsetMap := make(map[int]bool) - for _, offset := range keyColOffsets { - offsetMap[offset] = true - } - for _, offset := range pe.ColumnOffset { - if _, ok := offsetMap[offset]; !ok { - return condPruneResult, false, nil, nil - } + keyColOffsets := getPartitionKeyColOffsets(lookUpContent[0].keyColIDs, partitionTbl) + if len(keyColOffsets) == 0 { + return condPruneResult, false, nil, nil } locateKey := make([]types.Datum, len(partitionTbl.Cols())) @@ -4156,8 +4152,9 @@ func (builder *dataReaderBuilder) buildTableReaderForIndexJoin(ctx context.Conte usedPartitions[p.GetPhysicalID()] = p } var kvRanges []kv.KeyRange + keyColOffsets := getPartitionKeyColOffsets(lookUpContents[0].keyColIDs, pt) if v.IsCommonHandle { - if len(lookUpContents) > 0 && keyColumnsIncludeAllPartitionColumns(lookUpContents[0].keyColIDs, pt) { + if len(lookUpContents) > 0 && len(keyColOffsets) > 0 { locateKey := make([]types.Datum, e.Schema().Len()) kvRanges = make([]kv.KeyRange, 0, len(lookUpContents)) // lookUpContentsByPID groups lookUpContents by pid(partition) so that kv ranges for same partition can be merged. @@ -4204,7 +4201,7 @@ func (builder *dataReaderBuilder) buildTableReaderForIndexJoin(ctx context.Conte handles, lookUpContents := dedupHandles(lookUpContents) - if len(lookUpContents) > 0 && keyColumnsIncludeAllPartitionColumns(lookUpContents[0].keyColIDs, pt) { + if len(lookUpContents) > 0 && len(keyColOffsets) > 0 { locateKey := make([]types.Datum, e.Schema().Len()) kvRanges = make([]kv.KeyRange, 0, len(lookUpContents)) for _, content := range lookUpContents { diff --git a/table/table.go b/table/table.go index 396c2a77d75ca..813131df90896 100644 --- a/table/table.go +++ b/table/table.go @@ -246,7 +246,6 @@ type PartitionedTable interface { GetPartitionByRow(sessionctx.Context, []types.Datum) (PhysicalTable, error) GetAllPartitionIDs() []int64 GetPartitionColumnNames() []model.CIStr - GetPartitionColumnIDs() map[int64]struct{} CheckForExchangePartition(ctx sessionctx.Context, pi *model.PartitionInfo, r []types.Datum, pid int64) error } diff --git a/table/tables/partition.go b/table/tables/partition.go index 8e197087d3aad..6a0b315b856e9 100644 --- a/table/tables/partition.go +++ b/table/tables/partition.go @@ -962,29 +962,6 @@ func (t *partitionedTable) GetPartitionColumnNames() []model.CIStr { return colNames } -// GetPartitionColumnIDs returns the column IDs from the partition expression -// TODO: refactor and have the column ids or a hash on PartitionInfo instead -func (t *partitionedTable) GetPartitionColumnIDs() map[int64]struct{} { - var ids map[int64]struct{} - meta := t.Meta() - pi := meta.Partition - if len(pi.Columns) > 0 { - ids = make(map[int64]struct{}, len(pi.Columns)) - for _, name := range pi.Columns { - col := table.FindColLowerCase(t.Cols(), name.L) - ids[col.ID] = struct{}{} - } - return ids - } - - partitionCols := expression.ExtractColumns(t.partitionExpr.Expr) - ids = make(map[int64]struct{}, len(partitionCols)) - for _, col := range partitionCols { - ids[col.ID] = struct{}{} - } - return ids -} - // PartitionRecordKey is exported for test. func PartitionRecordKey(pid int64, handle int64) kv.Key { recordPrefix := tablecodec.GenTableRecordPrefix(pid)