executor: fix the wrong output when the whole data of partition table…

… falls in one region
pingcap · Apr 4, 2023 · 998c5f1 · 998c5f1
1 parent c9d322c
commit 998c5f1
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 9 deletions.
diff --git a/executor/builder.go b/executor/builder.go
@@ -4388,7 +4388,7 @@ func (builder *dataReaderBuilder) buildIndexReaderForIndexJoin(ctx context.Conte
 		if err != nil {
 			return nil, err
 		}
-		err = e.open(ctx, kvRanges)
+		err = e.open(ctx, kv.NewNonParitionedKeyRanges(kvRanges))
 		return e, err
 	}
 

diff --git a/executor/distsql.go b/executor/distsql.go
@@ -270,8 +270,9 @@ func (e *IndexReaderExecutor) Open(ctx context.Context) error {
 	}
 
 	sc := e.ctx.GetSessionVars().StmtCtx
-	var kvRanges []kv.KeyRange
+	var kvRanges *kv.KeyRanges
 	if len(e.partitions) > 0 {
+		keyRanges := make([][]kv.KeyRange, 0, len(e.partitions))
 		for _, p := range e.partitions {
 			partRange := e.ranges
 			if pRange, ok := e.partRangeMap[p.GetPhysicalID()]; ok {
@@ -281,10 +282,15 @@ func (e *IndexReaderExecutor) Open(ctx context.Context) error {
 			if err != nil {
 				return err
 			}
-			kvRanges = append(kvRanges, kvRange...)
+			keyRanges = append(keyRanges, kvRange)
+			e.kvRanges = append(e.kvRanges, kvRange...)
 		}
+		kvRanges = kv.NewPartitionedKeyRanges(keyRanges)
 	} else {
-		kvRanges, err = e.buildKeyRanges(sc, e.ranges, e.physicalTableID)
+		var keyRanges []kv.KeyRange
+		keyRanges, err = e.buildKeyRanges(sc, e.ranges, e.physicalTableID)
+		kvRanges = kv.NewNonParitionedKeyRanges(keyRanges)
+		e.kvRanges = keyRanges
 	}
 	if err != nil {
 		return err
@@ -293,7 +299,7 @@ func (e *IndexReaderExecutor) Open(ctx context.Context) error {
 	return e.open(ctx, kvRanges)
 }
 
-func (e *IndexReaderExecutor) open(ctx context.Context, kvRanges []kv.KeyRange) error {
+func (e *IndexReaderExecutor) open(ctx context.Context, kvRanges *kv.KeyRanges) error {
 	var err error
 	if e.corColInFilter {
 		e.dagPB.Executors, err = constructDistExec(e.ctx, e.plans)
@@ -306,7 +312,6 @@ func (e *IndexReaderExecutor) open(ctx context.Context, kvRanges []kv.KeyRange)
 		collExec := true
 		e.dagPB.CollectExecutionSummaries = &collExec
 	}
-	e.kvRanges = kvRanges
 	// Treat temporary table as dummy table, avoid sending distsql request to TiKV.
 	// In a test case IndexReaderExecutor is mocked and e.table is nil.
 	// Avoid sending distsql request to TIKV.
@@ -320,11 +325,11 @@ func (e *IndexReaderExecutor) open(ctx context.Context, kvRanges []kv.KeyRange)
 		e.memTracker = memory.NewTracker(e.id, -1)
 	}
 	e.memTracker.AttachTo(e.ctx.GetSessionVars().StmtCtx.MemTracker)
-	slices.SortFunc(kvRanges, func(i, j kv.KeyRange) bool {
+	kvRanges.SortByFunc(func(i, j kv.KeyRange) bool {
 		return bytes.Compare(i.StartKey, j.StartKey) < 0
 	})
 	var builder distsql.RequestBuilder
-	builder.SetKeyRanges(kvRanges).
+	builder.SetWrappedKeyRanges(kvRanges).
 		SetDAGRequest(e.dagPB).
 		SetStartTS(e.startTS).
 		SetDesc(e.desc).

diff --git a/executor/partition_table_test.go b/executor/partition_table_test.go
@@ -614,7 +614,7 @@ func TestOrderByandLimit(t *testing.T) {
 		x := rand.Intn(1099)
 		y := rand.Intn(2000) + 1
 		queryRangePartition := fmt.Sprintf("select /*+ LIMIT_TO_COP() */ a from trange use index(idx_a) where a > %v order by a limit %v;", x, y)
-		queryHashPartition := fmt.Sprintf("select /*+ LIMIT_TO_COP() */ a from trange use index(idx_a) where a > %v order by a limit %v;", x, y)
+		queryHashPartition := fmt.Sprintf("select /*+ LIMIT_TO_COP() */ a from thash use index(idx_a) where a > %v order by a limit %v;", x, y)
 		queryRegular := fmt.Sprintf("select a from tregular use index(idx_a) where a > %v order by a limit %v;", x, y)
 		require.True(t, tk.HasPlan(queryRangePartition, "IndexReader")) // check if indexReader is used
 		require.True(t, tk.HasPlan(queryHashPartition, "IndexReader"))