*: add explicit partition pruning to index joins (#32007)

executor/builder.go

@@ -3844,6 +3844,23 @@ func (builder *dataReaderBuilder) buildTableReaderForIndexJoin(ctx context.Conte
 		return nil, err
 	}
 	tbInfo := e.table.Meta()
+	tbl, _ := builder.is.TableByID(tbInfo.ID)
+	explicitPartitionSelection := make(map[int64]bool)
+	if pi := tbInfo.GetPartitionInfo(); pi != nil {
+		pt := tbl.(table.PartitionedTable)
+		if len(v.PartitionInfo.PartitionNames) > 0 {
+			for _, partName := range v.PartitionInfo.PartitionNames {
+				partID := pi.GetIDByName(partName.L)
+				if partID != 0 {
+					explicitPartitionSelection[partID] = true
+				}
+			}
+		} else {
+			for _, partID := range pt.GetAllPartitionIDs() {
+				explicitPartitionSelection[partID] = true
+			}
+		}
+	}
 	if v.IsCommonHandle {
 		if tbInfo.GetPartitionInfo() == nil || !builder.ctx.GetSessionVars().UseDynamicPartitionPrune() {
 			kvRanges, err := buildKvRangesForIndexJoin(e.ctx, getPhysicalTableID(e.table), -1, lookUpContents, indexRanges, keyOff2IdxOff, cwc, memTracker, interruptSignal)
@@ -3852,8 +3869,6 @@ func (builder *dataReaderBuilder) buildTableReaderForIndexJoin(ctx context.Conte
 			}
 			return builder.buildTableReaderFromKvRanges(ctx, e, kvRanges)
 		}
-
-		tbl, _ := builder.is.TableByID(tbInfo.ID)
 		pt := tbl.(table.PartitionedTable)
 		pe, err := tbl.(interface {
 			PartitionExpr() (*tables.PartitionExpr, error)
@@ -3875,6 +3890,9 @@ func (builder *dataReaderBuilder) buildTableReaderForIndexJoin(ctx context.Conte
 					return nil, err
 				}
 				pid := p.GetPhysicalID()
+				if _, ok := explicitPartitionSelection[pid]; !ok {
+					continue
+				}
 				tmp, err := buildKvRangesForIndexJoin(e.ctx, pid, -1, []*indexJoinLookUpContent{content}, indexRanges, keyOff2IdxOff, cwc, nil, interruptSignal)
 				if err != nil {
 					return nil, err
@@ -3890,6 +3908,9 @@ func (builder *dataReaderBuilder) buildTableReaderForIndexJoin(ctx context.Conte
 			kvRanges = make([]kv.KeyRange, 0, len(partitions)*len(lookUpContents))
 			for _, p := range partitions {
 				pid := p.GetPhysicalID()
+				if _, ok := explicitPartitionSelection[pid]; !ok {
+					continue
+				}
 				tmp, err := buildKvRangesForIndexJoin(e.ctx, pid, -1, lookUpContents, indexRanges, keyOff2IdxOff, cwc, memTracker, interruptSignal)
 				if err != nil {
 					return nil, err
@@ -3908,7 +3929,6 @@ func (builder *dataReaderBuilder) buildTableReaderForIndexJoin(ctx context.Conte
 		return builder.buildTableReaderFromHandles(ctx, e, handles, canReorderHandles)
 	}
 
-	tbl, _ := builder.is.TableByID(tbInfo.ID)
 	pt := tbl.(table.PartitionedTable)
 	pe, err := tbl.(interface {
 		PartitionExpr() (*tables.PartitionExpr, error)
@@ -3929,6 +3949,9 @@ func (builder *dataReaderBuilder) buildTableReaderForIndexJoin(ctx context.Conte
 				return nil, err
 			}
 			pid := p.GetPhysicalID()
+			if _, ok := explicitPartitionSelection[pid]; !ok {
+				continue
+			}
 			handle := kv.IntHandle(content.keys[0].GetInt64())
 			tmp := distsql.TableHandlesToKVRanges(pid, []kv.Handle{handle})
 			kvRanges = append(kvRanges, tmp...)
@@ -3941,6 +3964,9 @@ func (builder *dataReaderBuilder) buildTableReaderForIndexJoin(ctx context.Conte
 		}
 		for _, p := range partitions {
 			pid := p.GetPhysicalID()
+			if _, ok := explicitPartitionSelection[pid]; !ok {
+				continue
+			}
 			tmp := distsql.TableHandlesToKVRanges(pid, handles)
 			kvRanges = append(kvRanges, tmp...)
 		}

parser/model/model.go

@@ -846,6 +846,18 @@ func (pi *PartitionInfo) GetNameByID(id int64) string {
 	return ""
 }
 
+// GetIDByName gets the partition ID by Name.
+// 0 means Name not found!
+func (pi *PartitionInfo) GetIDByName(name string) int64 {
+	definitions := pi.Definitions
+	for i := range definitions {
+		if definitions[i].Name.L == strings.ToLower(name) {
+			return definitions[i].ID
+		}
+	}
+	return 0
+}
+
 func (pi *PartitionInfo) GetStateByID(id int64) SchemaState {
 	for _, pstate := range pi.States {
 		if pstate.ID == id {

planner/core/exhaust_physical_plans.go

@@ -987,6 +987,7 @@ func (p *LogicalJoin) constructInnerTableScanTask(
 		Columns:        ds.TblCols,
 		ColumnNames:    ds.names,
 	}
+	ts.PartitionInfo = copTask.partitionInfo
 	selStats := ts.stats.Scale(selectivity)
 	ts.addPushedDownSelection(copTask, selStats)
 	t := copTask.convertToRootTask(ds.ctx)

planner/core/partition_pruner_test.go

@@ -587,3 +587,24 @@ func (s *testPartitionPruneSuit) TestHashPartitionPruning(c *C) {
 	tk.MustExec("insert into t(col1, col3) values(0, 3522101843073676459);")
 	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 (s *testPartitionPruneSuit) TestIssue32007(c *C) {
+	tk := testkit.NewTestKit(c, s.store)
+	tk.MustExec("create database Issue32007")
+	defer tk.MustExec("drop database Issue32007")
+	tk.MustExec("USE Issue32007")
+	tk.MustExec("create table t1 (a int, b tinyint, primary key (a)) partition by range (a) (" +
+		"partition p0 values less than (5)," +
+		"partition p1 values less than (20)," +
+		"partition p2 values less than (30)," +
+		"partition p3 values less than (40)," +
+		"partition p4 values less than MAXVALUE)")
+	tk.MustExec("insert into t1 values (0, 0), (1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6), (7, 7), (10, 10), (11, 11), (12, 12), (13, 13), (14, 14), (15, 15), (20, 20), (21, 21), (22, 22), (23, 23), (24, 24), (25, 25), (30, 30), (31, 31), (32, 32), (33, 33), (34, 34), (35, 35), (36, 36), (40, 40), (50, 50), (80, 80), (90, 90), (100, 100)")
+	tk.MustExec("create table t3 (a int, b mediumint, primary key (a))")
+	tk.MustExec("insert into t3 values (0, 0), (1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6), (7, 7), (8, 8), (9, 9), (10, 10), (11, 11), (12, 12), (13, 13), (14, 14), (15, 15), (16, 16), (17, 17), (18, 18), (19, 19), (20, 20), (21, 21), (22, 22), (23, 23)")
+
+	tk.MustExec("set @@tidb_partition_prune_mode='static'")
+	tk.MustQuery("select * from t3 where t3.a <> ALL (select t1.a from t1 partition (p0)) order by t3.a").Sort().Check(testkit.Rows("10 10", "11 11", "12 12", "13 13", "14 14", "15 15", "16 16", "17 17", "18 18", "19 19", "20 20", "21 21", "22 22", "23 23", "5 5", "6 6", "7 7", "8 8", "9 9"))
+	tk.MustExec("set @@tidb_partition_prune_mode='dynamic'")
+	tk.MustQuery("select * from t3 where t3.a <> ALL (select t1.a from t1 partition (p0)) order by t3.a").Sort().Check(testkit.Rows("10 10", "11 11", "12 12", "13 13", "14 14", "15 15", "16 16", "17 17", "18 18", "19 19", "20 20", "21 21", "22 22", "23 23", "5 5", "6 6", "7 7", "8 8", "9 9"))
+}