diff --git a/internal/proxy/task_query.go b/internal/proxy/task_query.go index d20295f7a9..578fd2947f 100644 --- a/internal/proxy/task_query.go +++ b/internal/proxy/task_query.go @@ -139,11 +139,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 @@ -152,16 +147,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 3dd3c2ab55..420d9b226b 100644 --- a/internal/proxy/util.go +++ b/internal/proxy/util.go @@ -82,9 +82,24 @@ func isNumber(c uint8) bool { return true } -func validateTopKLimit(limit int64) error { - if limit <= 0 || limit > Params.QuotaConfig.TopKLimit { - return fmt.Errorf("should be in range [1, %d], but got %d", Params.QuotaConfig.TopKLimit, limit) +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 + if depth <= 0 || depth > Params.QuotaConfig.MaxQueryResultWindow { + return fmt.Errorf("(offset+limit) should be in range [1, %d], but got %d", Params.QuotaConfig.MaxQueryResultWindow, depth) + } + return nil +} + +func validateTopKLimit(topK int64) error { + if topK <= 0 || topK > Params.QuotaConfig.TopKLimit { + return fmt.Errorf("top k should be in range [1, %d], but got %d", Params.QuotaConfig.TopKLimit, topK) } return nil } diff --git a/internal/proxy/util_test.go b/internal/proxy/util_test.go index cc71bf84da..89bec2b021 100644 --- a/internal/proxy/util_test.go +++ b/internal/proxy/util_test.go @@ -1058,3 +1058,12 @@ func Test_TopKLimit(t *testing.T) { assert.Error(t, validateTopKLimit(16385)) assert.Error(t, validateTopKLimit(0)) } + +func Test_MaxQueryResultWindow(t *testing.T) { + Params.InitOnce() + 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)) +} diff --git a/internal/util/paramtable/quota_param.go b/internal/util/paramtable/quota_param.go index 98d2196393..08d6ac30d7 100644 --- a/internal/util/paramtable/quota_param.go +++ b/internal/util/paramtable/quota_param.go @@ -91,6 +91,9 @@ type quotaConfig struct { // limits MaxCollectionNum int MaxCollectionNumPerDB int + // Query limit, which applies on: + // maximum of offset + limit + MaxQueryResultWindow int64 // Search limit, which applies on: // maximum # of results to return (topK), and // maximum # of search requests (nq). @@ -171,6 +174,7 @@ func (p *quotaConfig) init(base *BaseTable) { // limits p.initMaxCollectionNum() p.initMaxCollectionNumPerDB() + p.initMaxQueryResultWindow() p.initTopKLimit() p.initNQLimit() @@ -622,6 +626,10 @@ func (p *quotaConfig) initMaxCollectionNumPerDB() { p.MaxCollectionNumPerDB = p.Base.ParseIntWithDefault("quotaAndLimits.limits.maxCollectionNumPerDB", 65536) } +func (p *quotaConfig) initMaxQueryResultWindow() { + p.MaxQueryResultWindow = p.Base.ParseInt64WithDefault("quotaAndLimits.limits.maxQueryResultWindow", 16384) +} + func (p *quotaConfig) initTopKLimit() { p.TopKLimit = p.Base.ParseInt64WithDefault("quotaAndLimits.limits.topK", 16384) }