enhance: speed up search iterator stage 1 (#37947)

issue: #37548 Signed-off-by: Patrick Weizhi Xu <[email protected]>
milvus-io · Dec 26, 2024 · 85f462b · 85f462b
1 parent f49d618
commit 85f462b
Show file tree

Hide file tree

Showing 23 changed files with 1,970 additions and 117 deletions.
diff --git a/go.mod b/go.mod
@@ -64,6 +64,7 @@ require (
 	github.com/cenkalti/backoff/v4 v4.2.1
 	github.com/cockroachdb/redact v1.1.3
 	github.com/goccy/go-json v0.10.3
+	github.com/google/uuid v1.6.0
 	github.com/greatroar/blobloom v0.0.0-00010101000000-000000000000
 	github.com/hashicorp/golang-lru/v2 v2.0.7
 	github.com/jolestar/go-commons-pool/v2 v2.1.2
@@ -144,7 +145,6 @@ require (
 	github.com/golang/snappy v0.0.4 // indirect
 	github.com/google/flatbuffers v2.0.8+incompatible // indirect
 	github.com/google/s2a-go v0.1.7 // indirect
-	github.com/google/uuid v1.6.0 // indirect
 	github.com/googleapis/enterprise-certificate-proxy v0.3.2 // indirect
 	github.com/googleapis/gax-go/v2 v2.12.5 // indirect
 	github.com/gorilla/websocket v1.4.2 // indirect

diff --git a/internal/core/src/common/QueryInfo.h b/internal/core/src/common/QueryInfo.h
@@ -24,6 +24,12 @@
 
 namespace milvus {
 
+struct SearchIteratorV2Info {
+    std::string token = "";
+    uint32_t batch_size = 0;
+    std::optional<float> last_bound = std::nullopt;
+};
+
 struct SearchInfo {
     int64_t topk_{0};
     int64_t group_size_{1};
@@ -36,6 +42,7 @@ struct SearchInfo {
     tracer::TraceContext trace_ctx_;
     bool materialized_view_involved = false;
     bool iterative_filter_execution = false;
+    std::optional<SearchIteratorV2Info> iterator_v2_info_ = std::nullopt;
 };
 
 using SearchInfoPtr = std::shared_ptr<SearchInfo>;

diff --git a/internal/core/src/index/Utils.cpp b/internal/core/src/index/Utils.cpp
@@ -362,4 +362,38 @@ ReadDataFromFD(int fd, void* buf, size_t size, size_t chunk_size) {
     }
 }
 
+bool
+CheckAndUpdateKnowhereRangeSearchParam(const SearchInfo& search_info,
+                                       const int64_t topk,
+                                       const MetricType& metric_type,
+                                       knowhere::Json& search_config) {
+    const auto radius =
+        index::GetValueFromConfig<float>(search_info.search_params_, RADIUS);
+    if (!radius.has_value()) {
+        return false;
+    }
+
+    search_config[RADIUS] = radius.value();
+    // `range_search_k` is only used as one of the conditions for iterator early termination.
+    // not gurantee to return exactly `range_search_k` results, which may be more or less.
+    // set it to -1 will return all results in the range.
+    search_config[knowhere::meta::RANGE_SEARCH_K] = topk;
+
+    const auto range_filter =
+        GetValueFromConfig<float>(search_info.search_params_, RANGE_FILTER);
+    if (range_filter.has_value()) {
+        search_config[RANGE_FILTER] = range_filter.value();
+        CheckRangeSearchParam(
+            search_config[RADIUS], search_config[RANGE_FILTER], metric_type);
+    }
+
+    const auto page_retain_order =
+        GetValueFromConfig<bool>(search_info.search_params_, PAGE_RETAIN_ORDER);
+    if (page_retain_order.has_value()) {
+        search_config[knowhere::meta::RETAIN_ITERATOR_ORDER] =
+            page_retain_order.value();
+    }
+    return true;
+}
+
 }  // namespace milvus::index
diff --git a/internal/core/src/index/Utils.h b/internal/core/src/index/Utils.h
@@ -30,6 +30,8 @@
 
 #include "common/Types.h"
 #include "common/FieldData.h"
+#include "common/QueryInfo.h"
+#include "common/RangeSearchHelper.h"
 #include "index/IndexInfo.h"
 #include "storage/Types.h"
 
@@ -147,4 +149,10 @@ AssembleIndexDatas(std::map<std::string, FieldDataChannelPtr>& index_datas,
 void
 ReadDataFromFD(int fd, void* buf, size_t size, size_t chunk_size = 0x7ffff000);
 
+bool
+CheckAndUpdateKnowhereRangeSearchParam(const SearchInfo& search_info,
+                                       const int64_t topk,
+                                       const MetricType& metric_type,
+                                       knowhere::Json& search_config);
+
 }  // namespace milvus::index
diff --git a/internal/core/src/index/VectorDiskIndex.cpp b/internal/core/src/index/VectorDiskIndex.cpp
@@ -266,32 +266,9 @@ VectorDiskAnnIndex<T>::Query(const DatasetPtr dataset,
     search_config[DISK_ANN_PREFIX_PATH] = local_index_path_prefix;
 
     auto final = [&] {
-        auto radius =
-            GetValueFromConfig<float>(search_info.search_params_, RADIUS);
-        if (radius.has_value()) {
-            search_config[RADIUS] = radius.value();
-            // `range_search_k` is only used as one of the conditions for iterator early termination.
-            // not gurantee to return exactly `range_search_k` results, which may be more or less.
-            // set it to -1 will return all results in the range.
-            search_config[knowhere::meta::RANGE_SEARCH_K] = topk;
-            auto range_filter = GetValueFromConfig<float>(
-                search_info.search_params_, RANGE_FILTER);
-            if (range_filter.has_value()) {
-                search_config[RANGE_FILTER] = range_filter.value();
-                CheckRangeSearchParam(search_config[RADIUS],
-                                      search_config[RANGE_FILTER],
-                                      GetMetricType());
-            }
-
-            auto page_retain_order = GetValueFromConfig<bool>(
-                search_info.search_params_, PAGE_RETAIN_ORDER);
-            if (page_retain_order.has_value()) {
-                search_config[knowhere::meta::RETAIN_ITERATOR_ORDER] =
-                    page_retain_order.value();
-            }
-
+        if (CheckAndUpdateKnowhereRangeSearchParam(
+                search_info, topk, GetMetricType(), search_config)) {
             auto res = index_.RangeSearch(dataset, search_config, bitset);
-
             if (!res.has_value()) {
                 PanicInfo(ErrorCode::UnexpectedError,
                           fmt::format("failed to range search: {}: {}",

diff --git a/internal/core/src/index/VectorMemIndex.cpp b/internal/core/src/index/VectorMemIndex.cpp
@@ -380,16 +380,8 @@ VectorMemIndex<T>::Query(const DatasetPtr dataset,
     // TODO :: check dim of search data
     auto final = [&] {
         auto index_type = GetIndexType();
-        if (CheckKeyInConfig(search_conf, RADIUS)) {
-            if (CheckKeyInConfig(search_conf, RANGE_FILTER)) {
-                CheckRangeSearchParam(search_conf[RADIUS],
-                                      search_conf[RANGE_FILTER],
-                                      GetMetricType());
-            }
-            // `range_search_k` is only used as one of the conditions for iterator early termination.
-            // not gurantee to return exactly `range_search_k` results, which may be more or less.
-            // set it to -1 will return all results in the range.
-            search_conf[knowhere::meta::RANGE_SEARCH_K] = topk;
+        if (CheckAndUpdateKnowhereRangeSearchParam(
+                search_info, topk, GetMetricType(), search_conf)) {
             milvus::tracer::AddEvent("start_knowhere_index_range_search");
             auto res = index_.RangeSearch(dataset, search_conf, bitset);
             milvus::tracer::AddEvent("finish_knowhere_index_range_search");