[feat][benchmark] Support vector search nprobe and ef parameter.

DataCanvasIO · Mar 26, 2024 · 472d3d4 · 472d3d4
1 parent ff6b83c
commit 472d3d4
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 2 deletions.
diff --git a/src/benchmark/benchmark.cc b/src/benchmark/benchmark.cc
@@ -108,6 +108,8 @@ DEFINE_double(vector_search_radius, 0.1, "Vector search flag radius");
 DECLARE_uint32(vector_put_batch_size);
 DECLARE_uint32(vector_arrange_concurrency);
 DECLARE_bool(vector_search_arrange_data);
+DECLARE_uint32(vector_search_nprobe);
+DECLARE_uint32(vector_search_ef);
 
 DECLARE_string(benchmark);
 DECLARE_uint32(key_size);
@@ -898,7 +900,8 @@ void Environment::PrintParam() {
                            FLAGS_vector_search_enable_range_search ? "true" : "false")
             << '\n';
   std::cout << fmt::format("{:<34}: {:>32}", "vector_search_radius", FLAGS_vector_search_radius) << '\n';
-
+  std::cout << fmt::format("{:<34}: {:>32}", "vector_search_nprobe", FLAGS_vector_search_nprobe) << '\n';
+  std::cout << fmt::format("{:<34}: {:>32}", "vector_search_ef", FLAGS_vector_search_ef) << '\n';
   std::cout << '\n';
 }
 

diff --git a/src/benchmark/main.cc b/src/benchmark/main.cc
@@ -74,6 +74,8 @@ static std::string GetUsageMessage() {
   message += "\n  --vector_search_use_brute_force vector search flag use_brute_force, default(false)";
   message += "\n  --vector_search_enable_range_search vector search flag enable_range_search, default(false)";
   message += "\n  --vector_search_radius vector search flag radius, default(0.1)";
+  message += "\n  --vector_search_nprobe vector search flag nprobe, default(80)";
+  message += "\n  --vector_search_ef vector search flag ef, default(128)";
 
   return message;
 }

diff --git a/src/benchmark/operation.cc b/src/benchmark/operation.cc
@@ -75,6 +75,9 @@ DEFINE_validator(vector_search_filter_source, [](const char*, const std::string&
   return filter_source.empty() || filter_source == "SCALAR" || filter_source == "TABLE" || filter_source == "VECTOR_ID";
 });
 
+DEFINE_uint32(vector_search_nprobe, 80, "Vector search nprobe param");
+DEFINE_uint32(vector_search_ef, 128, "Vector search ef param");
+
 DECLARE_uint32(vector_dimension);
 DECLARE_string(vector_value_type);
 
@@ -1141,7 +1144,8 @@ Operation::Result VectorSearchOperation::ExecuteManualData(VectorIndexEntryPtr e
   search_param.with_table_data = FLAGS_vector_search_with_table_data;
   search_param.use_brute_force = FLAGS_vector_search_use_brute_force;
   search_param.topk = FLAGS_vector_search_topk;
-  search_param.extra_params.insert(std::make_pair(sdk::SearchExtraParamType::kEfSearch, 128));
+  search_param.extra_params.insert(std::make_pair(sdk::SearchExtraParamType::kEfSearch, FLAGS_vector_search_ef));
+  search_param.extra_params.insert(std::make_pair(sdk::SearchExtraParamType::kNprobe, FLAGS_vector_search_nprobe));
 
   std::string filter_type = dingodb::Helper::ToUpper(FLAGS_vector_search_filter_type);
   if (filter_type == "PRE") {