search allocator

Signed-off-by: zourunxin.zrx <zourunxin.zrx@oceanbase.com>

Author: zourunxin.zrx

include/vsag/index.h

@@ -171,6 +171,15 @@ class Index {
         throw std::runtime_error("Index doesn't support new filter");
     }
 
+    virtual tl::expected<DatasetPtr, Error>
+    KnnSearch(const DatasetPtr& query,
+              int64_t k,
+              const std::string& parameters,
+              const FilterPtr& filter,
+              Allocator *allocator) const {
+        throw std::runtime_error("Index doesn't support new filter");
+    }
+
     /**
       * @brief Performing single KNN search on index
       *

src/algorithm/hnswlib/algorithm_interface.h

@@ -46,6 +46,7 @@ class AlgorithmInterface {
               size_t ef,
               const vsag::FilterPtr is_id_allowed = nullptr,
               float skip_ratio = 0.9f,
+              vsag::Allocator *allocator = nullptr,
               vsag::IteratorFilterContext* iter_ctx = nullptr,
               bool is_last_filter = false) const = 0;
 

src/algorithm/hnswlib/hnswalg.cpp

@@ -439,14 +439,16 @@ HierarchicalNSW::searchBaseLayerST(InnerIdType ep_id,
                                    size_t ef,
                                    const vsag::FilterPtr is_id_allowed,
                                    const float skip_ratio,
+                                   vsag::Allocator *allocator,
                                    vsag::IteratorFilterContext* iter_ctx) const {
     vsag::LinearCongruentialGenerator generator;
     VisitedListPtr vl = visited_list_pool_->getFreeVisitedList();
     vl_type* visited_array = vl->mass;
     vl_type visited_array_tag = vl->curV;
+    vsag::Allocator *search_allocator = allocator == nullptr ? allocator_ : allocator;
 
-    MaxHeap top_candidates(allocator_);
-    MaxHeap candidate_set(allocator_);
+    MaxHeap top_candidates(search_allocator);
+    MaxHeap candidate_set(search_allocator);
 
     float valid_ratio = is_id_allowed ? is_id_allowed->ValidRatio() : 1.0F;
     float skip_threshold = valid_ratio == 1.0F ? 0 : (1 - ((1 - valid_ratio) * skip_ratio));
@@ -1477,16 +1479,18 @@ HierarchicalNSW::searchKnn(const void* query_data,
                            uint64_t ef,
                            const vsag::FilterPtr is_id_allowed,
                            const float skip_ratio,
+                           vsag::Allocator *allocator,
                            vsag::IteratorFilterContext* iter_ctx,
                            bool is_last_filter) const {
     std::shared_lock resize_lock(resize_mutex_);
     std::priority_queue<std::pair<float, LabelType>> result;
     if (cur_element_count_ == 0)
         return result;
 
+    vsag::Allocator *search_alloctor = allocator == nullptr ? allocator_ : allocator;
     std::shared_ptr<float[]> normalize_query;
     normalizeVector(query_data, normalize_query);
-    MaxHeap top_candidates(allocator_);
+    MaxHeap top_candidates(search_alloctor);
     if (iter_ctx != nullptr && !iter_ctx->IsFirstUsed()) {
         if (iter_ctx->Empty())
             return result;
@@ -1504,6 +1508,7 @@ HierarchicalNSW::searchKnn(const void* query_data,
                                                         std::max(ef, k),
                                                         is_id_allowed,
                                                         skip_ratio,
+                                                        allocator,
                                                         iter_ctx);
     } else {
         int64_t currObj;
@@ -1545,10 +1550,10 @@ HierarchicalNSW::searchKnn(const void* query_data,
 
         if (num_deleted_ == 0) {
             top_candidates = searchBaseLayerST<false, true>(
-                currObj, query_data, std::max(ef, k), is_id_allowed, skip_ratio, iter_ctx);
+                currObj, query_data, std::max(ef, k), is_id_allowed, skip_ratio, allocator, iter_ctx);
         } else {
             top_candidates = searchBaseLayerST<true, true>(
-                currObj, query_data, std::max(ef, k), is_id_allowed, skip_ratio, iter_ctx);
+                currObj, query_data, std::max(ef, k), is_id_allowed, skip_ratio, allocator, iter_ctx);
         }
     }
 
@@ -1669,6 +1674,7 @@ HierarchicalNSW::searchBaseLayerST<false, false>(
     size_t ef,
     const vsag::FilterPtr is_id_allowed,
     const float skip_ratio,
+    vsag::Allocator *allocator,
     vsag::IteratorFilterContext* iter_ctx = nullptr) const;
 
 template MaxHeap

src/algorithm/hnswlib/hnswalg.h

@@ -255,6 +255,7 @@ class HierarchicalNSW : public AlgorithmInterface<float> {
                       size_t ef,
                       const vsag::FilterPtr is_id_allowed = nullptr,
                       const float skip_ratio = 0.9f,
+                      vsag::Allocator *allocator = nullptr,
                       vsag::IteratorFilterContext* iter_ctx = nullptr) const;
 
     template <bool has_deletions, bool collect_metrics = false>
@@ -411,6 +412,7 @@ class HierarchicalNSW : public AlgorithmInterface<float> {
               uint64_t ef,
               const vsag::FilterPtr is_id_allowed = nullptr,
               const float skip_ratio = 0.9f,
+              vsag::Allocator *allocator = nullptr,
               vsag::IteratorFilterContext* iter_ctx = nullptr,
               bool is_last_filter = false) const override;
 

src/algorithm/hnswlib/hnswalg_static.h

@@ -1467,6 +1467,7 @@ class StaticHierarchicalNSW : public AlgorithmInterface<float> {
               uint64_t ef,
               const vsag::FilterPtr is_id_allowed = nullptr,
               const float skip_ratio = 0.9f,
+              vsag::Allocator *allocator = nullptr,
               vsag::IteratorFilterContext* iter_ctx = nullptr,
               bool is_last_filter = false) const override {
         std::priority_queue<std::pair<float, LabelType>> result;

src/index/hnsw.cpp

@@ -204,6 +204,7 @@ HNSW::knn_search(const DatasetPtr& query,
                  int64_t k,
                  const std::string& parameters,
                  const FilterPtr& filter_ptr,
+                 vsag::Allocator *allocator,
                  vsag::IteratorContext** iter_ctx,
                  bool is_last_filter) const {
 #ifndef ENABLE_TESTS
@@ -216,6 +217,7 @@ HNSW::knn_search(const DatasetPtr& query,
             ret->Dim(0)->NumElements(1);
             return ret;
         }
+        vsag::Allocator *search_alloctor = allocator == nullptr ? allocator_.get() : allocator;
 
         // check query vector
         CHECK_ARGUMENT(query->GetNumElements() == 1, "query dataset should contain 1 vector only");
@@ -238,7 +240,7 @@ HNSW::knn_search(const DatasetPtr& query,
 
         if (iter_ctx != nullptr && *iter_ctx == nullptr) {
             auto* filter_context = new IteratorFilterContext();
-            filter_context->init(alg_hnsw_->getMaxElements(), params.ef_search, allocator_.get());
+            filter_context->init(alg_hnsw_->getMaxElements(), params.ef_search, search_alloctor);
             *iter_ctx = filter_context;
         }
         IteratorFilterContext* iter_filter_ctx = nullptr;
@@ -260,6 +262,7 @@ HNSW::knn_search(const DatasetPtr& query,
                                            std::max(params.ef_search, k),
                                            filter_ptr,
                                            params.skip_ratio,
+                                           allocator,
                                            iter_filter_ctx,
                                            is_last_filter);
         } catch (const std::runtime_error& e) {
@@ -299,7 +302,7 @@ HNSW::knn_search(const DatasetPtr& query,
             results.pop();
         }
         auto [dataset_results, dists, ids] =
-            CreateFastDataset(static_cast<int64_t>(results.size()), allocator_.get());
+            CreateFastDataset(static_cast<int64_t>(results.size()), search_alloctor);
 
         for (auto j = static_cast<int64_t>(results.size() - 1); j >= 0; --j) {
             dists[j] = results.top().first;

src/index/hnsw.h

@@ -111,6 +111,15 @@ class HNSW : public Index {
         SAFE_CALL(return this->knn_search(query, k, parameters, filter));
     }
 
+    tl::expected<DatasetPtr, Error>
+    KnnSearch(const DatasetPtr& query,
+              int64_t k,
+              const std::string& parameters,
+              const FilterPtr& filter,
+              Allocator *allocator) const override {
+        SAFE_CALL(return this->knn_search(query, k, parameters, filter, allocator));
+    }
+
     tl::expected<DatasetPtr, Error>
     KnnSearch(const DatasetPtr& query,
               int64_t k,
@@ -119,7 +128,7 @@ class HNSW : public Index {
               vsag::IteratorContext*& filter_ctx,
               bool is_last_search) const override {
         SAFE_CALL(
-            return this->knn_search(query, k, parameters, filter, &filter_ctx, is_last_search));
+            return this->knn_search(query, k, parameters, filter, nullptr, &filter_ctx, is_last_search));
     }
 
     tl::expected<DatasetPtr, Error>
@@ -282,6 +291,7 @@ class HNSW : public Index {
                int64_t k,
                const std::string& parameters,
                const FilterPtr& filter_ptr,
+               vsag::Allocator *allocator = nullptr,
                vsag::IteratorContext** iter_ctx = nullptr,
                bool is_last_filter = false) const;