From c8415453a1fd82d00e8de06eeed1dab436b36cfb Mon Sep 17 00:00:00 2001
From: Gao <chao.gao@zilliz.com>
Date: Wed, 25 Oct 2023 19:54:12 +0800
Subject: [PATCH] Support prepare config for index (#27920)

Signed-off-by: chasingegg <chao.gao@zilliz.com>
---
 .../querynodev2/segments/load_index_info.go    |  5 +++++
 .../{disk_index_params.go => index_params.go}  | 12 +++++++++++-
 ...dex_params_test.go => index_params_test.go} | 18 ++++++++++++++++++
 pkg/util/paramtable/autoindex_param.go         |  7 +++++++
 pkg/util/paramtable/autoindex_param_test.go    | 12 ++++++++++++
 5 files changed, 53 insertions(+), 1 deletion(-)
 rename pkg/util/indexparams/{disk_index_params.go => index_params.go} (96%)
 rename pkg/util/indexparams/{disk_index_params_test.go => index_params_test.go} (96%)

diff --git a/internal/querynodev2/segments/load_index_info.go b/internal/querynodev2/segments/load_index_info.go
index 125203d86b..3cb008d1c6 100644
--- a/internal/querynodev2/segments/load_index_info.go
+++ b/internal/querynodev2/segments/load_index_info.go
@@ -80,6 +80,11 @@ func (li *LoadIndexInfo) appendLoadIndexInfo(indexInfo *querypb.FieldIndexInfo,
 		}
 	}
 
+	err = indexparams.AppendPrepareLoadParams(paramtable.Get(), indexParams)
+	if err != nil {
+		return err
+	}
+
 	for key, value := range indexParams {
 		err = li.appendIndexParam(key, value)
 		if err != nil {
diff --git a/pkg/util/indexparams/disk_index_params.go b/pkg/util/indexparams/index_params.go
similarity index 96%
rename from pkg/util/indexparams/disk_index_params.go
rename to pkg/util/indexparams/index_params.go
index 564baef151..b26b2593b7 100644
--- a/pkg/util/indexparams/disk_index_params.go
+++ b/pkg/util/indexparams/index_params.go
@@ -208,7 +208,7 @@ func SetDiskIndexBuildParams(indexParams map[string]string, fieldDataSize int64)
 	}
 
 	searchCacheBudgetGBRatioStr, ok := indexParams[SearchCacheBudgetRatioKey]
-	// set generate cache size when cache ratio param set
+	// set generate cache size when cache ratio param not set
 	if ok {
 		SearchCacheBudgetGBRatio, err := strconv.ParseFloat(searchCacheBudgetGBRatioStr, 64)
 		if err != nil {
@@ -279,3 +279,13 @@ func SetDiskIndexLoadParams(params *paramtable.ComponentParam, indexParams map[s
 
 	return nil
 }
+
+func AppendPrepareLoadParams(params *paramtable.ComponentParam, indexParams map[string]string) error {
+	if params.AutoIndexConfig.Enable.GetAsBool() { // `enable` only for cloud instance.
+		// override prepare params by
+		for k, v := range params.AutoIndexConfig.PrepareParams.GetAsJSONMap() {
+			indexParams[k] = v
+		}
+	}
+	return nil
+}
diff --git a/pkg/util/indexparams/disk_index_params_test.go b/pkg/util/indexparams/index_params_test.go
similarity index 96%
rename from pkg/util/indexparams/disk_index_params_test.go
rename to pkg/util/indexparams/index_params_test.go
index 833301bab6..1bff36d5c8 100644
--- a/pkg/util/indexparams/disk_index_params_test.go
+++ b/pkg/util/indexparams/index_params_test.go
@@ -486,3 +486,21 @@ func TestBigDataIndex_parse(t *testing.T) {
 		assert.Error(t, err)
 	})
 }
+
+func TestAppendPrepareInfo_parse(t *testing.T) {
+	t.Run("parse prepare info", func(t *testing.T) {
+		var params paramtable.ComponentParam
+		params.Init(paramtable.NewBaseTable(paramtable.SkipRemote(true)))
+		params.Save(params.AutoIndexConfig.Enable.Key, "true")
+		mapString := make(map[string]string)
+		mapString["key1"] = "value1"
+		str, err := json.Marshal(mapString)
+		assert.NoError(t, err)
+		params.Save(params.AutoIndexConfig.PrepareParams.Key, string(str))
+
+		resultMapString := make(map[string]string)
+		err = AppendPrepareLoadParams(&params, resultMapString)
+		assert.NoError(t, err)
+		assert.Equal(t, resultMapString["key1"], "value1")
+	})
+}
diff --git a/pkg/util/paramtable/autoindex_param.go b/pkg/util/paramtable/autoindex_param.go
index cdc4f5289f..70d355381a 100644
--- a/pkg/util/paramtable/autoindex_param.go
+++ b/pkg/util/paramtable/autoindex_param.go
@@ -31,6 +31,7 @@ type autoIndexConfig struct {
 	Enable ParamItem `refreshable:"true"`
 
 	IndexParams           ParamItem  `refreshable:"true"`
+	PrepareParams         ParamItem  `refreshable:"true"`
 	ExtraParams           ParamItem  `refreshable:"true"`
 	IndexType             ParamItem  `refreshable:"true"`
 	AutoIndexTypeName     ParamItem  `refreshable:"true"`
@@ -54,6 +55,12 @@ func (p *autoIndexConfig) init(base *BaseTable) {
 	}
 	p.IndexParams.Init(base.mgr)
 
+	p.PrepareParams = ParamItem{
+		Key:     "autoIndex.params.prepare",
+		Version: "2.3.2",
+	}
+	p.PrepareParams.Init(base.mgr)
+
 	p.ExtraParams = ParamItem{
 		Key:     "autoIndex.params.extra",
 		Version: "2.2.0",
diff --git a/pkg/util/paramtable/autoindex_param_test.go b/pkg/util/paramtable/autoindex_param_test.go
index 1c4f262a47..83520ad3ef 100644
--- a/pkg/util/paramtable/autoindex_param_test.go
+++ b/pkg/util/paramtable/autoindex_param_test.go
@@ -65,6 +65,18 @@ func TestAutoIndexParams_build(t *testing.T) {
 		assert.Equal(t, "IVF_FLAT", CParams.AutoIndexConfig.IndexType.GetValue())
 		assert.Equal(t, strconv.Itoa(map2["nlist"].(int)), CParams.AutoIndexConfig.IndexParams.GetAsJSONMap()["nlist"])
 	})
+
+	t.Run("test parsePrepareParams success", func(t *testing.T) {
+		var err error
+		map1 := map[string]any{
+			"key1": 25,
+		}
+		var jsonStrBytes []byte
+		jsonStrBytes, err = json.Marshal(map1)
+		assert.NoError(t, err)
+		bt.Save(CParams.AutoIndexConfig.IndexParams.Key, string(jsonStrBytes))
+		assert.Equal(t, strconv.Itoa(map1["key1"].(int)), CParams.AutoIndexConfig.IndexParams.GetAsJSONMap()["key1"])
+	})
 }
 
 func Test_autoIndexConfig_panicIfNotValid(t *testing.T) {