set default search list size by k in diskann (#909)

Signed-off-by: cqy123456 <[email protected]>
milvus-io · May 29, 2023 · 0b47731 · 0b47731
1 parent 98e2a16
commit 0b47731
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 4 deletions.
diff --git a/knowhere/index/vector_index/IndexDiskANNConfig.cpp b/knowhere/index/vector_index/IndexDiskANNConfig.cpp
@@ -64,6 +64,7 @@ static constexpr uint32_t kBuildNumThreadsMaxValue = 128;
 static constexpr uint32_t kDiskPqBytesMinValue = 0;
 static constexpr std::optional<uint32_t> kDiskPqBytesMaxValue = std::nullopt;
 static constexpr uint32_t kSearchListSizeMaxValue = 200;
+static constexpr uint32_t kKThreshold = 16;
 static constexpr uint32_t kBeamwidthMinValue = 1;
 static constexpr uint32_t kBeamwidthMaxValue = 128;
 static constexpr float kFilterThresholdMinValue = -1;
@@ -218,10 +219,14 @@ to_json(Config& config, const DiskANNQueryConfig& query_conf) {
 void
 from_json(const Config& config, DiskANNQueryConfig& query_conf) {
     CheckNumericParamAndSet<uint64_t>(config, kK, kKMinValue, kKMaxValue, query_conf.k);
-    // The search_list_size should be no less than the k.
-    CheckNumericParamAndSet<uint32_t>(config, kSearchListSize, query_conf.k,
-                                      std::max(kSearchListSizeMaxValue, static_cast<uint32_t>(10 * query_conf.k)),
-                                      query_conf.search_list_size);
+    if (config.contains(kSearchListSize)) {
+        // The search_list_size should be no less than the k.
+        CheckNumericParamAndSet<uint32_t>(config, kSearchListSize, query_conf.k,
+                                        std::max(kSearchListSizeMaxValue, static_cast<uint32_t>(10 * query_conf.k)),
+                                        query_conf.search_list_size);
+    } else {
+        query_conf.search_list_size = query_conf.k < kKThreshold ? kKThreshold : query_conf.k;
+    }
     CheckNumericParamAndSet<uint32_t>(config, kBeamwidth, kBeamwidthMinValue, kBeamwidthMaxValue, query_conf.beamwidth);
     if (config.contains(kFilterThreshold)) {
         CheckNumericParamAndSet<float>(config, kFilterThreshold, kFilterThresholdMinValue, kFilterThresholdMaxValue,

diff --git a/unittest/test_diskann.cpp b/unittest/test_diskann.cpp
@@ -455,6 +455,24 @@ TEST_P(DiskANNTest, meta_test) {
     EXPECT_EQ(num_rows_, diskann->Count());
 }
 
+TEST_P(DiskANNTest, search_without_search_list_size) {
+    knowhere::Config cfg;
+    cfg.clear();
+    knowhere::DiskANNQueryConfig::Set(cfg, query_conf);
+    auto search_list_size = knowhere::DiskANNQueryConfig::Get(cfg).search_list_size;
+    EXPECT_EQ(search_list_size, kK * 10);
+
+    auto& query_cfg = cfg["diskANN_query_config"];
+    query_cfg.erase("search_list_size");
+    // check generate default search_list size by k
+    query_cfg["k"] = 1;
+    search_list_size = knowhere::DiskANNQueryConfig::Get(cfg).search_list_size;
+    EXPECT_EQ(search_list_size, 16);
+    query_cfg["k"] = 32;
+    search_list_size = knowhere::DiskANNQueryConfig::Get(cfg).search_list_size;
+    EXPECT_EQ(search_list_size, 32);
+}
+
 TEST_P(DiskANNTest, knn_search_test) {
     knowhere::Config cfg;
     knowhere::DiskANNQueryConfig::Set(cfg, query_conf);