diff --git a/distsql/request_builder.go b/distsql/request_builder.go index 439c6ecd8e7fe..8d0fae3a464f1 100644 --- a/distsql/request_builder.go +++ b/distsql/request_builder.go @@ -159,6 +159,7 @@ func (builder *RequestBuilder) SetDAGRequest(dag *tipb.DAGRequest) *RequestBuild if limit != nil && limit.Limit < estimatedRegionRowCount { builder.Request.Concurrency = 1 } + builder.Request.LimitSize = limit.GetLimit() } return builder } diff --git a/distsql/request_builder_test.go b/distsql/request_builder_test.go index 74bdf723216f1..1de6b11dd5244 100644 --- a/distsql/request_builder_test.go +++ b/distsql/request_builder_test.go @@ -676,6 +676,7 @@ func TestScanLimitConcurrency(t *testing.T) { Build() require.NoError(t, err) require.Equal(t, tt.concurrency, actual.Concurrency) + require.Equal(t, actual.LimitSize, tt.limit) }) } } diff --git a/kv/kv.go b/kv/kv.go index 9239cc514c7b2..057c22fb7312b 100644 --- a/kv/kv.go +++ b/kv/kv.go @@ -573,6 +573,8 @@ type Request struct { StoreBatchSize int // ResourceGroupName is the name of the bind resource group. ResourceGroupName string + // LimitSize indicates whether the request is scan and limit + LimitSize uint64 } // CoprRequestAdjuster is used to check and adjust a copr request according to specific rules. diff --git a/store/copr/coprocessor.go b/store/copr/coprocessor.go index ab3d14e47f60a..573f746a6b0d3 100644 --- a/store/copr/coprocessor.go +++ b/store/copr/coprocessor.go @@ -402,7 +402,13 @@ func buildCopTasks(bo *Backoffer, ranges *KeyRanges, opt *buildCopTaskOpt) ([]*c } i = nextI if req.Paging.Enable { - pagingSize = paging.GrowPagingSize(pagingSize, req.Paging.MaxPagingSize) + if req.LimitSize != 0 && req.LimitSize < pagingSize { + // disable paging for small limit. + task.paging = false + task.pagingSize = 0 + } else { + pagingSize = paging.GrowPagingSize(pagingSize, req.Paging.MaxPagingSize) + } } } } diff --git a/store/copr/coprocessor_test.go b/store/copr/coprocessor_test.go index f7b15ebfd682d..c94d441932d8c 100644 --- a/store/copr/coprocessor_test.go +++ b/store/copr/coprocessor_test.go @@ -514,8 +514,6 @@ func TestBuildPagingTasks(t *testing.T) { req := &kv.Request{} req.Paging.Enable = true req.Paging.MinPagingSize = paging.MinPagingSize - flashReq := &kv.Request{} - flashReq.StoreType = kv.TiFlash tasks, err := buildTestCopTasks(bo, cache, buildCopRanges("a", "c"), req, nil) require.NoError(t, err) require.Len(t, tasks, 1) @@ -525,6 +523,37 @@ func TestBuildPagingTasks(t *testing.T) { require.Equal(t, tasks[0].pagingSize, paging.MinPagingSize) } +func TestBuildPagingTasksDisablePagingForSmallLimit(t *testing.T) { + mockClient, cluster, pdClient, err := testutils.NewMockTiKV("", nil) + require.NoError(t, err) + defer func() { + pdClient.Close() + err = mockClient.Close() + require.NoError(t, err) + }() + _, regionIDs, _ := testutils.BootstrapWithMultiRegions(cluster, []byte("g"), []byte("n"), []byte("t")) + + pdCli := tikv.NewCodecPDClient(tikv.ModeTxn, pdClient) + defer pdCli.Close() + + cache := NewRegionCache(tikv.NewRegionCache(pdCli)) + defer cache.Close() + + bo := backoff.NewBackofferWithVars(context.Background(), 3000, nil) + + req := &kv.Request{} + req.Paging.Enable = true + req.Paging.MinPagingSize = paging.MinPagingSize + req.LimitSize = 1 + tasks, err := buildTestCopTasks(bo, cache, buildCopRanges("a", "c"), req, nil) + require.NoError(t, err) + require.Len(t, tasks, 1) + require.Len(t, tasks, 1) + taskEqual(t, tasks[0], regionIDs[0], 0, "a", "c") + require.False(t, tasks[0].paging) + require.Equal(t, tasks[0].pagingSize, uint64(0)) +} + func toCopRange(r kv.KeyRange) *coprocessor.KeyRange { coprRange := coprocessor.KeyRange{} coprRange.Start = r.StartKey