index_params_test.go 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608
  1. // Licensed to the LF AI & Data foundation under one
  2. // or more contributor license agreements. See the NOTICE file
  3. // distributed with this work for additional information
  4. // regarding copyright ownership. The ASF licenses this file
  5. // to you under the Apache License, Version 2.0 (the
  6. // "License"); you may not use this file except in compliance
  7. // with the License. You may obtain a copy of the License at
  8. //
  9. // http://www.apache.org/licenses/LICENSE-2.0
  10. //
  11. // Unless required by applicable law or agreed to in writing, software
  12. // distributed under the License is distributed on an "AS IS" BASIS,
  13. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14. // See the License for the specific language governing permissions and
  15. // limitations under the License.
  16. package indexparams
  17. import (
  18. "encoding/json"
  19. "fmt"
  20. "strconv"
  21. "testing"
  22. "github.com/stretchr/testify/assert"
  23. "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
  24. "github.com/milvus-io/milvus/pkg/common"
  25. "github.com/milvus-io/milvus/pkg/util/hardware"
  26. "github.com/milvus-io/milvus/pkg/util/paramtable"
  27. )
  28. func TestDiskIndexParams(t *testing.T) {
  29. t.Run("fill index params without auto index param", func(t *testing.T) {
  30. var params paramtable.ComponentParam
  31. params.Init(paramtable.NewBaseTable(paramtable.SkipRemote(true)))
  32. indexParams := make(map[string]string)
  33. err := FillDiskIndexParams(&params, indexParams)
  34. assert.NoError(t, err)
  35. pqCodeBudgetGBRatio, err := strconv.ParseFloat(indexParams[PQCodeBudgetRatioKey], 64)
  36. assert.NoError(t, err)
  37. assert.Equal(t, 0.125, pqCodeBudgetGBRatio)
  38. buildNumThreadsRatio, err := strconv.ParseFloat(indexParams[NumBuildThreadRatioKey], 64)
  39. assert.NoError(t, err)
  40. assert.Equal(t, 1.0, buildNumThreadsRatio)
  41. searchCacheBudgetRatio, err := strconv.ParseFloat(indexParams[SearchCacheBudgetRatioKey], 64)
  42. assert.NoError(t, err)
  43. assert.Equal(t, 0.10, searchCacheBudgetRatio)
  44. })
  45. t.Run("fill index params with auto index", func(t *testing.T) {
  46. var params paramtable.ComponentParam
  47. params.Init(paramtable.NewBaseTable(paramtable.SkipRemote(true)))
  48. params.Save(params.AutoIndexConfig.Enable.Key, "true")
  49. mapString := make(map[string]string)
  50. mapString[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
  51. mapString[PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 4}"
  52. str, err := json.Marshal(mapString)
  53. assert.NoError(t, err)
  54. params.Save(params.AutoIndexConfig.ExtraParams.Key, string(str))
  55. indexParams := make(map[string]string)
  56. indexParams["max_degree"] = "56"
  57. indexParams["search_list_size"] = "100"
  58. indexParams[common.IndexTypeKey] = "DISKANN"
  59. str, err = json.Marshal(indexParams)
  60. assert.NoError(t, err)
  61. params.Save(params.AutoIndexConfig.IndexParams.Key, string(str))
  62. indexParams = make(map[string]string)
  63. err = FillDiskIndexParams(&params, indexParams)
  64. assert.NoError(t, err)
  65. pqCodeBudgetGBRatio, err := strconv.ParseFloat(indexParams[PQCodeBudgetRatioKey], 64)
  66. assert.NoError(t, err)
  67. assert.Equal(t, 0.125, pqCodeBudgetGBRatio)
  68. buildNumThreadsRatio, err := strconv.ParseFloat(indexParams[NumBuildThreadRatioKey], 64)
  69. assert.NoError(t, err)
  70. assert.Equal(t, 1.0, buildNumThreadsRatio)
  71. })
  72. t.Run("fill index params with wrong auto index param", func(t *testing.T) {
  73. var params paramtable.ComponentParam
  74. params.Init(paramtable.NewBaseTable(paramtable.SkipRemote(true)))
  75. params.Save(params.AutoIndexConfig.Enable.Key, "true")
  76. // ExtraParams wrong
  77. params.Save(params.AutoIndexConfig.ExtraParams.Key, "")
  78. indexParams := make(map[string]string)
  79. indexParams["max_degree"] = "56"
  80. indexParams["search_list_size"] = "100"
  81. indexParams[common.IndexTypeKey] = "DISKANN"
  82. str, err := json.Marshal(indexParams)
  83. assert.NoError(t, err)
  84. params.Save(params.AutoIndexConfig.IndexParams.Key, string(str))
  85. indexParams = make(map[string]string)
  86. err = FillDiskIndexParams(&params, indexParams)
  87. assert.Error(t, err)
  88. // IndexParams wrong
  89. mapString := make(map[string]string)
  90. mapString[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
  91. mapString[PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 4}"
  92. str, err = json.Marshal(mapString)
  93. assert.NoError(t, err)
  94. params.Save(params.AutoIndexConfig.ExtraParams.Key, string(str))
  95. indexParams = make(map[string]string)
  96. indexParams["max_degree"] = "56"
  97. indexParams["search_list"] = "100" // should be search_list_size
  98. indexParams[common.IndexTypeKey] = "DISKANN"
  99. str, err = json.Marshal(indexParams)
  100. assert.NoError(t, err)
  101. params.Save(params.AutoIndexConfig.IndexParams.Key, string(str))
  102. indexParams = make(map[string]string)
  103. err = FillDiskIndexParams(&params, indexParams)
  104. assert.Error(t, err)
  105. })
  106. t.Run("patch index build params", func(t *testing.T) {
  107. var params paramtable.ComponentParam
  108. params.Init(paramtable.NewBaseTable(paramtable.SkipRemote(true)))
  109. indexParams := make([]*commonpb.KeyValuePair, 0, 3)
  110. indexParams = append(indexParams,
  111. &commonpb.KeyValuePair{
  112. Key: PQCodeBudgetRatioKey,
  113. Value: "0.125",
  114. })
  115. indexParams = append(indexParams,
  116. &commonpb.KeyValuePair{
  117. Key: NumBuildThreadRatioKey,
  118. Value: "1.0",
  119. })
  120. indexParams = append(indexParams,
  121. &commonpb.KeyValuePair{
  122. Key: BeamWidthRatioKey,
  123. Value: "4.0",
  124. })
  125. indexParams, err := UpdateDiskIndexBuildParams(&params, indexParams)
  126. assert.NoError(t, err)
  127. assert.True(t, len(indexParams) == 4)
  128. val := GetIndexParams(indexParams, SearchCacheBudgetRatioKey)
  129. cfgVal, cfgErr := strconv.ParseFloat(params.CommonCfg.SearchCacheBudgetGBRatio.GetValue(), 64)
  130. assert.NoError(t, cfgErr)
  131. iVal, iErr := strconv.ParseFloat(val, 64)
  132. assert.NoError(t, iErr)
  133. assert.Equal(t, cfgVal, iVal)
  134. params.Save(params.AutoIndexConfig.Enable.Key, "true")
  135. jsonStr := `
  136. {
  137. "build_ratio": "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}",
  138. "prepare_ratio": "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 8}",
  139. "beamwidth_ratio": "8.0"
  140. }
  141. `
  142. params.Save(params.AutoIndexConfig.ExtraParams.Key, jsonStr)
  143. autoParams := make([]*commonpb.KeyValuePair, 0, 3)
  144. autoParams = append(autoParams,
  145. &commonpb.KeyValuePair{
  146. Key: PQCodeBudgetRatioKey,
  147. Value: "0.125",
  148. })
  149. autoParams = append(autoParams,
  150. &commonpb.KeyValuePair{
  151. Key: NumBuildThreadRatioKey,
  152. Value: "1.0",
  153. })
  154. autoParams = append(autoParams,
  155. &commonpb.KeyValuePair{
  156. Key: BeamWidthRatioKey,
  157. Value: "4.0",
  158. })
  159. autoParams, err = UpdateDiskIndexBuildParams(&params, autoParams)
  160. assert.NoError(t, err)
  161. assert.True(t, len(autoParams) == 4)
  162. val = GetIndexParams(autoParams, SearchCacheBudgetRatioKey)
  163. iVal, iErr = strconv.ParseFloat(val, 64)
  164. assert.NoError(t, iErr)
  165. assert.Equal(t, 0.225, iVal)
  166. newJSONStr := `
  167. {
  168. "build_ratio": "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}",
  169. "prepare_ratio": "{\"search_cache_budget_gb\": 0.325, \"num_threads\": 8}",
  170. "beamwidth_ratio": "8.0"
  171. }
  172. `
  173. params.Save(params.AutoIndexConfig.ExtraParams.Key, newJSONStr)
  174. autoParams, err = UpdateDiskIndexBuildParams(&params, autoParams)
  175. assert.NoError(t, err)
  176. assert.True(t, len(autoParams) == 4)
  177. val = GetIndexParams(autoParams, SearchCacheBudgetRatioKey)
  178. iVal, iErr = strconv.ParseFloat(val, 64)
  179. assert.NoError(t, iErr)
  180. assert.Equal(t, 0.325, iVal)
  181. })
  182. t.Run("set disk index build params", func(t *testing.T) {
  183. indexParams := make(map[string]string)
  184. indexParams[PQCodeBudgetRatioKey] = "0.125"
  185. indexParams[NumBuildThreadRatioKey] = "1.0"
  186. indexParams[SearchCacheBudgetRatioKey] = "0.125"
  187. err := SetDiskIndexBuildParams(indexParams, 100)
  188. assert.NoError(t, err)
  189. _, ok := indexParams[SearchCacheBudgetKey]
  190. assert.True(t, ok)
  191. indexParams[SearchCacheBudgetRatioKey] = "aabb"
  192. err = SetDiskIndexBuildParams(indexParams, 100)
  193. assert.Error(t, err)
  194. delete(indexParams, SearchCacheBudgetRatioKey)
  195. delete(indexParams, SearchCacheBudgetKey)
  196. err = SetDiskIndexBuildParams(indexParams, 100)
  197. assert.NoError(t, err)
  198. _, ok = indexParams[PQCodeBudgetKey]
  199. assert.True(t, ok)
  200. _, ok = indexParams[BuildDramBudgetKey]
  201. assert.True(t, ok)
  202. _, ok = indexParams[NumBuildThreadKey]
  203. assert.True(t, ok)
  204. _, ok = indexParams[SearchCacheBudgetKey]
  205. assert.False(t, ok)
  206. })
  207. t.Run("set disk index load params without auto index param", func(t *testing.T) {
  208. var params paramtable.ComponentParam
  209. params.Init(paramtable.NewBaseTable(paramtable.SkipRemote(true)))
  210. indexParams := make(map[string]string)
  211. err := SetDiskIndexLoadParams(&params, indexParams, 100)
  212. assert.Error(t, err)
  213. indexParams[common.DimKey] = "128"
  214. err = SetDiskIndexLoadParams(&params, indexParams, 100)
  215. assert.NoError(t, err)
  216. searchCacheBudget, ok := indexParams[SearchCacheBudgetKey]
  217. assert.True(t, ok)
  218. searchCacheBudgetRatio, err := strconv.ParseFloat(params.CommonCfg.SearchCacheBudgetGBRatio.GetValue(), 64)
  219. assert.NoError(t, err)
  220. assert.Equal(t, fmt.Sprintf("%f", float32(getRowDataSizeOfFloatVector(100, 128))*float32(searchCacheBudgetRatio)/(1<<30)), searchCacheBudget)
  221. numLoadThread, ok := indexParams[NumLoadThreadKey]
  222. assert.True(t, ok)
  223. numLoadThreadRatio, err := strconv.ParseFloat(params.CommonCfg.LoadNumThreadRatio.GetValue(), 64)
  224. assert.NoError(t, err)
  225. expectedNumLoadThread := int(float32(hardware.GetCPUNum()) * float32(numLoadThreadRatio))
  226. if expectedNumLoadThread > MaxLoadThread {
  227. expectedNumLoadThread = MaxLoadThread
  228. }
  229. assert.Equal(t, strconv.Itoa(expectedNumLoadThread), numLoadThread)
  230. beamWidth, ok := indexParams[BeamWidthKey]
  231. assert.True(t, ok)
  232. beamWidthRatio, err := strconv.ParseFloat(params.CommonCfg.BeamWidthRatio.GetValue(), 64)
  233. assert.NoError(t, err)
  234. expectedBeamWidth := int(float32(hardware.GetCPUNum()) * float32(beamWidthRatio))
  235. if expectedBeamWidth > MaxBeamWidth {
  236. expectedBeamWidth = MaxBeamWidth
  237. }
  238. assert.Equal(t, strconv.Itoa(expectedBeamWidth), beamWidth)
  239. params.Save(params.CommonCfg.SearchCacheBudgetGBRatio.Key, "w1")
  240. err = SetDiskIndexLoadParams(&params, indexParams, 100)
  241. assert.Error(t, err)
  242. params.Save(params.CommonCfg.SearchCacheBudgetGBRatio.Key, "0.1")
  243. params.Save(params.CommonCfg.LoadNumThreadRatio.Key, "w1")
  244. err = SetDiskIndexLoadParams(&params, indexParams, 100)
  245. assert.Error(t, err)
  246. params.Save(params.CommonCfg.LoadNumThreadRatio.Key, "8.0")
  247. params.Save(params.CommonCfg.BeamWidthRatio.Key, "w1")
  248. err = SetDiskIndexLoadParams(&params, indexParams, 100)
  249. assert.Error(t, err)
  250. })
  251. t.Run("set disk index load params with auto index param", func(t *testing.T) {
  252. var params paramtable.ComponentParam
  253. params.Init(paramtable.NewBaseTable(paramtable.SkipRemote(true)))
  254. params.Save(params.AutoIndexConfig.Enable.Key, "true")
  255. mapString := make(map[string]string)
  256. mapString[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
  257. mapString[PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 4}"
  258. str, err := json.Marshal(mapString)
  259. assert.NoError(t, err)
  260. params.Save(params.AutoIndexConfig.ExtraParams.Key, string(str))
  261. extraParams, err := NewBigDataExtraParamsFromJSON(params.AutoIndexConfig.ExtraParams.GetValue())
  262. assert.NoError(t, err)
  263. indexParams := make(map[string]string)
  264. err = SetDiskIndexLoadParams(&params, indexParams, 100)
  265. assert.Error(t, err)
  266. indexParams[common.DimKey] = "128"
  267. err = SetDiskIndexLoadParams(&params, indexParams, 100)
  268. assert.NoError(t, err)
  269. searchCacheBudget, ok := indexParams[SearchCacheBudgetKey]
  270. assert.True(t, ok)
  271. assert.Equal(t, fmt.Sprintf("%f", float32(getRowDataSizeOfFloatVector(100, 128))*float32(extraParams.SearchCacheBudgetGBRatio)/(1<<30)), searchCacheBudget)
  272. numLoadThread, ok := indexParams[NumLoadThreadKey]
  273. assert.True(t, ok)
  274. expectedNumLoadThread := int(float32(hardware.GetCPUNum()) * float32(extraParams.LoadNumThreadRatio))
  275. if expectedNumLoadThread > MaxLoadThread {
  276. expectedNumLoadThread = MaxLoadThread
  277. }
  278. assert.Equal(t, strconv.Itoa(expectedNumLoadThread), numLoadThread)
  279. beamWidth, ok := indexParams[BeamWidthKey]
  280. assert.True(t, ok)
  281. expectedBeamWidth := int(float32(hardware.GetCPUNum()) * float32(extraParams.BeamWidthRatio))
  282. if expectedBeamWidth > MaxBeamWidth {
  283. expectedBeamWidth = MaxBeamWidth
  284. }
  285. assert.Equal(t, strconv.Itoa(expectedBeamWidth), beamWidth)
  286. })
  287. t.Run("set disk index load params with wrong autoindex param", func(t *testing.T) {
  288. var params paramtable.ComponentParam
  289. params.Init(paramtable.NewBaseTable(paramtable.SkipRemote(true)))
  290. params.Save(params.AutoIndexConfig.Enable.Key, "true")
  291. // ExtraParams wrong
  292. params.Save(params.AutoIndexConfig.ExtraParams.Key, "")
  293. indexParams := make(map[string]string)
  294. indexParams["max_degree"] = "56"
  295. indexParams["search_list_size"] = "100"
  296. indexParams[common.IndexTypeKey] = "DISKANN"
  297. str, err := json.Marshal(indexParams)
  298. assert.NoError(t, err)
  299. params.Save(params.AutoIndexConfig.IndexParams.Key, string(str))
  300. indexParams[common.DimKey] = "128"
  301. err = SetDiskIndexLoadParams(&params, indexParams, 100)
  302. assert.Error(t, err)
  303. indexParams = make(map[string]string)
  304. err = SetDiskIndexLoadParams(&params, indexParams, 100)
  305. assert.Error(t, err)
  306. // IndexParams wrong
  307. mapString := make(map[string]string)
  308. mapString[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
  309. mapString[PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 4}"
  310. str, err = json.Marshal(mapString)
  311. assert.NoError(t, err)
  312. params.Save(params.AutoIndexConfig.ExtraParams.Key, string(str))
  313. indexParams = make(map[string]string)
  314. indexParams["max_degree"] = "56"
  315. indexParams["search_list"] = "100" // should be search_list_size
  316. indexParams[common.IndexTypeKey] = "DISKANN"
  317. str, err = json.Marshal(indexParams)
  318. assert.NoError(t, err)
  319. params.Save(params.AutoIndexConfig.IndexParams.Key, string(str))
  320. indexParams = make(map[string]string)
  321. err = SetDiskIndexLoadParams(&params, indexParams, 100)
  322. assert.Error(t, err)
  323. })
  324. }
  325. func TestBigDataIndex_parse(t *testing.T) {
  326. t.Run("parse normal", func(t *testing.T) {
  327. mapString := make(map[string]string)
  328. mapString[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
  329. mapString[PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 8}"
  330. extraParams, err := NewBigDataExtraParamsFromMap(mapString)
  331. assert.NoError(t, err)
  332. assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
  333. assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
  334. assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
  335. assert.Equal(t, 0.225, extraParams.SearchCacheBudgetGBRatio)
  336. mapString = make(map[string]string)
  337. mapString[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1, \"search_cache_budget_gb\": 0.20}"
  338. mapString[PrepareRatioKey] = "{\"num_threads\": 8}"
  339. extraParams, err = NewBigDataExtraParamsFromMap(mapString)
  340. assert.NoError(t, err)
  341. assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
  342. assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
  343. assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
  344. assert.Equal(t, 0.20, extraParams.SearchCacheBudgetGBRatio)
  345. })
  346. t.Run("parse with build_ratio partial or wrong", func(t *testing.T) {
  347. mapString := make(map[string]string)
  348. mapString[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.15}"
  349. mapString[PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 8}"
  350. extraParams, err := NewBigDataExtraParamsFromMap(mapString)
  351. assert.NoError(t, err)
  352. assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
  353. assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
  354. assert.Equal(t, 0.15, extraParams.PQCodeBudgetGBRatio)
  355. assert.Equal(t, 0.225, extraParams.SearchCacheBudgetGBRatio)
  356. mapString = make(map[string]string)
  357. mapString[BuildRatioKey] = "{\"num_threads\": 2}"
  358. mapString[PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 8}"
  359. extraParams, err = NewBigDataExtraParamsFromMap(mapString)
  360. assert.NoError(t, err)
  361. assert.Equal(t, 2.0, extraParams.BuildNumThreadsRatio)
  362. assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
  363. assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
  364. assert.Equal(t, 0.225, extraParams.SearchCacheBudgetGBRatio)
  365. mapString = make(map[string]string)
  366. mapString[BuildRatioKey] = "{\"num_threads\": 2"
  367. mapString[PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 8}"
  368. _, err = NewBigDataExtraParamsFromMap(mapString)
  369. assert.Error(t, err)
  370. mapString = make(map[string]string)
  371. mapString[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
  372. mapString[PrepareRatioKey] = "{\"num_threads\": 8}"
  373. extraParams, err = NewBigDataExtraParamsFromMap(mapString)
  374. assert.NoError(t, err)
  375. assert.Equal(t, 0.10, extraParams.SearchCacheBudgetGBRatio)
  376. })
  377. t.Run("parse with prepare_ratio partial or wrong", func(t *testing.T) {
  378. mapString := make(map[string]string)
  379. mapString[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
  380. mapString[PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.25}"
  381. extraParams, err := NewBigDataExtraParamsFromMap(mapString)
  382. assert.NoError(t, err)
  383. assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
  384. assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
  385. assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
  386. assert.Equal(t, 0.25, extraParams.SearchCacheBudgetGBRatio)
  387. mapString = make(map[string]string)
  388. mapString[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
  389. mapString[PrepareRatioKey] = "{\"num_threads\": 4}"
  390. extraParams, err = NewBigDataExtraParamsFromMap(mapString)
  391. assert.NoError(t, err)
  392. assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
  393. assert.Equal(t, 4.0, extraParams.LoadNumThreadRatio)
  394. assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
  395. assert.Equal(t, 0.10, extraParams.SearchCacheBudgetGBRatio)
  396. mapString = make(map[string]string)
  397. mapString[BuildRatioKey] = "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
  398. mapString[PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.225"
  399. _, err = NewBigDataExtraParamsFromMap(mapString)
  400. assert.Error(t, err)
  401. })
  402. t.Run("parse with beamwidth wrong", func(t *testing.T) {
  403. mapString := make(map[string]string)
  404. mapString[BeamWidthRatioKey] = "aa"
  405. extraParams, err := NewBigDataExtraParamsFromMap(mapString)
  406. assert.NoError(t, err)
  407. assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
  408. assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
  409. assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
  410. assert.Equal(t, 0.10, extraParams.SearchCacheBudgetGBRatio)
  411. })
  412. t.Run("parse with partial", func(t *testing.T) {
  413. mapString := make(map[string]string)
  414. mapString[PrepareRatioKey] = "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 8}"
  415. extraParams, err := NewBigDataExtraParamsFromMap(mapString)
  416. assert.NoError(t, err)
  417. assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
  418. assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
  419. assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
  420. assert.Equal(t, 0.225, extraParams.SearchCacheBudgetGBRatio)
  421. })
  422. t.Run("parse with empty", func(t *testing.T) {
  423. mapString := make(map[string]string)
  424. extraParams, err := NewBigDataExtraParamsFromMap(mapString)
  425. assert.NoError(t, err)
  426. assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
  427. assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
  428. assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
  429. assert.Equal(t, 0.10, extraParams.SearchCacheBudgetGBRatio)
  430. })
  431. t.Run("parse with nil", func(t *testing.T) {
  432. extraParams, err := NewBigDataExtraParamsFromMap(nil)
  433. assert.NoError(t, err)
  434. assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
  435. assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
  436. assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
  437. assert.Equal(t, 0.10, extraParams.SearchCacheBudgetGBRatio)
  438. })
  439. t.Run("new from json normal", func(t *testing.T) {
  440. jsonStr := `
  441. {
  442. "build_ratio": "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}",
  443. "prepare_ratio": "{\"search_cache_budget_gb\": 0.225, \"num_threads\": 8}",
  444. "beamwidth_ratio": "8.0"
  445. }
  446. `
  447. extraParams, err := NewBigDataExtraParamsFromJSON(jsonStr)
  448. assert.NoError(t, err)
  449. assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
  450. assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
  451. assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
  452. assert.Equal(t, 0.225, extraParams.SearchCacheBudgetGBRatio)
  453. assert.Equal(t, 8.0, extraParams.BeamWidthRatio)
  454. })
  455. t.Run("new from json partial", func(t *testing.T) {
  456. jsonStr := `
  457. {
  458. "build_ratio": "{\"pq_code_budget_gb\": 0.125, \"num_threads\": 1}"
  459. }
  460. `
  461. extraParams, err := NewBigDataExtraParamsFromJSON(jsonStr)
  462. assert.NoError(t, err)
  463. assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
  464. assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
  465. assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
  466. assert.Equal(t, 0.10, extraParams.SearchCacheBudgetGBRatio)
  467. assert.Equal(t, 4.0, extraParams.BeamWidthRatio)
  468. })
  469. t.Run("new from json empty", func(t *testing.T) {
  470. jsonStr := `
  471. {
  472. }
  473. `
  474. extraParams, err := NewBigDataExtraParamsFromJSON(jsonStr)
  475. assert.NoError(t, err)
  476. assert.Equal(t, 1.0, extraParams.BuildNumThreadsRatio)
  477. assert.Equal(t, 8.0, extraParams.LoadNumThreadRatio)
  478. assert.Equal(t, 0.125, extraParams.PQCodeBudgetGBRatio)
  479. assert.Equal(t, 0.10, extraParams.SearchCacheBudgetGBRatio)
  480. assert.Equal(t, 4.0, extraParams.BeamWidthRatio)
  481. })
  482. t.Run("new from json invalid1", func(t *testing.T) {
  483. jsonStr := `
  484. { x
  485. }
  486. `
  487. _, err := NewBigDataExtraParamsFromJSON(jsonStr)
  488. assert.Error(t, err)
  489. })
  490. t.Run("new from json invalid1", func(t *testing.T) {
  491. jsonStr := `
  492. ""
  493. `
  494. _, err := NewBigDataExtraParamsFromJSON(jsonStr)
  495. assert.Error(t, err)
  496. })
  497. }
  498. func TestAppendPrepareInfo_parse(t *testing.T) {
  499. t.Run("parse load info", func(t *testing.T) {
  500. var params paramtable.ComponentParam
  501. params.Init(paramtable.NewBaseTable(paramtable.SkipRemote(true)))
  502. params.Save(params.AutoIndexConfig.Enable.Key, "true")
  503. mapString := make(map[string]string)
  504. mapString["key1"] = "value1"
  505. str, err := json.Marshal(mapString)
  506. assert.NoError(t, err)
  507. params.Save(params.AutoIndexConfig.PrepareParams.Key, string(str))
  508. mapString2 := make(map[string]string)
  509. mapString2["key2"] = "value2"
  510. str2, err2 := json.Marshal(mapString2)
  511. assert.NoError(t, err2)
  512. params.Save(params.AutoIndexConfig.LoadAdaptParams.Key, string(str2))
  513. resultMapString := make(map[string]string)
  514. err = AppendPrepareLoadParams(&params, resultMapString)
  515. assert.NoError(t, err)
  516. assert.Equal(t, resultMapString["key1"], "value1")
  517. assert.Equal(t, resultMapString["key2"], "value2")
  518. })
  519. }