test/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2024-12-03 04:19:18 +08:00
Code Issues Actions 13 Packages Projects Releases Wiki Activity

enhance: use configed max topk for iterator when input topk exceeds(#34292) (#34293)

Browse Source 
related: #34292
pr: https://github.com/milvus-io/milvus/pull/34290

Signed-off-by: MrPresent-Han <chun.han@gmail.com>
Co-authored-by: MrPresent-Han <chun.han@gmail.com>
This commit is contained in:
Chun Han 2024-07-04 09:58:09 +08:00 committed by GitHub
parent bc1746f96c
commit 014cb7b071
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 31 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
internal/proxy/search_util.go
View File

@ -28,6 +28,9 @@ type rankParams struct {
// parseSearchInfo returns QueryInfo and offset
func parseSearchInfo(searchParamsPair []*commonpb.KeyValuePair, schema *schemapb.CollectionSchema, ignoreOffset bool) (*planpb.QueryInfo, int64, error) {
//0. parse iterator field
isIterator, _ := funcutil.GetAttrByKeyFromRepeatedKV(IteratorField, searchParamsPair)
// 1. parse offset and real topk
topKStr, err := funcutil.GetAttrByKeyFromRepeatedKV(TopKKey, searchParamsPair)
if err != nil {
@ -38,7 +41,13 @@ func parseSearchInfo(searchParamsPair []*commonpb.KeyValuePair, schema *schemapb
return nil, 0, fmt.Errorf("%s [%s] is invalid", TopKKey, topKStr)
}
if err := validateLimit(topK); err != nil {
return nil, 0, fmt.Errorf("%s [%d] is invalid, %w", TopKKey, topK, err)
if isIterator == "True" {
topK = Params.QuotaConfig.TopKLimit.GetAsInt64()
//1. if the request is from iterator, we set topK to QuotaLimit as the iterator can resolve too large topK problem
//2. GetAsInt64 has cached inside, no need to worry about cpu cost for parsing here
} else {
return nil, 0, fmt.Errorf("%s [%d] is invalid, %w", TopKKey, topK, err)
}
}
var offset int64
@ -109,8 +118,7 @@ func parseSearchInfo(searchParamsPair []*commonpb.KeyValuePair, schema *schemapb
}
}
// 6. parse iterator tag, prevent trying to groupBy when doing iteration or doing range-search
isIterator, _ := funcutil.GetAttrByKeyFromRepeatedKV(IteratorField, searchParamsPair)
// 6. disable groupBy for iterator and range search
if isIterator == "True" && groupByFieldId > 0 {
return nil, 0, merr.WrapErrParameterInvalid("", "",
"Not allowed to do groupBy when doing iteration")

20
internal/proxy/task_search_test.go
View File

@ -2104,6 +2104,26 @@ func TestTaskSearch_parseQueryInfo(t *testing.T) {
assert.Nil(t, info)
assert.ErrorIs(t, err, merr.ErrParameterInvalid)
})
t.Run("check iterator and topK", func(t *testing.T) {
normalParam := getValidSearchParams()
normalParam = append(normalParam, &commonpb.KeyValuePair{
Key: IteratorField,
Value: "True",
})
resetSearchParamsValue(normalParam, TopKKey, `1024000`)
fields := make([]*schemapb.FieldSchema, 0)
fields = append(fields, &schemapb.FieldSchema{
FieldID: int64(101),
Name: "string_field",
})
schema := &schemapb.CollectionSchema{
Fields: fields,
}
info, _, err := parseSearchInfo(normalParam, schema, false)
assert.NotNil(t, info)
assert.NoError(t, err)
assert.Equal(t, Params.QuotaConfig.TopKLimit.GetAsInt64(), info.Topk)
})
}
func getSearchResultData(nq, topk int64) *schemapb.SearchResultData {