From 0c3338a6445b1e5a7d4c0259d2a8882f94afcc6f Mon Sep 17 00:00:00 2001
From: xuhuaiyu <391585975@qq.com>
Date: Mon, 20 Feb 2023 17:44:39 +0800
Subject: [PATCH] distsql: refine code for setting LimitSize of kvRequest

---
 distsql/request_builder.go | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/distsql/request_builder.go b/distsql/request_builder.go
index 8d0fae3a464f1..772b20cac8180 100644
--- a/distsql/request_builder.go
+++ b/distsql/request_builder.go
@@ -153,13 +153,15 @@ func (builder *RequestBuilder) SetDAGRequest(dag *tipb.DAGRequest) *RequestBuild
 		builder.Request.Cacheable = true
 		builder.Request.Data, builder.err = dag.Marshal()
 	}
-	// When the DAG is just simple scan and small limit, set concurrency to 1 would be sufficient.
-	if len(dag.Executors) == 2 && dag.Executors[1].GetLimit() != nil {
-		limit := dag.Executors[1].GetLimit()
-		if limit != nil && limit.Limit < estimatedRegionRowCount {
-			builder.Request.Concurrency = 1
-		}
+	if execCnt := len(dag.Executors); execCnt != 0 && dag.Executors[execCnt-1].GetLimit() != nil {
+		limit := dag.Executors[execCnt-1].GetLimit()
 		builder.Request.LimitSize = limit.GetLimit()
+		// When the DAG is just simple scan and small limit, set concurrency to 1 would be sufficient.
+		if execCnt == 2 {
+			if limit.Limit < estimatedRegionRowCount {
+				builder.Request.Concurrency = 1
+			}
+		}
 	}
 	return builder
 }