milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

enhance: support disable search optimization (#32143) 

master pr: https://github.com/milvus-io/milvus/pull/32141
2.4 pr: https://github.com/milvus-io/milvus/pull/32142

Signed-off-by: chasingegg <chao.gao@zilliz.com>
pull/32393/head
Gao 2024-04-17 10:45:20 +08:00 committed by GitHub
parent 0d000ae3b7
commit 89d95901c5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 52 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -12,6 +12,7 @@ import (
"github.com/milvus-io/milvus/pkg/common"
"github.com/milvus-io/milvus/pkg/log"
"github.com/milvus-io/milvus/pkg/util/merr"
"github.com/milvus-io/milvus/pkg/util/paramtable"
)
// QueryHook is the interface for search/query parameter optimizer.
@ -23,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
}

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

@ -13,6 +13,7 @@ import (
"github.com/milvus-io/milvus/internal/proto/querypb"
"github.com/milvus-io/milvus/pkg/common"
"github.com/milvus-io/milvus/pkg/util/merr"
"github.com/milvus-io/milvus/pkg/util/paramtable"
)
type QueryHookSuite struct {
@ -30,15 +31,20 @@ func (suite *QueryHookSuite) TearDownTest() {
func (suite *QueryHookSuite) TestOptimizeSearchParam() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
paramtable.Init()
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{
@ -63,7 +69,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{
@ -89,13 +125,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{},
@ -114,6 +154,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 }()
@ -126,13 +168,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{