Fix diskann param parse bug (#22498)

Signed-off-by: chasingegg <chao.gao@zilliz.com>
pull/22545/head
Gao 2023-03-02 19:59:50 +08:00 committed by GitHub
parent 897c7f33f7
commit aa19243c22
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 46 additions and 6 deletions

View File

@ -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)

View File

@ -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(&params, 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(&params, 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"