diff --git a/internal/util/indexparams/disk_index_params.go b/internal/util/indexparams/disk_index_params.go index 0f67d83c13..8b8298d2c7 100644 --- a/internal/util/indexparams/disk_index_params.go +++ b/internal/util/indexparams/disk_index_params.go @@ -17,7 +17,6 @@ package indexparams import ( - "encoding/json" "fmt" "strconv" "unsafe" @@ -76,8 +75,10 @@ func FillDiskIndexParams(params *paramtable.ComponentParam, indexParams map[stri if !ok { return fmt.Errorf("index param search_list_size not exist") } - extraParams := autoindex.BigDataIndexExtraParams{} - json.Unmarshal([]byte(params.AutoIndexConfig.ExtraParams.GetValue()), &extraParams) + extraParams, err := autoindex.NewBigDataExtraParamsFromJSON(params.AutoIndexConfig.ExtraParams.GetValue()) + if err != nil { + return err + } pqCodeBudgetGBRatio = fmt.Sprintf("%f", extraParams.PQCodeBudgetGBRatio) buildNumThreadsRatio = fmt.Sprintf("%f", extraParams.BuildNumThreadsRatio) searchCacheBudgetGBRatio = fmt.Sprintf("%f", extraParams.SearchCacheBudgetGBRatio) diff --git a/internal/util/indexparams/disk_index_params_test.go b/internal/util/indexparams/disk_index_params_test.go index 2353ebe7a0..2eaf60103a 100644 --- a/internal/util/indexparams/disk_index_params_test.go +++ b/internal/util/indexparams/disk_index_params_test.go @@ -65,9 +65,8 @@ func TestDiskIndexParams(t *testing.T) { mapString := make(map[string]string) mapString[autoindex.BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}" mapString[autoindex.PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 4}" - extraParams, err := autoindex.NewBigDataExtraParamsFromMap(mapString) - assert.NoError(t, err) - str, err := json.Marshal(extraParams) + + str, err := json.Marshal(mapString) assert.NoError(t, err) params.Save(params.AutoIndexConfig.ExtraParams.Key, string(str)) indexParams := make(map[string]string) @@ -103,6 +102,46 @@ func TestDiskIndexParams(t *testing.T) { assert.Equal(t, 4.0, beamWidthRatio) }) + t.Run("fill index params with wrong auto index param", func(t *testing.T) { + var params paramtable.ComponentParam + params.Init() + params.Save(params.AutoIndexConfig.Enable.Key, "true") + // ExtraParams wrong + params.Save(params.AutoIndexConfig.ExtraParams.Key, "") + indexParams := make(map[string]string) + indexParams["max_degree"] = "56" + indexParams["search_list_size"] = "100" + indexParams["index_type"] = "DISKANN" + str, err := json.Marshal(indexParams) + assert.NoError(t, err) + params.Save(params.AutoIndexConfig.IndexParams.Key, string(str)) + + indexParams = make(map[string]string) + err = FillDiskIndexParams(¶ms, indexParams) + assert.Error(t, err) + + // IndexParams wrong + mapString := make(map[string]string) + mapString[autoindex.BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}" + mapString[autoindex.PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 4}" + + str, err = json.Marshal(mapString) + assert.NoError(t, err) + params.Save(params.AutoIndexConfig.ExtraParams.Key, string(str)) + + indexParams = make(map[string]string) + indexParams["max_degree"] = "56" + indexParams["search_list"] = "100" // should be search_list_size + indexParams["index_type"] = "DISKANN" + str, err = json.Marshal(indexParams) + assert.NoError(t, err) + params.Save(params.AutoIndexConfig.IndexParams.Key, string(str)) + + indexParams = make(map[string]string) + err = FillDiskIndexParams(¶ms, indexParams) + assert.Error(t, err) + }) + t.Run("set disk index build params", func(t *testing.T) { indexParams := make(map[string]string) indexParams[PQCodeBudgetRatioKey] = "0.125"