From ccf3f0066fd569476b647db569b4c885d892bc71 Mon Sep 17 00:00:00 2001
From: chyezh <ye.zhen@zilliz.com>
Date: Sun, 25 Jun 2023 14:42:43 +0800
Subject: [PATCH] [Pick] Enable max result window limit (#24986)

Signed-off-by: chyezh <ye.zhen@zilliz.com>
---
 internal/proxy/task_query.go       | 16 +++-------------
 internal/proxy/util.go             | 22 +++++++++++++++++++---
 internal/proxy/util_test.go        |  9 +++++++++
 pkg/util/paramtable/quota_param.go | 10 ++++++++++
 4 files changed, 41 insertions(+), 16 deletions(-)

diff --git a/internal/proxy/task_query.go b/internal/proxy/task_query.go
index 9de6f97e24..28b40c991a 100644
--- a/internal/proxy/task_query.go
+++ b/internal/proxy/task_query.go
@@ -145,11 +145,6 @@ func parseQueryParams(queryParamsPair []*commonpb.KeyValuePair) (*queryParams, e
 	if err != nil {
 		return nil, fmt.Errorf("%s [%s] is invalid", LimitKey, limitStr)
 	}
-	if limit != 0 {
-		if err := validateTopKLimit(limit); err != nil {
-			return nil, fmt.Errorf("%s [%d] is invalid, %w", LimitKey, limit, err)
-		}
-	}
 
 	offsetStr, err := funcutil.GetAttrByKeyFromRepeatedKV(OffsetKey, queryParamsPair)
 	// if offset is provided
@@ -158,16 +153,11 @@ func parseQueryParams(queryParamsPair []*commonpb.KeyValuePair) (*queryParams, e
 		if err != nil {
 			return nil, fmt.Errorf("%s [%s] is invalid", OffsetKey, offsetStr)
 		}
-
-		if offset != 0 {
-			if err := validateTopKLimit(offset); err != nil {
-				return nil, fmt.Errorf("%s [%d] is invalid, %w", OffsetKey, offset, err)
-			}
-		}
 	}
 
-	if err = validateTopKLimit(limit + offset); err != nil {
-		return nil, fmt.Errorf("invalid limit[%d] + offset[%d], %w", limit, offset, err)
+	// validate max result window.
+	if err = validateMaxQueryResultWindow(offset, limit); err != nil {
+		return nil, fmt.Errorf("invalid max query result window, %w", err)
 	}
 
 	return &queryParams{
diff --git a/internal/proxy/util.go b/internal/proxy/util.go
index a5a36474a9..df9b6cd92a 100644
--- a/internal/proxy/util.go
+++ b/internal/proxy/util.go
@@ -81,10 +81,26 @@ func isNumber(c uint8) bool {
 	return true
 }
 
-func validateTopKLimit(limit int64) error {
+func validateMaxQueryResultWindow(offset int64, limit int64) error {
+	if offset < 0 {
+		return fmt.Errorf("%s [%d] is invalid, should be gte than 0", OffsetKey, offset)
+	}
+	if limit <= 0 {
+		return fmt.Errorf("%s [%d] is invalid, should be greater than 0", LimitKey, limit)
+	}
+
+	depth := offset + limit
+	maxQueryResultWindow := Params.QuotaConfig.MaxQueryResultWindow.GetAsInt64()
+	if depth <= 0 || depth > maxQueryResultWindow {
+		return fmt.Errorf("(offset+limit) should be in range [1, %d], but got %d", maxQueryResultWindow, depth)
+	}
+	return nil
+}
+
+func validateTopKLimit(topK int64) error {
 	topKLimit := Params.QuotaConfig.TopKLimit.GetAsInt64()
-	if limit <= 0 || limit > topKLimit {
-		return fmt.Errorf("should be in range [1, %d], but got %d", topKLimit, limit)
+	if topK <= 0 || topK > topKLimit {
+		return fmt.Errorf("top k should be in range [1, %d], but got %d", topKLimit, topK)
 	}
 	return nil
 }
diff --git a/internal/proxy/util_test.go b/internal/proxy/util_test.go
index 46626406a6..68b21c2f97 100644
--- a/internal/proxy/util_test.go
+++ b/internal/proxy/util_test.go
@@ -1766,6 +1766,15 @@ func Test_TopKLimit(t *testing.T) {
 	assert.Error(t, validateTopKLimit(0))
 }
 
+func Test_MaxQueryResultWindow(t *testing.T) {
+	paramtable.Init()
+	assert.Nil(t, validateMaxQueryResultWindow(0, 16384))
+	assert.Nil(t, validateMaxQueryResultWindow(0, 1))
+	assert.Error(t, validateMaxQueryResultWindow(0, 16385))
+	assert.Error(t, validateMaxQueryResultWindow(0, 0))
+	assert.Error(t, validateMaxQueryResultWindow(1, 0))
+}
+
 func Test_GetPartitionProgressFailed(t *testing.T) {
 	qc := types.NewMockQueryCoord(t)
 	qc.EXPECT().ShowPartitions(mock.Anything, mock.Anything).Return(&querypb.ShowPartitionsResponse{
diff --git a/pkg/util/paramtable/quota_param.go b/pkg/util/paramtable/quota_param.go
index 88cfcf2bfe..d0ebb0796d 100644
--- a/pkg/util/paramtable/quota_param.go
+++ b/pkg/util/paramtable/quota_param.go
@@ -91,6 +91,7 @@ type quotaConfig struct {
 	MaxCollectionNumPerDB ParamItem `refreshable:"true"`
 	TopKLimit             ParamItem `refreshable:"true"`
 	NQLimit               ParamItem `refreshable:"true"`
+	MaxQueryResultWindow  ParamItem `refreshable:"true"`
 
 	// limit writing
 	ForceDenyWriting                     ParamItem `refreshable:"true"`
@@ -762,6 +763,15 @@ Check https://milvus.io/docs/limitations.md for more details.`,
 	}
 	p.NQLimit.Init(base.mgr)
 
+	p.MaxQueryResultWindow = ParamItem{
+		Key:          "quotaAndLimits.limits.maxQueryResultWindow",
+		Version:      "2.3.0",
+		DefaultValue: "16384",
+		FallbackKeys: []string{},
+		Doc:          `Query limit, which applies on: maximum of offset + limit`,
+	}
+	p.MaxQueryResultWindow.Init(base.mgr)
+
 	// limit writing
 	p.ForceDenyWriting = ParamItem{
 		Key:          "quotaAndLimits.limitWriting.forceDeny",