From 8a630f733a4710c6547287d623d8429b2e25bc7e Mon Sep 17 00:00:00 2001 From: Gao Date: Thu, 28 Dec 2023 10:04:46 +0800 Subject: [PATCH] enhance: add new optimize param for queryhook (#29495) add a flag to indicate if we use search param optimizations, default is on --------- Signed-off-by: chasingegg --- internal/querynodev2/optimizers/query_hook.go | 12 +++++++----- internal/querynodev2/optimizers/query_hook_test.go | 3 +++ pkg/common/common.go | 11 ++++++----- pkg/util/paramtable/autoindex_param.go | 11 ++++++++++- 4 files changed, 26 insertions(+), 11 deletions(-) diff --git a/internal/querynodev2/optimizers/query_hook.go b/internal/querynodev2/optimizers/query_hook.go index faaf990d1d..f9a6564a16 100644 --- a/internal/querynodev2/optimizers/query_hook.go +++ b/internal/querynodev2/optimizers/query_hook.go @@ -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. @@ -57,11 +58,12 @@ func OptimizeSearchParams(ctx context.Context, req *querypb.SearchRequest, query withFilter := (plan.GetVectorAnns().GetPredicates() != nil) queryInfo := plan.GetVectorAnns().GetQueryInfo() params := map[string]any{ - common.TopKKey: queryInfo.GetTopk(), - common.SearchParamKey: queryInfo.GetSearchParams(), - common.SegmentNumKey: estSegmentNum, - common.WithFilterKey: withFilter, - common.CollectionKey: req.GetReq().GetCollectionID(), + common.TopKKey: queryInfo.GetTopk(), + common.SearchParamKey: queryInfo.GetSearchParams(), + common.SegmentNumKey: estSegmentNum, + common.WithFilterKey: withFilter, + common.WithOptimizeKey: paramtable.Get().AutoIndexConfig.EnableOptimize.GetAsBool(), + common.CollectionKey: req.GetReq().GetCollectionID(), } err := queryHook.Run(params) if err != nil { diff --git a/internal/querynodev2/optimizers/query_hook_test.go b/internal/querynodev2/optimizers/query_hook_test.go index 132619b5e3..8e88618bd0 100644 --- a/internal/querynodev2/optimizers/query_hook_test.go +++ b/internal/querynodev2/optimizers/query_hook_test.go @@ -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,6 +31,8 @@ func (suite *QueryHookSuite) TearDownTest() { func (suite *QueryHookSuite) TestOptimizeSearchParam() { ctx, cancel := context.WithCancel(context.Background()) defer cancel() + paramtable.Init() + paramtable.Get().Save(paramtable.Get().AutoIndexConfig.EnableOptimize.Key, "true") suite.Run("normal_run", func() { mockHook := NewMockQueryHook(suite.T()) diff --git a/pkg/common/common.go b/pkg/common/common.go index 946611e5dc..3cc05025a4 100644 --- a/pkg/common/common.go +++ b/pkg/common/common.go @@ -88,11 +88,12 @@ const ( // Search, Index parameter keys const ( - TopKKey = "topk" - SearchParamKey = "search_param" - SegmentNumKey = "segment_num" - WithFilterKey = "with_filter" - CollectionKey = "collection" + TopKKey = "topk" + SearchParamKey = "search_param" + SegmentNumKey = "segment_num" + WithFilterKey = "with_filter" + WithOptimizeKey = "with_optimize" + CollectionKey = "collection" IndexParamsKey = "params" IndexTypeKey = "index_type" diff --git a/pkg/util/paramtable/autoindex_param.go b/pkg/util/paramtable/autoindex_param.go index 70d355381a..7273325dae 100644 --- a/pkg/util/paramtable/autoindex_param.go +++ b/pkg/util/paramtable/autoindex_param.go @@ -28,7 +28,8 @@ import ( // ///////////////////////////////////////////////////////////////////////////// // --- common --- type autoIndexConfig struct { - Enable ParamItem `refreshable:"true"` + Enable ParamItem `refreshable:"true"` + EnableOptimize ParamItem `refreshable:"true"` IndexParams ParamItem `refreshable:"true"` PrepareParams ParamItem `refreshable:"true"` @@ -48,6 +49,14 @@ func (p *autoIndexConfig) init(base *BaseTable) { } p.Enable.Init(base.mgr) + p.EnableOptimize = ParamItem{ + Key: "autoIndex.optimize", + Version: "2.4.0", + DefaultValue: "true", + PanicIfEmpty: true, + } + p.EnableOptimize.Init(base.mgr) + p.IndexParams = ParamItem{ Key: "autoIndex.params.build", Version: "2.2.0",