test/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2024-11-30 02:48:45 +08:00
Code Issues Actions 11 Packages Projects Releases Wiki Activity

enhance: support disable search optimization (#32141)

Browse Source 
Signed-off-by: chasingegg <chao.gao@zilliz.com>
This commit is contained in:
Gao 2024-04-16 10:51:20 +08:00 committed by GitHub
parent 1af084ea6b
commit 55d894bd5e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 49 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
internal/querynodev2/optimizers/query_hook.go
View File

@ -24,8 +24,8 @@ type QueryHook interface {
}
func OptimizeSearchParams(ctx context.Context, req *querypb.SearchRequest, queryHook QueryHook, numSegments int) (*querypb.SearchRequest, error) {
// no hook applied, just return
if queryHook == nil {
// no hook applied or disabled, just return
if queryHook == nil || !paramtable.Get().AutoIndexConfig.Enable.GetAsBool() {
return req, nil
}

50
internal/querynodev2/optimizers/query_hook_test.go
View File

@ -35,13 +35,17 @@ func (suite *QueryHookSuite) TestOptimizeSearchParam() {
paramtable.Get().Save(paramtable.Get().AutoIndexConfig.EnableOptimize.Key, "true")
suite.Run("normal_run", func() {
paramtable.Get().Save(paramtable.Get().AutoIndexConfig.Enable.Key, "true")
mockHook := NewMockQueryHook(suite.T())
mockHook.EXPECT().Run(mock.Anything).Run(func(params map[string]any) {
params[common.TopKKey] = int64(50)
params[common.SearchParamKey] = `{"param": 2}`
}).Return(nil)
suite.queryHook = mockHook
defer func() { suite.queryHook = nil }()
defer func() {
paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.Enable.Key)
suite.queryHook = nil
}()
plan := &planpb.PlanNode{
Node: &planpb.PlanNode_VectorAnns{
@ -66,7 +70,37 @@ func (suite *QueryHookSuite) TestOptimizeSearchParam() {
suite.verifyQueryInfo(req, 50, `{"param": 2}`)
})
suite.Run("disable optimization", func() {
mockHook := NewMockQueryHook(suite.T())
suite.queryHook = mockHook
defer func() { suite.queryHook = nil }()
plan := &planpb.PlanNode{
Node: &planpb.PlanNode_VectorAnns{
VectorAnns: &planpb.VectorANNS{
QueryInfo: &planpb.QueryInfo{
Topk: 100,
SearchParams: `{"param": 1}`,
},
},
},
}
bs, err := proto.Marshal(plan)
suite.Require().NoError(err)
req, err := OptimizeSearchParams(ctx, &querypb.SearchRequest{
Req: &internalpb.SearchRequest{
SerializedExprPlan: bs,
},
TotalChannelNum: 2,
}, suite.queryHook, 2)
suite.NoError(err)
suite.verifyQueryInfo(req, 100, `{"param": 1}`)
})
suite.Run("no_hook", func() {
paramtable.Get().Save(paramtable.Get().AutoIndexConfig.Enable.Key, "true")
defer paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.Enable.Key)
suite.queryHook = nil
plan := &planpb.PlanNode{
Node: &planpb.PlanNode_VectorAnns{
@ -92,13 +126,17 @@ func (suite *QueryHookSuite) TestOptimizeSearchParam() {
})
suite.Run("other_plannode", func() {
paramtable.Get().Save(paramtable.Get().AutoIndexConfig.Enable.Key, "true")
mockHook := NewMockQueryHook(suite.T())
mockHook.EXPECT().Run(mock.Anything).Run(func(params map[string]any) {
params[common.TopKKey] = int64(50)
params[common.SearchParamKey] = `{"param": 2}`
}).Return(nil).Maybe()
suite.queryHook = mockHook
defer func() { suite.queryHook = nil }()
defer func() {
paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.Enable.Key)
suite.queryHook = nil
}()
plan := &planpb.PlanNode{
Node: &planpb.PlanNode_Query{},
@ -117,6 +155,8 @@ func (suite *QueryHookSuite) TestOptimizeSearchParam() {
})
suite.Run("no_serialized_plan", func() {
paramtable.Get().Save(paramtable.Get().AutoIndexConfig.Enable.Key, "true")
defer paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.Enable.Key)
mockHook := NewMockQueryHook(suite.T())
suite.queryHook = mockHook
defer func() { suite.queryHook = nil }()
@ -129,13 +169,17 @@ func (suite *QueryHookSuite) TestOptimizeSearchParam() {
})
suite.Run("hook_run_error", func() {
paramtable.Get().Save(paramtable.Get().AutoIndexConfig.Enable.Key, "true")
mockHook := NewMockQueryHook(suite.T())
mockHook.EXPECT().Run(mock.Anything).Run(func(params map[string]any) {
params[common.TopKKey] = int64(50)
params[common.SearchParamKey] = `{"param": 2}`
}).Return(merr.WrapErrServiceInternal("mocked"))
suite.queryHook = mockHook
defer func() { suite.queryHook = nil }()
defer func() {
paramtable.Get().Reset(paramtable.Get().AutoIndexConfig.Enable.Key)
suite.queryHook = nil
}()
plan := &planpb.PlanNode{
Node: &planpb.PlanNode_VectorAnns{