123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608 |
- // Licensed to the LF AI & Data foundation under one
- // or more contributor license agreements. See the NOTICE file
- // distributed with this work for additional information
- // regarding copyright ownership. The ASF licenses this file
- // to you under the Apache License, Version 2.0 (the
- // "License"); you may not use this file except in compliance
- // with the License. You may obtain a copy of the License at
- //
- // http://www.apache.org/licenses/LICENSE-2.0
- //
- // Unless required by applicable law or agreed to in writing, software
- // distributed under the License is distributed on an "AS IS" BASIS,
- // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- // See the License for the specific language governing permissions and
- // limitations under the License.
- package indexparams
- import (
- "encoding/json"
- "fmt"
- "strconv"
- "testing"
- "github.com/stretchr/testify/assert"
- "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
- "github.com/milvus-io/milvus/pkg/common"
- "github.com/milvus-io/milvus/pkg/util/hardware"
- "github.com/milvus-io/milvus/pkg/util/paramtable"
- )
- func TestDiskIndexParams(t *testing.T) {
- t.Run("fill index params without auto index param", func(t *testing.T) {
- var params paramtable.ComponentParam
- params.Init(paramtable.NewBaseTable(paramtable.SkipRemote(true)))
- indexParams := make(map[string]string)
- err := FillDiskIndexParams(¶ms, indexParams)
- assert.NoError(t, err)
- pqCodeBudgetGBRatio, err := strconv.ParseFloat(indexParams[PQCodeBudgetRatioKey], 64)
- assert.NoError(t, err)
- assert.Equal(t, 0.125, pqCodeBudgetGBRatio)
- buildNumThreadsRatio, err := strconv.ParseFloat(indexParams[NumBuildThreadRatioKey], 64)
- assert.NoError(t, err)
- assert.Equal(t, 1.0, buildNumThreadsRatio)
- searchCacheBudgetRatio, err := strconv.ParseFloat(indexParams[SearchCacheBudgetRatioKey], 64)
- assert.NoError(t, err)
- assert.Equal(t, 0.10, searchCacheBudgetRatio)
- })
- t.Run("fill index params with auto index", 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[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
- mapString[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_size"] = "100"
- indexParams[common.IndexTypeKey] = "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.NoError(t, err)
- pqCodeBudgetGBRatio, err := strconv.ParseFloat(indexParams[PQCodeBudgetRatioKey], 64)
- assert.NoError(t, err)
- assert.Equal(t, 0.125, pqCodeBudgetGBRatio)
- buildNumThreadsRatio, err := strconv.ParseFloat(indexParams[NumBuildThreadRatioKey], 64)
- assert.NoError(t, err)
- assert.Equal(t, 1.0, buildNumThreadsRatio)
- })
- t.Run("fill index params with wrong auto index param", func(t *testing.T) {
- var params paramtable.ComponentParam
- params.Init(paramtable.NewBaseTable(paramtable.SkipRemote(true)))
- 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[common.IndexTypeKey] = "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[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
- mapString[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[common.IndexTypeKey] = "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("patch index build params", func(t *testing.T) {
- var params paramtable.ComponentParam
- params.Init(paramtable.NewBaseTable(paramtable.SkipRemote(true)))
- indexParams := make([]*commonpb.KeyValuePair, 0, 3)
- indexParams = append(indexParams,
- &commonpb.KeyValuePair{
- Key: PQCodeBudgetRatioKey,
- Value: "0.125",
- })
- indexParams = append(indexParams,
- &commonpb.KeyValuePair{
- Key: NumBuildThreadRatioKey,
- Value: "1.0",
- })
- indexParams = append(indexParams,
- &commonpb.KeyValuePair{
- Key: BeamWidthRatioKey,
- Value: "4.0",
- })
- indexParams, err := UpdateDiskIndexBuildParams(¶ms, indexParams)
- assert.NoError(t, err)
- assert.True(t, len(indexParams) == 4)
- val := GetIndexParams(indexParams, SearchCacheBudgetRatioKey)
- cfgVal, cfgErr := strconv.ParseFloat(params.CommonCfg.SearchCacheBudgetGBRatio.GetValue(), 64)
- assert.NoError(t, cfgErr)
- iVal, iErr := strconv.ParseFloat(val, 64)
- assert.NoError(t, iErr)
- assert.Equal(t, cfgVal, iVal)
- params.Save(params.AutoIndexConfig.Enable.Key, "true")
- jsonStr := `
- {
- "build_ratio": "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}",
- "prepare_ratio": "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 8}",
- "beamwidth_ratio": "8.0"
- }
- `
- params.Save(params.AutoIndexConfig.ExtraParams.Key, jsonStr)
- autoParams := make([]*commonpb.KeyValuePair, 0, 3)
- autoParams = append(autoParams,
- &commonpb.KeyValuePair{
- Key: PQCodeBudgetRatioKey,
- Value: "0.125",
- })
- autoParams = append(autoParams,
- &commonpb.KeyValuePair{
- Key: NumBuildThreadRatioKey,
- Value: "1.0",
- })
- autoParams = append(autoParams,
- &commonpb.KeyValuePair{
- Key: BeamWidthRatioKey,
- Value: "4.0",
- })
- autoParams, err = UpdateDiskIndexBuildParams(¶ms, autoParams)
- assert.NoError(t, err)
- assert.True(t, len(autoParams) == 4)
- val = GetIndexParams(autoParams, SearchCacheBudgetRatioKey)
- iVal, iErr = strconv.ParseFloat(val, 64)
- assert.NoError(t, iErr)
- assert.Equal(t, 0.225, iVal)
- newJSONStr := `
- {
- "build_ratio": "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}",
- "prepare_ratio": "{\"search_cache_budget_gb\": 0.325, \"num_threads\": 8}",
- "beamwidth_ratio": "8.0"
- }
- `
- params.Save(params.AutoIndexConfig.ExtraParams.Key, newJSONStr)
- autoParams, err = UpdateDiskIndexBuildParams(¶ms, autoParams)
- assert.NoError(t, err)
- assert.True(t, len(autoParams) == 4)
- val = GetIndexParams(autoParams, SearchCacheBudgetRatioKey)
- iVal, iErr = strconv.ParseFloat(val, 64)
- assert.NoError(t, iErr)
- assert.Equal(t, 0.325, iVal)
- })
- t.Run("set disk index build params", func(t *testing.T) {
- indexParams := make(map[string]string)
- indexParams[PQCodeBudgetRatioKey] = "0.125"
- indexParams[NumBuildThreadRatioKey] = "1.0"
- indexParams[SearchCacheBudgetRatioKey] = "0.125"
- err := SetDiskIndexBuildParams(indexParams, 100)
- assert.NoError(t, err)
- _, ok := indexParams[SearchCacheBudgetKey]
- assert.True(t, ok)
- indexParams[SearchCacheBudgetRatioKey] = "aabb"
- err = SetDiskIndexBuildParams(indexParams, 100)
- assert.Error(t, err)
- delete(indexParams, SearchCacheBudgetRatioKey)
- delete(indexParams, SearchCacheBudgetKey)
- err = SetDiskIndexBuildParams(indexParams, 100)
- assert.NoError(t, err)
- _, ok = indexParams[PQCodeBudgetKey]
- assert.True(t, ok)
- _, ok = indexParams[BuildDramBudgetKey]
- assert.True(t, ok)
- _, ok = indexParams[NumBuildThreadKey]
- assert.True(t, ok)
- _, ok = indexParams[SearchCacheBudgetKey]
- assert.False(t, ok)
- })
- t.Run("set disk index load params without auto index param", func(t *testing.T) {
- var params paramtable.ComponentParam
- params.Init(paramtable.NewBaseTable(paramtable.SkipRemote(true)))
- indexParams := make(map[string]string)
- err := SetDiskIndexLoadParams(¶ms, indexParams, 100)
- assert.Error(t, err)
- indexParams[common.DimKey] = "128"
- err = SetDiskIndexLoadParams(¶ms, indexParams, 100)
- assert.NoError(t, err)
- searchCacheBudget, ok := indexParams[SearchCacheBudgetKey]
- assert.True(t, ok)
- searchCacheBudgetRatio, err := strconv.ParseFloat(params.CommonCfg.SearchCacheBudgetGBRatio.GetValue(), 64)
- assert.NoError(t, err)
- assert.Equal(t, fmt.Sprintf("%f", float32(getRowDataSizeOfFloatVector(100, 128))*float32(searchCacheBudgetRatio)/(1<<30)), searchCacheBudget)
- numLoadThread, ok := indexParams[NumLoadThreadKey]
- assert.True(t, ok)
- numLoadThreadRatio, err := strconv.ParseFloat(params.CommonCfg.LoadNumThreadRatio.GetValue(), 64)
- assert.NoError(t, err)
- expectedNumLoadThread := int(float32(hardware.GetCPUNum()) * float32(numLoadThreadRatio))
- if expectedNumLoadThread > MaxLoadThread {
- expectedNumLoadThread = MaxLoadThread
- }
- assert.Equal(t, strconv.Itoa(expectedNumLoadThread), numLoadThread)
- beamWidth, ok := indexParams[BeamWidthKey]
- assert.True(t, ok)
- beamWidthRatio, err := strconv.ParseFloat(params.CommonCfg.BeamWidthRatio.GetValue(), 64)
- assert.NoError(t, err)
- expectedBeamWidth := int(float32(hardware.GetCPUNum()) * float32(beamWidthRatio))
- if expectedBeamWidth > MaxBeamWidth {
- expectedBeamWidth = MaxBeamWidth
- }
- assert.Equal(t, strconv.Itoa(expectedBeamWidth), beamWidth)
- params.Save(params.CommonCfg.SearchCacheBudgetGBRatio.Key, "w1")
- err = SetDiskIndexLoadParams(¶ms, indexParams, 100)
- assert.Error(t, err)
- params.Save(params.CommonCfg.SearchCacheBudgetGBRatio.Key, "0.1")
- params.Save(params.CommonCfg.LoadNumThreadRatio.Key, "w1")
- err = SetDiskIndexLoadParams(¶ms, indexParams, 100)
- assert.Error(t, err)
- params.Save(params.CommonCfg.LoadNumThreadRatio.Key, "8.0")
- params.Save(params.CommonCfg.BeamWidthRatio.Key, "w1")
- err = SetDiskIndexLoadParams(¶ms, indexParams, 100)
- assert.Error(t, err)
- })
- t.Run("set disk index load params with auto index param", 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[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
- mapString[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))
- extraParams, err := NewBigDataExtraParamsFromJSON(params.AutoIndexConfig.ExtraParams.GetValue())
- assert.NoError(t, err)
- indexParams := make(map[string]string)
- err = SetDiskIndexLoadParams(¶ms, indexParams, 100)
- assert.Error(t, err)
- indexParams[common.DimKey] = "128"
- err = SetDiskIndexLoadParams(¶ms, indexParams, 100)
- assert.NoError(t, err)
- searchCacheBudget, ok := indexParams[SearchCacheBudgetKey]
- assert.True(t, ok)
- assert.Equal(t, fmt.Sprintf("%f", float32(getRowDataSizeOfFloatVector(100, 128))*float32(extraParams.SearchCacheBudgetGBRatio)/(1<<30)), searchCacheBudget)
- numLoadThread, ok := indexParams[NumLoadThreadKey]
- assert.True(t, ok)
- expectedNumLoadThread := int(float32(hardware.GetCPUNum()) * float32(extraParams.LoadNumThreadRatio))
- if expectedNumLoadThread > MaxLoadThread {
- expectedNumLoadThread = MaxLoadThread
- }
- assert.Equal(t, strconv.Itoa(expectedNumLoadThread), numLoadThread)
- beamWidth, ok := indexParams[BeamWidthKey]
- assert.True(t, ok)
- expectedBeamWidth := int(float32(hardware.GetCPUNum()) * float32(extraParams.BeamWidthRatio))
- if expectedBeamWidth > MaxBeamWidth {
- expectedBeamWidth = MaxBeamWidth
- }
- assert.Equal(t, strconv.Itoa(expectedBeamWidth), beamWidth)
- })
- t.Run("set disk index load params with wrong autoindex param", func(t *testing.T) {
- var params paramtable.ComponentParam
- params.Init(paramtable.NewBaseTable(paramtable.SkipRemote(true)))
- 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[common.IndexTypeKey] = "DISKANN"
- str, err := json.Marshal(indexParams)
- assert.NoError(t, err)
- params.Save(params.AutoIndexConfig.IndexParams.Key, string(str))
- indexParams[common.DimKey] = "128"
- err = SetDiskIndexLoadParams(¶ms, indexParams, 100)
- assert.Error(t, err)
- indexParams = make(map[string]string)
- err = SetDiskIndexLoadParams(¶ms, indexParams, 100)
- assert.Error(t, err)
- // IndexParams wrong
- mapString := make(map[string]string)
- mapString[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
- mapString[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[common.IndexTypeKey] = "DISKANN"
- str, err = json.Marshal(indexParams)
- assert.NoError(t, err)
- params.Save(params.AutoIndexConfig.IndexParams.Key, string(str))
- indexParams = make(map[string]string)
- err = SetDiskIndexLoadParams(¶ms, indexParams, 100)
- assert.Error(t, err)
- })
- }
- func TestBigDataIndex_parse(t *testing.T) {
- t.Run("parse normal", func(t *testing.T) {
- mapString := make(map[string]string)
- mapString[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
- mapString[PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 8}"
- extraParams, err := NewBigDataExtraParamsFromMap(mapString)
- assert.NoError(t, err)
- assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
- assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
- assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
- assert.Equal(t, 0.225, extraParams.SearchCacheBudgetGBRatio)
- mapString = make(map[string]string)
- mapString[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1, \"search_cache_budget_gb\": 0.20}"
- mapString[PrepareRatioKey] = "{\"num_threads\": 8}"
- extraParams, err = NewBigDataExtraParamsFromMap(mapString)
- assert.NoError(t, err)
- assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
- assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
- assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
- assert.Equal(t, 0.20, extraParams.SearchCacheBudgetGBRatio)
- })
- t.Run("parse with build_ratio partial or wrong", func(t *testing.T) {
- mapString := make(map[string]string)
- mapString[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.15}"
- mapString[PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 8}"
- extraParams, err := NewBigDataExtraParamsFromMap(mapString)
- assert.NoError(t, err)
- assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
- assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
- assert.Equal(t, 0.15, extraParams.PQCodeBudgetGBRatio)
- assert.Equal(t, 0.225, extraParams.SearchCacheBudgetGBRatio)
- mapString = make(map[string]string)
- mapString[BuildRatioKey] = "{\"num_threads\": 2}"
- mapString[PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 8}"
- extraParams, err = NewBigDataExtraParamsFromMap(mapString)
- assert.NoError(t, err)
- assert.Equal(t, 2.0, extraParams.BuildNumThreadsRatio)
- assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
- assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
- assert.Equal(t, 0.225, extraParams.SearchCacheBudgetGBRatio)
- mapString = make(map[string]string)
- mapString[BuildRatioKey] = "{\"num_threads\": 2"
- mapString[PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 8}"
- _, err = NewBigDataExtraParamsFromMap(mapString)
- assert.Error(t, err)
- mapString = make(map[string]string)
- mapString[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
- mapString[PrepareRatioKey] = "{\"num_threads\": 8}"
- extraParams, err = NewBigDataExtraParamsFromMap(mapString)
- assert.NoError(t, err)
- assert.Equal(t, 0.10, extraParams.SearchCacheBudgetGBRatio)
- })
- t.Run("parse with prepare_ratio partial or wrong", func(t *testing.T) {
- mapString := make(map[string]string)
- mapString[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
- mapString[PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.25}"
- extraParams, err := NewBigDataExtraParamsFromMap(mapString)
- assert.NoError(t, err)
- assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
- assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
- assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
- assert.Equal(t, 0.25, extraParams.SearchCacheBudgetGBRatio)
- mapString = make(map[string]string)
- mapString[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
- mapString[PrepareRatioKey] = "{\"num_threads\": 4}"
- extraParams, err = NewBigDataExtraParamsFromMap(mapString)
- assert.NoError(t, err)
- assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
- assert.Equal(t, 4.0, extraParams.LoadNumThreadRatio)
- assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
- assert.Equal(t, 0.10, extraParams.SearchCacheBudgetGBRatio)
- mapString = make(map[string]string)
- mapString[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
- mapString[PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.225"
- _, err = NewBigDataExtraParamsFromMap(mapString)
- assert.Error(t, err)
- })
- t.Run("parse with beamwidth wrong", func(t *testing.T) {
- mapString := make(map[string]string)
- mapString[BeamWidthRatioKey] = "aa"
- extraParams, err := NewBigDataExtraParamsFromMap(mapString)
- assert.NoError(t, err)
- assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
- assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
- assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
- assert.Equal(t, 0.10, extraParams.SearchCacheBudgetGBRatio)
- })
- t.Run("parse with partial", func(t *testing.T) {
- mapString := make(map[string]string)
- mapString[PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 8}"
- extraParams, err := NewBigDataExtraParamsFromMap(mapString)
- assert.NoError(t, err)
- assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
- assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
- assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
- assert.Equal(t, 0.225, extraParams.SearchCacheBudgetGBRatio)
- })
- t.Run("parse with empty", func(t *testing.T) {
- mapString := make(map[string]string)
- extraParams, err := NewBigDataExtraParamsFromMap(mapString)
- assert.NoError(t, err)
- assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
- assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
- assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
- assert.Equal(t, 0.10, extraParams.SearchCacheBudgetGBRatio)
- })
- t.Run("parse with nil", func(t *testing.T) {
- extraParams, err := NewBigDataExtraParamsFromMap(nil)
- assert.NoError(t, err)
- assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
- assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
- assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
- assert.Equal(t, 0.10, extraParams.SearchCacheBudgetGBRatio)
- })
- t.Run("new from json normal", func(t *testing.T) {
- jsonStr := `
- {
- "build_ratio": "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}",
- "prepare_ratio": "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 8}",
- "beamwidth_ratio": "8.0"
- }
- `
- extraParams, err := NewBigDataExtraParamsFromJSON(jsonStr)
- assert.NoError(t, err)
- assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
- assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
- assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
- assert.Equal(t, 0.225, extraParams.SearchCacheBudgetGBRatio)
- assert.Equal(t, 8.0, extraParams.BeamWidthRatio)
- })
- t.Run("new from json partial", func(t *testing.T) {
- jsonStr := `
- {
- "build_ratio": "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
- }
- `
- extraParams, err := NewBigDataExtraParamsFromJSON(jsonStr)
- assert.NoError(t, err)
- assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
- assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
- assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
- assert.Equal(t, 0.10, extraParams.SearchCacheBudgetGBRatio)
- assert.Equal(t, 4.0, extraParams.BeamWidthRatio)
- })
- t.Run("new from json empty", func(t *testing.T) {
- jsonStr := `
- {
- }
- `
- extraParams, err := NewBigDataExtraParamsFromJSON(jsonStr)
- assert.NoError(t, err)
- assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
- assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
- assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
- assert.Equal(t, 0.10, extraParams.SearchCacheBudgetGBRatio)
- assert.Equal(t, 4.0, extraParams.BeamWidthRatio)
- })
- t.Run("new from json invalid1", func(t *testing.T) {
- jsonStr := `
- { x
- }
- `
- _, err := NewBigDataExtraParamsFromJSON(jsonStr)
- assert.Error(t, err)
- })
- t.Run("new from json invalid1", func(t *testing.T) {
- jsonStr := `
- ""
- `
- _, err := NewBigDataExtraParamsFromJSON(jsonStr)
- assert.Error(t, err)
- })
- }
- func TestAppendPrepareInfo_parse(t *testing.T) {
- t.Run("parse load 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))
- mapString2 := make(map[string]string)
- mapString2["key2"] = "value2"
- str2, err2 := json.Marshal(mapString2)
- assert.NoError(t, err2)
- params.Save(params.AutoIndexConfig.LoadAdaptParams.Key, string(str2))
- resultMapString := make(map[string]string)
- err = AppendPrepareLoadParams(¶ms, resultMapString)
- assert.NoError(t, err)
- assert.Equal(t, resultMapString["key1"], "value1")
- assert.Equal(t, resultMapString["key2"], "value2")
- })
- }
|