From e7a86e0eeb2cdb44213764809ab5342d1b3ecd3e Mon Sep 17 00:00:00 2001
From: bb7133 <bb7133@gmail.com>
Date: Mon, 28 Mar 2022 00:31:08 +0800
Subject: [PATCH] executor: fix KVRange bug for index join with dynamic
 partition pruning

---
 executor/builder.go                   |  8 +++++++-
 planner/core/partition_pruner_test.go | 19 +++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/executor/builder.go b/executor/builder.go
index f5a4a629c3660..b602616c703d6 100644
--- a/executor/builder.go
+++ b/executor/builder.go
@@ -3921,6 +3921,8 @@ func (builder *dataReaderBuilder) buildTableReaderForIndexJoin(ctx context.Conte
 		if len(lookUpContents) > 0 && keyColumnsIncludeAllPartitionColumns(lookUpContents[0].keyCols, pe) {
 			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.
+			lookUpContentsByPID := make(map[int64][]*indexJoinLookUpContent)
 			for _, content := range lookUpContents {
 				for i, date := range content.keys {
 					locateKey[content.keyCols[i]] = date
@@ -3933,7 +3935,11 @@ func (builder *dataReaderBuilder) buildTableReaderForIndexJoin(ctx context.Conte
 				if _, ok := usedPartitions[pid]; !ok {
 					continue
 				}
-				tmp, err := buildKvRangesForIndexJoin(e.ctx, pid, -1, []*indexJoinLookUpContent{content}, indexRanges, keyOff2IdxOff, cwc, nil, interruptSignal)
+				lookUpContentsByPID[pid] = append(lookUpContentsByPID[pid], content)
+			}
+			for pid, contents := range lookUpContentsByPID {
+				// buildKvRanges for each partition.
+				tmp, err := buildKvRangesForIndexJoin(e.ctx, pid, -1, contents, indexRanges, keyOff2IdxOff, cwc, nil, interruptSignal)
 				if err != nil {
 					return nil, err
 				}
diff --git a/planner/core/partition_pruner_test.go b/planner/core/partition_pruner_test.go
index ac4673efaa1ff..5853438545d08 100644
--- a/planner/core/partition_pruner_test.go
+++ b/planner/core/partition_pruner_test.go
@@ -718,3 +718,22 @@ func TestIssue32007(t *testing.T) {
 	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"))
 }
+
+func TestIssue33231(t *testing.T) {
+	store, clean := testkit.CreateMockStore(t)
+	defer clean()
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("create database issue33231")
+	tk.MustExec("use issue33231")
+	tk.MustExec("set @@session.tidb_partition_prune_mode = 'dynamic';")
+	tk.MustExec("create table t1 (c_int int, c_str varchar(40), primary key (c_int, c_str) clustered, key(c_int) ) partition by hash (c_int) partitions 4;")
+	tk.MustExec("create table t2 like t1;")
+	tk.MustExec("insert into t1 values(6, 'beautiful curran');")
+	tk.MustExec("insert into t1 values(7, 'epic kalam');")
+	tk.MustExec("insert into t1 values(7, 'affectionate curie');")
+	tk.MustExec("insert into t2 values(6, 'vigorous rhodes');")
+	tk.MustExec("insert into t2 values(7, 'sweet aryabhata');")
+	tk.MustQuery("select /*+ INL_JOIN(t2) */ * from t1, t2 where t1.c_int = t2.c_int and t1.c_str <= t2.c_str and t2.c_int in (6, 7, 6);").
+		Sort().
+		Check(testkit.Rows("6 beautiful curran 6 vigorous rhodes", "7 affectionate curie 7 sweet aryabhata", "7 epic kalam 7 sweet aryabhata"))
+}