From 4fc740c69f4ec3df74478785126ba291cd6328e1 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Wed, 21 Oct 2026 12:38:46 -0700 Subject: [PATCH 01/33] Add filtered search on the dynamic Vamana search buffer Reuse the MutableBuffer valid/invalid mechanism: a FilteredValidBuilder marks an entry invalid if it is deleted or the predicate rejects its external id. New MutableVamanaIndex::search overload taking a predicate, padding rows with id -1 / +inf and returning per-query found counts. Supplement path checks the predicate before computing the distance. Orchestrator virtual appended last. --- include/svs/index/vamana/dynamic_index.h | 122 +++++++++++++++- include/svs/index/vamana/extensions.h | 26 ++++ include/svs/orchestrators/dynamic_vamana.h | 58 ++++++++ tests/svs/index/vamana/dynamic_index.cpp | 157 +++++++++++++++++++++ 4 files changed, 361 insertions(+), 2 deletions(-) diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 9336f82e1..5c75d4445 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -18,6 +18,7 @@ // stdlib #include +#include #include #include @@ -109,11 +110,49 @@ class ValidBuilder { const std::vector& status_; }; +// Like `ValidBuilder`, but an entry is also invalid if `pred(external_id)` is false. +template class FilteredValidBuilder { + public: + FilteredValidBuilder( + const std::vector& status, + const IDTranslator& translator, + const Pred& pred + ) + : status_{status} + , translator_{translator} + , pred_{pred} {} + + template bool is_valid(I internal_id, size_t external_id) const { + return getindex(status_, internal_id) != SlotMetadata::Deleted && + pred_(external_id); + } + + template + PredicatedSearchNeighbor operator()(I i, float distance) const { + bool valid = getindex(status_, i) != SlotMetadata::Deleted && + pred_(translator_.get_external(i)); + return PredicatedSearchNeighbor(i, distance, valid); + } + + private: + const std::vector& status_; + const IDTranslator& translator_; + const Pred& pred_; +}; + +// Data types for which the filtered `search` overload is valid (identity accessor). +template inline constexpr bool filtered_search_supported = false; +template +inline constexpr bool filtered_search_supported> = true; + template class MutableVamanaIndex { friend class MultiMutableVamanaIndex; public: + // Present only on indexes that support the filtered `search` overload. + static constexpr bool has_filtered_search = filtered_search_supported; + // Traits static constexpr bool supports_insertions = true; static constexpr bool supports_deletions = true; @@ -575,11 +614,13 @@ class MutableVamanaIndex { // This is an internal method, mostly used to help implement the batch iterator. ValidBuilder internal_search_builder() const { return ValidBuilder{status_}; } + template auto greedy_search_closure( + const Builder& builder, GreedySearchPrefetchParameters prefetch_parameters, const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) ) const { - return [&, prefetch_parameters]( + return [&, builder, prefetch_parameters]( const auto& query, auto& accessor, auto& distance, auto& buffer ) { // Perform the greedy search using the provided resources. @@ -591,7 +632,7 @@ class MutableVamanaIndex { distance, buffer, vamana::EntryPointInitializer{lib::as_const_span(entry_point_)}, - internal_search_builder(), + builder, prefetch_parameters, cancel ); @@ -601,6 +642,15 @@ class MutableVamanaIndex { }; } + auto greedy_search_closure( + GreedySearchPrefetchParameters prefetch_parameters, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + ) const { + return greedy_search_closure( + internal_search_builder(), prefetch_parameters, cancel + ); + } + // Single Search template void search( @@ -667,6 +717,74 @@ class MutableVamanaIndex { translate_to_external(results.indices()); } + /// + /// @brief Filtered top-k search: only live entries whose external id satisfies `pred` + /// are returned. Invalid entries are navigated through but do not count toward the + /// search window. + /// + /// @param pred Callable `bool(size_t external_id)`. Must be thread safe. + /// @param found Resized to `queries.size()`; receives the number of results per query. + /// + /// Rows with fewer than `n_neighbors` results are padded with id `size_t(-1)` and + /// distance `+inf`. Requires `filtered_search_supported`. + /// + template + requires(std::is_invocable_r_v) + void search( + QueryResultView results, + const Queries& queries, + const search_parameters_type& sp, + const Pred& pred, + std::vector& found + ) { + static_assert( + filtered_search_supported, + "Data type needs a custom search path (e.g. reranking)." + ); + found.assign(queries.size(), 0); + threads::parallel_for( + threadpool_, + threads::StaticPartition{queries.size()}, + [&](const auto is, uint64_t SVS_UNUSED(tid)) { + size_t num_neighbors = results.n_neighbors(); + auto buffer = search_buffer_type{ + sp.buffer_config_.get_total_capacity() < num_neighbors + ? SearchBufferConfig{num_neighbors} + : sp.buffer_config_, + distance::comparator(distance_), + sp.search_buffer_visited_set_}; + auto prefetch_parameters = GreedySearchPrefetchParameters{ + sp.prefetch_lookahead_, sp.prefetch_step_}; + auto distance = extensions::per_thread_batch_search_setup(data_, distance_); + auto accessor = data::GetDatumAccessor(); + auto builder = FilteredValidBuilder{status_, translator_, pred}; + // The closure holds `cancel` by reference, so it must outlive the closure. + const lib::DefaultPredicate cancel = lib::Returns(lib::Const()); + auto search = greedy_search_closure(builder, prefetch_parameters, cancel); + + for (auto i : is) { + auto query = queries.get_datum(i); + search(query, accessor, distance, buffer); + extensions::check_and_supplement_search_buffer( + *this, buffer, query, builder + ); + // Everything in the buffer is valid here (cleanup() ran, and + // supplemented entries are checked before insertion). + size_t n = std::min(num_neighbors, buffer.size()); + for (size_t j = 0; j < n; ++j) { + results.index(i, j) = translate_internal_id(buffer[j].id()); + results.distance(i, j) = buffer[j].distance(); + } + for (size_t j = n; j < num_neighbors; ++j) { + results.index(i, j) = static_cast(-1); + results.distance(i, j) = std::numeric_limits::infinity(); + } + found[i] = n; + } + } + ); + } + /// /// @brief Return a unique instance of the distance function. /// diff --git a/include/svs/index/vamana/extensions.h b/include/svs/index/vamana/extensions.h index 96ffeaa06..5b73cf7e9 100644 --- a/include/svs/index/vamana/extensions.h +++ b/include/svs/index/vamana/extensions.h @@ -449,6 +449,32 @@ void check_and_supplement_search_buffer( } } +/// Filtered variant: only ids accepted by `builder` are inserted, and the predicate is +/// evaluated before the distance is computed. Does not change the unfiltered path above. +template +void check_and_supplement_search_buffer( + const Index& index, + SearchBuffer& search_buffer, + const Query& query, + const Builder& builder +) { + if (search_buffer.valid() < search_buffer.target_window() && + search_buffer.valid() < index.size()) { + search_buffer.sort(); + for (auto external_id : index.external_ids()) { + auto internal_id = index.translate_external_id(external_id); + if (!builder.is_valid(internal_id, external_id)) { + continue; + } + auto dist = index.get_distance(external_id, query); + search_buffer.insert(builder(internal_id, dist)); + if (search_buffer.valid() >= search_buffer.target_window()) { + break; + } + } + } +} + /// Customization point object for processing single queries. inline constexpr VamanaSingleSearchType single_search{}; diff --git a/include/svs/orchestrators/dynamic_vamana.h b/include/svs/orchestrators/dynamic_vamana.h index 067c7c486..e82c5c0f9 100644 --- a/include/svs/orchestrators/dynamic_vamana.h +++ b/include/svs/orchestrators/dynamic_vamana.h @@ -22,6 +22,7 @@ #include "svs/orchestrators/vamana.h" // stdlib +#include #include namespace svs { @@ -56,6 +57,17 @@ class DynamicVamanaInterface : public VamanaInterface { // Check if the index has been consolidated - no `Deleted` entries remain. virtual bool is_consolidated() const = 0; + + // Filtered top-k on the search buffer; `filter` takes an external id. Returns false, + // leaving `results` untouched, if the data type does not support it. + // Keep last: new virtuals must be appended. + virtual bool filtered_search( + svs::QueryResultView results, + AnonymousArray<2> queries, + const index::vamana::VamanaSearchParameters& search_parameters, + const std::function& filter, + std::vector& found + ) = 0; }; template @@ -109,6 +121,31 @@ class DynamicVamanaImpl : public VamanaImpl results, + AnonymousArray<2> queries, + const index::vamana::VamanaSearchParameters& search_parameters, + const std::function& filter, + std::vector& found + ) override { + if constexpr (requires { requires Impl::has_filtered_search; }) { + lib::match( + QueryTypes{}, + queries.type(), + [&](lib::Type SVS_UNUSED(type)) { + const auto view = data::ConstSimpleDataView(queries); + impl().search(results, view, search_parameters, filter, found); + }, + [&](svs::DataType data_type) { + throw ANNEXCEPTION("Unsupported query datatype! Got: {}.", data_type); + } + ); + return true; + } else { + return false; + } + } }; // Forward Declaractions. @@ -609,6 +646,27 @@ class DynamicVamana : public manager::IndexManager { /// @brief Check if the index has been consolidated - no `Deleted` entries remain. bool is_consolidated() const { return impl_->is_consolidated(); } + + /// + /// @brief Filtered top-k search on the search buffer, see + /// ``MutableVamanaIndex::search`` (filtered overload). + /// + /// @param filter Takes an external id. + /// @param found Receives the number of results per query. Rows are padded with id + /// ``size_t(-1)`` and distance ``+inf``. + /// @return False if the data type is unsupported, in which case nothing is written. + template + bool filtered_search( + QueryResultView results, + data::ConstSimpleDataView queries, + const VamanaSearchParameters& search_parameters, + const std::function& filter, + std::vector& found + ) { + return impl_->filtered_search( + results, AnonymousArray<2>(queries), search_parameters, filter, found + ); + } }; /// diff --git a/tests/svs/index/vamana/dynamic_index.cpp b/tests/svs/index/vamana/dynamic_index.cpp index 94daddcf1..7a7fadabe 100644 --- a/tests/svs/index/vamana/dynamic_index.cpp +++ b/tests/svs/index/vamana/dynamic_index.cpp @@ -23,11 +23,14 @@ #include #include #include +#include #include #include +#include #include #include #include +#include // svs #include "svs/core/recall.h" @@ -464,3 +467,157 @@ CATCH_TEST_CASE("MutableVamana Index Memory Usage", "[graph_index][dynamic_index const size_t usage = index.get_memory_breakdown().total(); CATCH_REQUIRE(usage == expected_total_bytes); } + +CATCH_TEST_CASE( + "MutableVamana Index Filtered Search", "[graph_index][dynamic_index][filtered]" +) { + const size_t num_threads = 2; + const size_t n = 3000; + const size_t dim = 32; + const size_t num_queries = 50; + const size_t k = 10; + using Distance = svs::distance::DistanceL2; + + std::mt19937 rng(42); + std::normal_distribution normal; + auto data = svs::data::BlockedData(n, dim); + std::vector> raw(n, std::vector(dim)); + for (size_t i = 0; i < n; ++i) { + for (auto& x : raw[i]) { + x = normal(rng); + } + data.set_datum(i, raw[i]); + } + auto queries = svs::data::SimpleData(num_queries, dim); + for (size_t i = 0; i < num_queries; ++i) { + std::vector q(dim); + for (auto& x : q) { + x = normal(rng); + } + queries.set_datum(i, q); + } + + // External ids differ from internal ids. + std::vector ids(n); + std::iota(ids.begin(), ids.end(), 7); + svs::index::vamana::VamanaBuildParameters parameters{1.2, 32, 64, 64, 32, true}; + auto index = svs::index::vamana::MutableVamanaIndex( + parameters, std::move(data), ids, Distance(), num_threads + ); + + auto sp = svs::index::vamana::VamanaSearchParameters{}; + sp.buffer_config_ = svs::index::vamana::SearchBufferConfig{50}; + + // Brute force over live ids passing `pred`. + std::unordered_set deleted; + auto brute_force = [&](size_t q, auto&& pred) { + std::vector> all; + for (size_t i = 0; i < n; ++i) { + size_t e = ids[i]; + if (deleted.contains(e) || !pred(e)) { + continue; + } + float d = 0; + for (size_t j = 0; j < dim; ++j) { + float diff = queries.get_datum(q)[j] - raw[i][j]; + d += diff * diff; + } + all.emplace_back(d, e); + } + std::sort(all.begin(), all.end()); + std::vector out; + for (size_t i = 0; i < std::min(k, all.size()); ++i) { + out.push_back(all[i].second); + } + return out; + }; + + auto check = [&](auto&& pred, bool expect_full) { + auto results = svs::QueryResult(num_queries, k); + std::vector found; + index.search(results.view(), queries.cview(), sp, pred, found); + CATCH_REQUIRE(found.size() == num_queries); + size_t hits = 0, total = 0; + for (size_t q = 0; q < num_queries; ++q) { + auto gt = brute_force(q, pred); + CATCH_REQUIRE(found[q] <= k); + if (expect_full) { + CATCH_REQUIRE(found[q] == k); + } + for (size_t j = 0; j < found[q]; ++j) { + size_t id = results.index(q, j); + CATCH_REQUIRE(pred(id)); + CATCH_REQUIRE(!deleted.contains(id)); + hits += std::count(gt.begin(), gt.end(), id); + } + for (size_t j = found[q]; j < k; ++j) { + CATCH_REQUIRE(results.index(q, j) == static_cast(-1)); + CATCH_REQUIRE( + results.distance(q, j) == std::numeric_limits::infinity() + ); + } + total += gt.size(); + } + return double(hits) / double(total); + }; + + auto mod10 = [](size_t e) { return e % 10 == 0; }; + + CATCH_SECTION("recall against brute force") { + CATCH_REQUIRE(check(mod10, true) >= 0.95); + } + + CATCH_SECTION("deleted ids are never returned") { + std::vector to_delete; + for (size_t i = 0; i < n; i += 3) { + to_delete.push_back(ids[i]); + } + index.delete_entries(to_delete); + deleted.insert(to_delete.begin(), to_delete.end()); + CATCH_REQUIRE(check(mod10, true) >= 0.9); + CATCH_REQUIRE(check([](size_t) { return true; }, true) >= 0.9); + } + + CATCH_SECTION("fewer than k passing ids are padded") { + auto few = [](size_t e) { return e == 100 || e == 200 || e == 300; }; + check(few, false); + auto results = svs::QueryResult(num_queries, k); + std::vector found; + index.search(results.view(), queries.cview(), sp, few, found); + for (size_t q = 0; q < num_queries; ++q) { + CATCH_REQUIRE(found[q] == 3); + } + auto none = [](size_t) { return false; }; + index.search(results.view(), queries.cview(), sp, none, found); + for (size_t q = 0; q < num_queries; ++q) { + CATCH_REQUIRE(found[q] == 0); + } + } + + CATCH_SECTION("orchestrator overload") { + auto dyn = svs::DynamicVamana::build( + parameters, + [&] { + auto d = svs::data::SimpleData(n, dim); + for (size_t i = 0; i < n; ++i) { + d.set_datum(i, raw[i]); + } + return d; + }(), + ids, + Distance(), + num_threads + ); + auto results = svs::QueryResult(num_queries, k); + std::vector found; + CATCH_REQUIRE(dyn.filtered_search( + results.view(), queries.cview(), sp, std::function(mod10), found + )); + for (size_t q = 0; q < num_queries; ++q) { + CATCH_REQUIRE(found[q] == k); + for (size_t j = 0; j < k; ++j) { + CATCH_REQUIRE(mod10(results.index(q, j))); + } + } + } +} From 8b0f4371abb15fb39fcad0655a3173212983e804 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Wed, 21 Oct 2026 14:53:14 -0700 Subject: [PATCH 02/33] Make filtered search an optional argument of the existing search Replace the separate filtered search overload with two optional trailing arguments on MutableVamanaIndex::search (an id predicate and a per-query result count). Without a predicate the existing code path runs unchanged. check_and_supplement_search_buffer gets an optional predicate that is checked before the distance is computed. The filter predicate type is now a type-erased std::function, so FilteredValidBuilder is no longer a template. The orchestrator keeps one appended virtual that forwards to the index search. --- include/svs/index/vamana/dynamic_index.h | 126 ++++++++------------- include/svs/index/vamana/extensions.h | 36 ++---- include/svs/orchestrators/dynamic_vamana.h | 22 ++-- tests/svs/index/vamana/dynamic_index.cpp | 54 ++++++++- 4 files changed, 120 insertions(+), 118 deletions(-) diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 5c75d4445..6a0c0cefc 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -111,22 +111,17 @@ class ValidBuilder { }; // Like `ValidBuilder`, but an entry is also invalid if `pred(external_id)` is false. -template class FilteredValidBuilder { +class FilteredValidBuilder { public: FilteredValidBuilder( const std::vector& status, const IDTranslator& translator, - const Pred& pred + const extensions::IDPredicate& pred ) : status_{status} , translator_{translator} , pred_{pred} {} - template bool is_valid(I internal_id, size_t external_id) const { - return getindex(status_, internal_id) != SlotMetadata::Deleted && - pred_(external_id); - } - template PredicatedSearchNeighbor operator()(I i, float distance) const { bool valid = getindex(status_, i) != SlotMetadata::Deleted && @@ -137,7 +132,7 @@ template class FilteredValidBuilder { private: const std::vector& status_; const IDTranslator& translator_; - const Pred& pred_; + const extensions::IDPredicate& pred_; }; // Data types for which the filtered `search` overload is valid (identity accessor). @@ -673,8 +668,17 @@ class MutableVamanaIndex { QueryResultView results, const Queries& queries, const search_parameters_type& sp, - const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()), + const extensions::IDPredicate* pred = nullptr, + std::vector* found = nullptr ) { + // Optional filter: only live entries with `(*pred)(external_id)` are returned; + // others are navigated through but do not count toward the search window. Rows + // with fewer than `n_neighbors` results are padded with id `size_t(-1)` and + // distance `+inf`; `found` (if given) receives the count per query. + if (found != nullptr) { + found->assign(queries.size(), 0); + } threads::parallel_for( threadpool_, threads::StaticPartition{queries.size()}, @@ -693,6 +697,39 @@ class MutableVamanaIndex { auto scratch = extensions::per_thread_batch_search_setup(data_, distance_); + if (pred != nullptr) { + if constexpr (filtered_search_supported) { + auto builder = FilteredValidBuilder{status_, translator_, *pred}; + auto search = + greedy_search_closure(builder, prefetch_parameters, cancel); + auto accessor = data::GetDatumAccessor(); + for (auto i : is) { + auto query = queries.get_datum(i); + search(query, accessor, scratch, buffer); + extensions::check_and_supplement_search_buffer( + *this, buffer, query, pred + ); + size_t n = std::min(num_neighbors, buffer.size()); + for (size_t j = 0; j < num_neighbors; ++j) { + bool ok = j < n; + results.index(i, j) = + ok ? translate_internal_id(buffer[j].id()) : size_t(-1); + results.distance(i, j) = + ok ? buffer[j].distance() + : std::numeric_limits::infinity(); + } + if (found != nullptr) { + (*found)[i] = n; + } + } + return; + } else { + throw ANNEXCEPTION( + "Filtered search is not supported for this data type." + ); + } + } + extensions::per_thread_batch_search( data_, buffer, @@ -708,7 +745,8 @@ class MutableVamanaIndex { ); // Check if request to cancel the search - if (cancel()) { + // (Filtered results already hold external ids.) + if (cancel() || pred != nullptr) { return; } @@ -717,74 +755,6 @@ class MutableVamanaIndex { translate_to_external(results.indices()); } - /// - /// @brief Filtered top-k search: only live entries whose external id satisfies `pred` - /// are returned. Invalid entries are navigated through but do not count toward the - /// search window. - /// - /// @param pred Callable `bool(size_t external_id)`. Must be thread safe. - /// @param found Resized to `queries.size()`; receives the number of results per query. - /// - /// Rows with fewer than `n_neighbors` results are padded with id `size_t(-1)` and - /// distance `+inf`. Requires `filtered_search_supported`. - /// - template - requires(std::is_invocable_r_v) - void search( - QueryResultView results, - const Queries& queries, - const search_parameters_type& sp, - const Pred& pred, - std::vector& found - ) { - static_assert( - filtered_search_supported, - "Data type needs a custom search path (e.g. reranking)." - ); - found.assign(queries.size(), 0); - threads::parallel_for( - threadpool_, - threads::StaticPartition{queries.size()}, - [&](const auto is, uint64_t SVS_UNUSED(tid)) { - size_t num_neighbors = results.n_neighbors(); - auto buffer = search_buffer_type{ - sp.buffer_config_.get_total_capacity() < num_neighbors - ? SearchBufferConfig{num_neighbors} - : sp.buffer_config_, - distance::comparator(distance_), - sp.search_buffer_visited_set_}; - auto prefetch_parameters = GreedySearchPrefetchParameters{ - sp.prefetch_lookahead_, sp.prefetch_step_}; - auto distance = extensions::per_thread_batch_search_setup(data_, distance_); - auto accessor = data::GetDatumAccessor(); - auto builder = FilteredValidBuilder{status_, translator_, pred}; - // The closure holds `cancel` by reference, so it must outlive the closure. - const lib::DefaultPredicate cancel = lib::Returns(lib::Const()); - auto search = greedy_search_closure(builder, prefetch_parameters, cancel); - - for (auto i : is) { - auto query = queries.get_datum(i); - search(query, accessor, distance, buffer); - extensions::check_and_supplement_search_buffer( - *this, buffer, query, builder - ); - // Everything in the buffer is valid here (cleanup() ran, and - // supplemented entries are checked before insertion). - size_t n = std::min(num_neighbors, buffer.size()); - for (size_t j = 0; j < n; ++j) { - results.index(i, j) = translate_internal_id(buffer[j].id()); - results.distance(i, j) = buffer[j].distance(); - } - for (size_t j = n; j < num_neighbors; ++j) { - results.index(i, j) = static_cast(-1); - results.distance(i, j) = std::numeric_limits::infinity(); - } - found[i] = n; - } - } - ); - } - /// /// @brief Return a unique instance of the distance function. /// diff --git a/include/svs/index/vamana/extensions.h b/include/svs/index/vamana/extensions.h index 5b73cf7e9..edcddec2c 100644 --- a/include/svs/index/vamana/extensions.h +++ b/include/svs/index/vamana/extensions.h @@ -27,6 +27,9 @@ #include "svs/lib/preprocessor.h" #include "svs/lib/threads.h" +// stdlib +#include + namespace svs::index::vamana::extensions { ///// @@ -425,48 +428,33 @@ struct VamanaSingleSearchType { } }; +/// Predicate on external ids used by filtered search. +using IDPredicate = std::function; + /// In rare cases, the search buffer may not be filled with enough results. /// This can occur in dynamic indexes when many vectors have been deleted /// and the graph becomes sparsely connected. It's a corner case and should /// not happen frequently, but when it does, we may need to supplement the buffer /// with additional results. +/// If `pred` is given, entries whose external id it rejects are skipped (checked before +/// the distance is computed). template -void check_and_supplement_search_buffer( - const Index& index, SearchBuffer& search_buffer, const Query& query -) { - if (search_buffer.valid() < search_buffer.target_window() && - search_buffer.valid() < index.size()) { - search_buffer.sort(); - for (auto external_id : index.external_ids()) { - auto internal_id = index.translate_external_id(external_id); - auto dist = index.get_distance(external_id, query); - auto builder = index.internal_search_builder(); - search_buffer.insert(builder(internal_id, dist)); - if (search_buffer.valid() >= search_buffer.target_window()) { - break; - } - } - } -} - -/// Filtered variant: only ids accepted by `builder` are inserted, and the predicate is -/// evaluated before the distance is computed. Does not change the unfiltered path above. -template void check_and_supplement_search_buffer( const Index& index, SearchBuffer& search_buffer, const Query& query, - const Builder& builder + const IDPredicate* pred = nullptr ) { if (search_buffer.valid() < search_buffer.target_window() && search_buffer.valid() < index.size()) { search_buffer.sort(); for (auto external_id : index.external_ids()) { - auto internal_id = index.translate_external_id(external_id); - if (!builder.is_valid(internal_id, external_id)) { + if (pred != nullptr && !(*pred)(external_id)) { continue; } + auto internal_id = index.translate_external_id(external_id); auto dist = index.get_distance(external_id, query); + auto builder = index.internal_search_builder(); search_buffer.insert(builder(internal_id, dist)); if (search_buffer.valid() >= search_buffer.target_window()) { break; diff --git a/include/svs/orchestrators/dynamic_vamana.h b/include/svs/orchestrators/dynamic_vamana.h index e82c5c0f9..403080185 100644 --- a/include/svs/orchestrators/dynamic_vamana.h +++ b/include/svs/orchestrators/dynamic_vamana.h @@ -58,9 +58,7 @@ class DynamicVamanaInterface : public VamanaInterface { // Check if the index has been consolidated - no `Deleted` entries remain. virtual bool is_consolidated() const = 0; - // Filtered top-k on the search buffer; `filter` takes an external id. Returns false, - // leaving `results` untouched, if the data type does not support it. - // Keep last: new virtuals must be appended. + // Filtered search; false if the data type does not support it. Keep last (ABI). virtual bool filtered_search( svs::QueryResultView results, AnonymousArray<2> queries, @@ -135,7 +133,14 @@ class DynamicVamanaImpl : public VamanaImpl(lib::Type SVS_UNUSED(type)) { const auto view = data::ConstSimpleDataView(queries); - impl().search(results, view, search_parameters, filter, found); + impl().search( + results, + view, + search_parameters, + lib::Returns(lib::Const()), + &filter, + &found + ); }, [&](svs::DataType data_type) { throw ANNEXCEPTION("Unsupported query datatype! Got: {}.", data_type); @@ -647,14 +652,7 @@ class DynamicVamana : public manager::IndexManager { /// @brief Check if the index has been consolidated - no `Deleted` entries remain. bool is_consolidated() const { return impl_->is_consolidated(); } - /// - /// @brief Filtered top-k search on the search buffer, see - /// ``MutableVamanaIndex::search`` (filtered overload). - /// - /// @param filter Takes an external id. - /// @param found Receives the number of results per query. Rows are padded with id - /// ``size_t(-1)`` and distance ``+inf``. - /// @return False if the data type is unsupported, in which case nothing is written. + /// @brief Filtered search (see ``MutableVamanaIndex::search``); false if unsupported. template bool filtered_search( QueryResultView results, diff --git a/tests/svs/index/vamana/dynamic_index.cpp b/tests/svs/index/vamana/dynamic_index.cpp index 7a7fadabe..c903b0470 100644 --- a/tests/svs/index/vamana/dynamic_index.cpp +++ b/tests/svs/index/vamana/dynamic_index.cpp @@ -508,6 +508,19 @@ CATCH_TEST_CASE( auto sp = svs::index::vamana::VamanaSearchParameters{}; sp.buffer_config_ = svs::index::vamana::SearchBufferConfig{50}; + auto search_filtered = [&](auto& results, + const svs::index::vamana::extensions::IDPredicate& pred, + std::vector& found) { + index.search( + results.view(), + queries.cview(), + sp, + svs::lib::Returns(svs::lib::Const()), + &pred, + &found + ); + }; + // Brute force over live ids passing `pred`. std::unordered_set deleted; auto brute_force = [&](size_t q, auto&& pred) { @@ -535,7 +548,7 @@ CATCH_TEST_CASE( auto check = [&](auto&& pred, bool expect_full) { auto results = svs::QueryResult(num_queries, k); std::vector found; - index.search(results.view(), queries.cview(), sp, pred, found); + search_filtered(results, pred, found); CATCH_REQUIRE(found.size() == num_queries); size_t hits = 0, total = 0; for (size_t q = 0; q < num_queries; ++q) { @@ -583,18 +596,51 @@ CATCH_TEST_CASE( check(few, false); auto results = svs::QueryResult(num_queries, k); std::vector found; - index.search(results.view(), queries.cview(), sp, few, found); + search_filtered(results, few, found); for (size_t q = 0; q < num_queries; ++q) { CATCH_REQUIRE(found[q] == 3); } auto none = [](size_t) { return false; }; - index.search(results.view(), queries.cview(), sp, none, found); + search_filtered(results, none, found); for (size_t q = 0; q < num_queries; ++q) { CATCH_REQUIRE(found[q] == 0); } } - CATCH_SECTION("orchestrator overload") { + CATCH_SECTION("unfiltered search is unchanged by the optional arguments") { + index.delete_entries(std::vector{ids[0], ids[1], ids[2]}); + auto plain = svs::QueryResult(num_queries, k); + auto with_defaults = svs::QueryResult(num_queries, k); + auto null_args = svs::QueryResult(num_queries, k); + index.search(plain.view(), queries.cview(), sp); + index.search( + with_defaults.view(), + queries.cview(), + sp, + svs::lib::Returns(svs::lib::Const()), + nullptr, + nullptr + ); + std::vector found; + index.search( + null_args.view(), + queries.cview(), + sp, + svs::lib::Returns(svs::lib::Const()), + nullptr, + &found + ); + CATCH_REQUIRE(found.size() == num_queries); + for (size_t q = 0; q < num_queries; ++q) { + for (size_t j = 0; j < k; ++j) { + CATCH_REQUIRE(plain.index(q, j) == with_defaults.index(q, j)); + CATCH_REQUIRE(plain.index(q, j) == null_args.index(q, j)); + CATCH_REQUIRE(plain.distance(q, j) == with_defaults.distance(q, j)); + } + } + } + + CATCH_SECTION("orchestrator entry point") { auto dyn = svs::DynamicVamana::build( parameters, [&] { From a4983729c5d5055e9271f55a34ebcf5b141bdaea Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Wed, 21 Oct 2026 15:46:49 -0700 Subject: [PATCH 03/33] Run filtered search through the existing batch search path The optional filter on MutableVamanaIndex::search now only chooses which builder the existing greedy search closure uses; both cases call the same per_thread_batch_search. FilteredValidBuilder derives from ValidBuilder. The buffer top-up in single_search applies the filter when the search closure carries one. The per-query result count out-parameter is removed (callers count non-padding ids). per_thread_batch_search now pads result rows with fewer than k entries (id size_t(-1), distance +inf) instead of reading past the end of the search buffer, and the post-search id translation skips the padding. --- include/svs/index/vamana/dynamic_index.h | 106 +++++++++------------ include/svs/index/vamana/extensions.h | 14 ++- include/svs/orchestrators/dynamic_vamana.h | 14 +-- tests/svs/index/vamana/dynamic_index.cpp | 56 ++++++----- 4 files changed, 96 insertions(+), 94 deletions(-) diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 6a0c0cefc..6fbe09c1d 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -111,30 +111,40 @@ class ValidBuilder { }; // Like `ValidBuilder`, but an entry is also invalid if `pred(external_id)` is false. -class FilteredValidBuilder { +class FilteredValidBuilder : public ValidBuilder { public: FilteredValidBuilder( const std::vector& status, const IDTranslator& translator, const extensions::IDPredicate& pred ) - : status_{status} + : ValidBuilder{status} , translator_{translator} , pred_{pred} {} template PredicatedSearchNeighbor operator()(I i, float distance) const { - bool valid = getindex(status_, i) != SlotMetadata::Deleted && - pred_(translator_.get_external(i)); - return PredicatedSearchNeighbor(i, distance, valid); + auto base = ValidBuilder::operator()(i, distance); + if (base.valid() && !pred_(translator_.get_external(i))) { + return PredicatedSearchNeighbor(i, distance, false); + } + return base; } private: - const std::vector& status_; const IDTranslator& translator_; const extensions::IDPredicate& pred_; }; +// A search closure that also carries the filter, so the buffer top-up can apply it too. +template struct FilteredSearch { + Search search; + const extensions::IDPredicate& pred; + template void operator()(Args&&... args) const { + search(std::forward(args)...); + } +}; + // Data types for which the filtered `search` overload is valid (identity accessor). template inline constexpr bool filtered_search_supported = false; template @@ -566,6 +576,9 @@ class MutableVamanaIndex { [&](const auto is, uint64_t /*tid*/) { for (auto i : is) { for (size_t j = 0, jmax = getsize<1>(ids); j < jmax; ++j) { + if (ids.at(i, j) == std::numeric_limits::max()) { + continue; // Padding from a search with fewer than k results. + } auto internal = lib::narrow_cast(ids.at(i, j)); ids.at(i, j) = translate_internal_id(internal); } @@ -669,16 +682,11 @@ class MutableVamanaIndex { const Queries& queries, const search_parameters_type& sp, const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()), - const extensions::IDPredicate* pred = nullptr, - std::vector* found = nullptr + const extensions::IDPredicate* pred = nullptr ) { - // Optional filter: only live entries with `(*pred)(external_id)` are returned; - // others are navigated through but do not count toward the search window. Rows - // with fewer than `n_neighbors` results are padded with id `size_t(-1)` and - // distance `+inf`; `found` (if given) receives the count per query. - if (found != nullptr) { - found->assign(queries.size(), 0); - } + // Optional filter (`pred` takes an external id): only live entries it accepts are + // returned. Rows with fewer than `n_neighbors` results are padded with id + // `size_t(-1)` and distance `+inf`. Only meant for `SimpleData`. threads::parallel_for( threadpool_, threads::StaticPartition{queries.size()}, @@ -697,56 +705,34 @@ class MutableVamanaIndex { auto scratch = extensions::per_thread_batch_search_setup(data_, distance_); - if (pred != nullptr) { - if constexpr (filtered_search_supported) { - auto builder = FilteredValidBuilder{status_, translator_, *pred}; - auto search = - greedy_search_closure(builder, prefetch_parameters, cancel); - auto accessor = data::GetDatumAccessor(); - for (auto i : is) { - auto query = queries.get_datum(i); - search(query, accessor, scratch, buffer); - extensions::check_and_supplement_search_buffer( - *this, buffer, query, pred - ); - size_t n = std::min(num_neighbors, buffer.size()); - for (size_t j = 0; j < num_neighbors; ++j) { - bool ok = j < n; - results.index(i, j) = - ok ? translate_internal_id(buffer[j].id()) : size_t(-1); - results.distance(i, j) = - ok ? buffer[j].distance() - : std::numeric_limits::infinity(); - } - if (found != nullptr) { - (*found)[i] = n; - } - } - return; - } else { - throw ANNEXCEPTION( - "Filtered search is not supported for this data type." - ); - } + auto run = [&](const auto& search) { + extensions::per_thread_batch_search( + data_, + buffer, + scratch, + queries, + results, + threads::UnitRange{is}, + search, + *this, + cancel + ); + }; + if (pred == nullptr) { + run(greedy_search_closure(prefetch_parameters, cancel)); + } else { + auto search = greedy_search_closure( + FilteredValidBuilder{status_, translator_, *pred}, + prefetch_parameters, + cancel + ); + run(FilteredSearch{search, *pred}); } - - extensions::per_thread_batch_search( - data_, - buffer, - scratch, - queries, - results, - threads::UnitRange{is}, - greedy_search_closure(prefetch_parameters, cancel), - *this, - cancel - ); } ); // Check if request to cancel the search - // (Filtered results already hold external ids.) - if (cancel() || pred != nullptr) { + if (cancel()) { return; } diff --git a/include/svs/index/vamana/extensions.h b/include/svs/index/vamana/extensions.h index edcddec2c..83b47557f 100644 --- a/include/svs/index/vamana/extensions.h +++ b/include/svs/index/vamana/extensions.h @@ -29,6 +29,7 @@ // stdlib #include +#include namespace svs::index::vamana::extensions { @@ -495,7 +496,11 @@ SVS_FORCE_INLINE void svs_invoke( search(query, accessor, distance, search_buffer); if constexpr (Index::needs_id_translation) { - check_and_supplement_search_buffer(index, search_buffer, query); + const IDPredicate* pred = nullptr; + if constexpr (requires { search.pred; }) { + pred = &search.pred; // Filtered search: `search` carries its filter. + } + check_and_supplement_search_buffer(index, search_buffer, query, pred); } } @@ -600,7 +605,12 @@ void svs_invoke( // Copy back results. for (size_t j = 0; j < num_neighbors; ++j) { - result.set(search_buffer[j], i, j); + if (j < search_buffer.size()) { + result.set(search_buffer[j], i, j); + } else { // Fewer than k results: pad. + result.index(i, j) = static_cast(-1); + result.distance(i, j) = std::numeric_limits::infinity(); + } } } } diff --git a/include/svs/orchestrators/dynamic_vamana.h b/include/svs/orchestrators/dynamic_vamana.h index 403080185..0af8f6486 100644 --- a/include/svs/orchestrators/dynamic_vamana.h +++ b/include/svs/orchestrators/dynamic_vamana.h @@ -63,8 +63,7 @@ class DynamicVamanaInterface : public VamanaInterface { svs::QueryResultView results, AnonymousArray<2> queries, const index::vamana::VamanaSearchParameters& search_parameters, - const std::function& filter, - std::vector& found + const std::function& filter ) = 0; }; @@ -124,8 +123,7 @@ class DynamicVamanaImpl : public VamanaImpl results, AnonymousArray<2> queries, const index::vamana::VamanaSearchParameters& search_parameters, - const std::function& filter, - std::vector& found + const std::function& filter ) override { if constexpr (requires { requires Impl::has_filtered_search; }) { lib::match( @@ -138,8 +136,7 @@ class DynamicVamanaImpl : public VamanaImpl()), - &filter, - &found + &filter ); }, [&](svs::DataType data_type) { @@ -658,11 +655,10 @@ class DynamicVamana : public manager::IndexManager { QueryResultView results, data::ConstSimpleDataView queries, const VamanaSearchParameters& search_parameters, - const std::function& filter, - std::vector& found + const std::function& filter ) { return impl_->filtered_search( - results, AnonymousArray<2>(queries), search_parameters, filter, found + results, AnonymousArray<2>(queries), search_parameters, filter ); } }; diff --git a/tests/svs/index/vamana/dynamic_index.cpp b/tests/svs/index/vamana/dynamic_index.cpp index c903b0470..0e9580f4c 100644 --- a/tests/svs/index/vamana/dynamic_index.cpp +++ b/tests/svs/index/vamana/dynamic_index.cpp @@ -508,6 +508,7 @@ CATCH_TEST_CASE( auto sp = svs::index::vamana::VamanaSearchParameters{}; sp.buffer_config_ = svs::index::vamana::SearchBufferConfig{50}; + // Runs the filtered search; `found` = number of non-padding ids per query. auto search_filtered = [&](auto& results, const svs::index::vamana::extensions::IDPredicate& pred, std::vector& found) { @@ -516,9 +517,14 @@ CATCH_TEST_CASE( queries.cview(), sp, svs::lib::Returns(svs::lib::Const()), - &pred, - &found + &pred ); + found.assign(num_queries, 0); + for (size_t q = 0; q < num_queries; ++q) { + for (size_t j = 0; j < k; ++j) { + found[q] += results.index(q, j) != static_cast(-1); + } + } }; // Brute force over live ids passing `pred`. @@ -607,35 +613,41 @@ CATCH_TEST_CASE( } } - CATCH_SECTION("unfiltered search is unchanged by the optional arguments") { + CATCH_SECTION("unfiltered search is unchanged by the optional argument") { index.delete_entries(std::vector{ids[0], ids[1], ids[2]}); auto plain = svs::QueryResult(num_queries, k); - auto with_defaults = svs::QueryResult(num_queries, k); - auto null_args = svs::QueryResult(num_queries, k); + auto with_null = svs::QueryResult(num_queries, k); index.search(plain.view(), queries.cview(), sp); index.search( - with_defaults.view(), + with_null.view(), queries.cview(), sp, svs::lib::Returns(svs::lib::Const()), - nullptr, nullptr ); - std::vector found; - index.search( - null_args.view(), - queries.cview(), - sp, - svs::lib::Returns(svs::lib::Const()), - nullptr, - &found - ); - CATCH_REQUIRE(found.size() == num_queries); for (size_t q = 0; q < num_queries; ++q) { for (size_t j = 0; j < k; ++j) { - CATCH_REQUIRE(plain.index(q, j) == with_defaults.index(q, j)); - CATCH_REQUIRE(plain.index(q, j) == null_args.index(q, j)); - CATCH_REQUIRE(plain.distance(q, j) == with_defaults.distance(q, j)); + CATCH_REQUIRE(plain.index(q, j) == with_null.index(q, j)); + CATCH_REQUIRE(plain.distance(q, j) == with_null.distance(q, j)); + } + } + } + + CATCH_SECTION("unfiltered search pads rows with fewer than k results") { + std::vector to_delete(ids.begin() + 4, ids.end()); + index.delete_entries(to_delete); + auto results = svs::QueryResult(num_queries, k); + index.search(results.view(), queries.cview(), sp); + for (size_t q = 0; q < num_queries; ++q) { + for (size_t j = 0; j < k; ++j) { + if (j < 4) { + CATCH_REQUIRE(results.index(q, j) < ids[4]); + } else { + CATCH_REQUIRE(results.index(q, j) == static_cast(-1)); + CATCH_REQUIRE( + results.distance(q, j) == std::numeric_limits::infinity() + ); + } } } } @@ -655,12 +667,10 @@ CATCH_TEST_CASE( num_threads ); auto results = svs::QueryResult(num_queries, k); - std::vector found; CATCH_REQUIRE(dyn.filtered_search( - results.view(), queries.cview(), sp, std::function(mod10), found + results.view(), queries.cview(), sp, std::function(mod10) )); for (size_t q = 0; q < num_queries; ++q) { - CATCH_REQUIRE(found[q] == k); for (size_t j = 0; j < k; ++j) { CATCH_REQUIRE(mod10(results.index(q, j))); } From 5cce383a267d76ba1a2a69a7dedf13a7361dd41d Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Wed, 21 Oct 2026 16:01:09 -0700 Subject: [PATCH 04/33] Filter through ValidBuilder in the existing dynamic search Give ValidBuilder an optional predicate and pass it down through internal_search_builder, greedy_search_closure and search() as optional arguments, instead of a separate builder and search path. Filtered-only checks: the supplement fallback skips failing ids, copy-out pads short rows with id -1 and the worst score, translation skips padded ids. --- include/svs/index/vamana/dynamic_index.h | 106 +++++++-------------- include/svs/index/vamana/extensions.h | 36 ++++--- include/svs/orchestrators/dynamic_vamana.h | 14 ++- tests/svs/index/vamana/dynamic_index.cpp | 2 +- 4 files changed, 64 insertions(+), 94 deletions(-) diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 6fbe09c1d..a78b212e7 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -18,7 +18,6 @@ // stdlib #include -#include #include #include @@ -95,69 +94,39 @@ inline constexpr std::string_view name(SlotMetadata metadata) { class ValidBuilder { public: - ValidBuilder(const std::vector& status) - : status_{status} {} + // Optional `pred` (with `translator`): an entry is also invalid if `pred(external_id)` + // is false. Invalid entries are navigated through but never returned. + ValidBuilder( + const std::vector& status, + const IDTranslator* translator = nullptr, + const std::function* pred = nullptr + ) + : status_{status} + , translator_{translator} + , pred_{pred} {} template constexpr PredicatedSearchNeighbor operator()(I i, float distance) const { bool invalid = getindex(status_, i) == SlotMetadata::Deleted; // This neighbor should be skipped if the metadata corresponding to the given index // marks this slot as deleted. + if (!invalid && pred_ != nullptr) { + invalid = !(*pred_)(translator_->get_external(i)); + } return PredicatedSearchNeighbor(i, distance, !invalid); } private: const std::vector& status_; + const IDTranslator* translator_; + const std::function* pred_; }; -// Like `ValidBuilder`, but an entry is also invalid if `pred(external_id)` is false. -class FilteredValidBuilder : public ValidBuilder { - public: - FilteredValidBuilder( - const std::vector& status, - const IDTranslator& translator, - const extensions::IDPredicate& pred - ) - : ValidBuilder{status} - , translator_{translator} - , pred_{pred} {} - - template - PredicatedSearchNeighbor operator()(I i, float distance) const { - auto base = ValidBuilder::operator()(i, distance); - if (base.valid() && !pred_(translator_.get_external(i))) { - return PredicatedSearchNeighbor(i, distance, false); - } - return base; - } - - private: - const IDTranslator& translator_; - const extensions::IDPredicate& pred_; -}; - -// A search closure that also carries the filter, so the buffer top-up can apply it too. -template struct FilteredSearch { - Search search; - const extensions::IDPredicate& pred; - template void operator()(Args&&... args) const { - search(std::forward(args)...); - } -}; - -// Data types for which the filtered `search` overload is valid (identity accessor). -template inline constexpr bool filtered_search_supported = false; -template -inline constexpr bool filtered_search_supported> = true; - template class MutableVamanaIndex { friend class MultiMutableVamanaIndex; public: - // Present only on indexes that support the filtered `search` overload. - static constexpr bool has_filtered_search = filtered_search_supported; - // Traits static constexpr bool supports_insertions = true; static constexpr bool supports_deletions = true; @@ -576,8 +545,8 @@ class MutableVamanaIndex { [&](const auto is, uint64_t /*tid*/) { for (auto i : is) { for (size_t j = 0, jmax = getsize<1>(ids); j < jmax; ++j) { - if (ids.at(i, j) == std::numeric_limits::max()) { - continue; // Padding from a search with fewer than k results. + if (ids.at(i, j) == static_cast(-1)) { + continue; // Padding from a filtered search. } auto internal = lib::narrow_cast(ids.at(i, j)); ids.at(i, j) = translate_internal_id(internal); @@ -620,15 +589,17 @@ class MutableVamanaIndex { // Return a `greedy_search` compatible builder for this index. // This is an internal method, mostly used to help implement the batch iterator. - ValidBuilder internal_search_builder() const { return ValidBuilder{status_}; } + ValidBuilder internal_search_builder(const std::function* pred = nullptr + ) const { + return ValidBuilder{status_, &translator_, pred}; + } - template auto greedy_search_closure( - const Builder& builder, GreedySearchPrefetchParameters prefetch_parameters, - const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()), + const std::function* pred = nullptr ) const { - return [&, builder, prefetch_parameters]( + return [&, prefetch_parameters, pred]( const auto& query, auto& accessor, auto& distance, auto& buffer ) { // Perform the greedy search using the provided resources. @@ -640,7 +611,7 @@ class MutableVamanaIndex { distance, buffer, vamana::EntryPointInitializer{lib::as_const_span(entry_point_)}, - builder, + internal_search_builder(pred), prefetch_parameters, cancel ); @@ -650,15 +621,6 @@ class MutableVamanaIndex { }; } - auto greedy_search_closure( - GreedySearchPrefetchParameters prefetch_parameters, - const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) - ) const { - return greedy_search_closure( - internal_search_builder(), prefetch_parameters, cancel - ); - } - // Single Search template void search( @@ -682,11 +644,11 @@ class MutableVamanaIndex { const Queries& queries, const search_parameters_type& sp, const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()), - const extensions::IDPredicate* pred = nullptr + const std::function* pred = nullptr ) { - // Optional filter (`pred` takes an external id): only live entries it accepts are - // returned. Rows with fewer than `n_neighbors` results are padded with id - // `size_t(-1)` and distance `+inf`. Only meant for `SimpleData`. + // Optional `pred`: only entries whose external id passes are returned. It is called + // from all search threads. Rows with fewer passing entries than requested are + // padded with id `size_t(-1)` and the worst score. threads::parallel_for( threadpool_, threads::StaticPartition{queries.size()}, @@ -718,15 +680,11 @@ class MutableVamanaIndex { cancel ); }; + auto search = greedy_search_closure(prefetch_parameters, cancel, pred); if (pred == nullptr) { - run(greedy_search_closure(prefetch_parameters, cancel)); + run(search); } else { - auto search = greedy_search_closure( - FilteredValidBuilder{status_, translator_, *pred}, - prefetch_parameters, - cancel - ); - run(FilteredSearch{search, *pred}); + run(extensions::FilteredSearch{search, pred}); } } ); diff --git a/include/svs/index/vamana/extensions.h b/include/svs/index/vamana/extensions.h index 83b47557f..883286b27 100644 --- a/include/svs/index/vamana/extensions.h +++ b/include/svs/index/vamana/extensions.h @@ -27,7 +27,6 @@ #include "svs/lib/preprocessor.h" #include "svs/lib/threads.h" -// stdlib #include #include @@ -429,22 +428,17 @@ struct VamanaSingleSearchType { } }; -/// Predicate on external ids used by filtered search. -using IDPredicate = std::function; - /// In rare cases, the search buffer may not be filled with enough results. /// This can occur in dynamic indexes when many vectors have been deleted /// and the graph becomes sparsely connected. It's a corner case and should /// not happen frequently, but when it does, we may need to supplement the buffer /// with additional results. -/// If `pred` is given, entries whose external id it rejects are skipped (checked before -/// the distance is computed). template void check_and_supplement_search_buffer( const Index& index, SearchBuffer& search_buffer, const Query& query, - const IDPredicate* pred = nullptr + const std::function* pred = nullptr ) { if (search_buffer.valid() < search_buffer.target_window() && search_buffer.valid() < index.size()) { @@ -464,6 +458,15 @@ void check_and_supplement_search_buffer( } } +/// A search closure that also carries the filter of a filtered search. +template struct FilteredSearch { + Search search; + const std::function* pred; + template void operator()(Args&&... args) const { + search(std::forward(args)...); + } +}; + /// Customization point object for processing single queries. inline constexpr VamanaSingleSearchType single_search{}; @@ -496,9 +499,9 @@ SVS_FORCE_INLINE void svs_invoke( search(query, accessor, distance, search_buffer); if constexpr (Index::needs_id_translation) { - const IDPredicate* pred = nullptr; + const std::function* pred = nullptr; if constexpr (requires { search.pred; }) { - pred = &search.pred; // Filtered search: `search` carries its filter. + pred = search.pred; // Filtered search: keep filtered-out ids out of the top-up. } check_and_supplement_search_buffer(index, search_buffer, query, pred); } @@ -605,12 +608,17 @@ void svs_invoke( // Copy back results. for (size_t j = 0; j < num_neighbors; ++j) { - if (j < search_buffer.size()) { - result.set(search_buffer[j], i, j); - } else { // Fewer than k results: pad. - result.index(i, j) = static_cast(-1); - result.distance(i, j) = std::numeric_limits::infinity(); + if constexpr (requires { search.pred; }) { + if (j >= search_buffer.size()) { // Filtered: fewer than k passed. + result.index(i, j) = static_cast(-1); + result.distance(i, j) = + distance::comparator(distance)(0.0f, 1.0f) + ? std::numeric_limits::infinity() + : -std::numeric_limits::infinity(); + continue; + } } + result.set(search_buffer[j], i, j); } } } diff --git a/include/svs/orchestrators/dynamic_vamana.h b/include/svs/orchestrators/dynamic_vamana.h index 0af8f6486..da4e01137 100644 --- a/include/svs/orchestrators/dynamic_vamana.h +++ b/include/svs/orchestrators/dynamic_vamana.h @@ -22,7 +22,6 @@ #include "svs/orchestrators/vamana.h" // stdlib -#include #include namespace svs { @@ -67,6 +66,11 @@ class DynamicVamanaInterface : public VamanaInterface { ) = 0; }; +// Filtered search is only enabled for uncompressed data. +template inline constexpr bool is_simple_data = false; +template +inline constexpr bool is_simple_data> = true; + template class DynamicVamanaImpl : public VamanaImpl { public: @@ -125,15 +129,14 @@ class DynamicVamanaImpl : public VamanaImpl& filter ) override { - if constexpr (requires { requires Impl::has_filtered_search; }) { + if constexpr (is_simple_data) { lib::match( QueryTypes{}, queries.type(), [&](lib::Type SVS_UNUSED(type)) { - const auto view = data::ConstSimpleDataView(queries); impl().search( results, - view, + data::ConstSimpleDataView(queries), search_parameters, lib::Returns(lib::Const()), &filter @@ -649,7 +652,8 @@ class DynamicVamana : public manager::IndexManager { /// @brief Check if the index has been consolidated - no `Deleted` entries remain. bool is_consolidated() const { return impl_->is_consolidated(); } - /// @brief Filtered search (see ``MutableVamanaIndex::search``); false if unsupported. + /// @brief Filtered search: only ids for which ``filter(id)`` is true are returned. + /// Returns false if the data type does not support it. template bool filtered_search( QueryResultView results, diff --git a/tests/svs/index/vamana/dynamic_index.cpp b/tests/svs/index/vamana/dynamic_index.cpp index 0e9580f4c..50fe5b282 100644 --- a/tests/svs/index/vamana/dynamic_index.cpp +++ b/tests/svs/index/vamana/dynamic_index.cpp @@ -510,7 +510,7 @@ CATCH_TEST_CASE( // Runs the filtered search; `found` = number of non-padding ids per query. auto search_filtered = [&](auto& results, - const svs::index::vamana::extensions::IDPredicate& pred, + const std::function& pred, std::vector& found) { index.search( results.view(), From c4d7396e54b4689b43e9b4a15dfc417d8d3a1c85 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Wed, 21 Oct 2026 22:12:04 -0700 Subject: [PATCH 05/33] Fix filtered test and orchestrator for the concurrent index The orchestrator reports filtered search as unsupported for indexes whose search has no filter argument (the concurrent index), instead of failing to compile. In the filtered test, drop the unfiltered-padding section (padding is filtered-only) and add an inner-product case checking padding uses -inf. --- include/svs/orchestrators/dynamic_vamana.h | 13 +++++++- tests/svs/index/vamana/dynamic_index.cpp | 35 ++++++++++++++-------- 2 files changed, 34 insertions(+), 14 deletions(-) diff --git a/include/svs/orchestrators/dynamic_vamana.h b/include/svs/orchestrators/dynamic_vamana.h index da4e01137..9619bdc7d 100644 --- a/include/svs/orchestrators/dynamic_vamana.h +++ b/include/svs/orchestrators/dynamic_vamana.h @@ -129,7 +129,18 @@ class DynamicVamanaImpl : public VamanaImpl& filter ) override { - if constexpr (is_simple_data) { + // Indexes whose `search` has no filter argument (e.g. the concurrent index) report + // "unsupported" like other data types. + constexpr bool supported = + is_simple_data && + requires( + Impl & idx, + svs::QueryResultView r, + data::ConstSimpleDataView q, + const index::vamana::VamanaSearchParameters& sp, + const std::function* pred + ) { idx.search(r, q, sp, lib::Returns(lib::Const()), pred); }; + if constexpr (supported) { lib::match( QueryTypes{}, queries.type(), diff --git a/tests/svs/index/vamana/dynamic_index.cpp b/tests/svs/index/vamana/dynamic_index.cpp index 50fe5b282..44ed215d7 100644 --- a/tests/svs/index/vamana/dynamic_index.cpp +++ b/tests/svs/index/vamana/dynamic_index.cpp @@ -633,21 +633,30 @@ CATCH_TEST_CASE( } } - CATCH_SECTION("unfiltered search pads rows with fewer than k results") { - std::vector to_delete(ids.begin() + 4, ids.end()); - index.delete_entries(to_delete); + CATCH_SECTION("inner product pads with the worst score (-inf)") { + auto data_ip = svs::data::BlockedData(n, dim); + for (size_t i = 0; i < n; ++i) { + data_ip.set_datum(i, raw[i]); + } + svs::index::vamana::VamanaBuildParameters parameters_ip{0.95, 32, 64, 64, 32, true}; + auto index_ip = svs::index::vamana::MutableVamanaIndex( + parameters_ip, std::move(data_ip), ids, svs::distance::DistanceIP(), num_threads + ); + std::function few = [](size_t e) { return e == 100 || e == 200; }; auto results = svs::QueryResult(num_queries, k); - index.search(results.view(), queries.cview(), sp); + index_ip.search( + results.view(), + queries.cview(), + sp, + svs::lib::Returns(svs::lib::Const()), + &few + ); for (size_t q = 0; q < num_queries; ++q) { - for (size_t j = 0; j < k; ++j) { - if (j < 4) { - CATCH_REQUIRE(results.index(q, j) < ids[4]); - } else { - CATCH_REQUIRE(results.index(q, j) == static_cast(-1)); - CATCH_REQUIRE( - results.distance(q, j) == std::numeric_limits::infinity() - ); - } + for (size_t j = 2; j < k; ++j) { + CATCH_REQUIRE(results.index(q, j) == static_cast(-1)); + CATCH_REQUIRE( + results.distance(q, j) == -std::numeric_limits::infinity() + ); } } } From 4ca5691bc06cc49fbedbae21c8cb68d551f07294 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Mon, 26 Oct 2026 12:27:51 -0700 Subject: [PATCH 06/33] Filtered search: two heaps inside the search buffer Experiment: with a filter, MutableBuffer keeps a min-heap of entries to visit and a max-heap of the best valid entries during the search, instead of one sorted list. greedy_search is unchanged; cleanup() sorts the results once at the end. --- include/svs/index/vamana/dynamic_index.h | 2 + .../svs/index/vamana/dynamic_search_buffer.h | 84 +++++++++++++++++++ tests/svs/index/vamana/search_buffer.cpp | 40 +++++++++ 3 files changed, 126 insertions(+) diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index a78b212e7..24d0acce3 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -661,6 +661,8 @@ class MutableVamanaIndex { : sp.buffer_config_, distance::comparator(distance_), sp.search_buffer_visited_set_}; + // Filtered: use two heaps in the buffer during the search. + buffer.enable_heap(pred != nullptr); auto prefetch_parameters = GreedySearchPrefetchParameters{ sp.prefetch_lookahead_, sp.prefetch_step_}; diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index 9e8f304e3..94953f03b 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -95,6 +95,15 @@ template > class MutableBuffer { vector_type candidates_{}; // An optional visited filter. std::optional visited_{std::nullopt}; + // Optional (`enable_heap`): two heaps while searching. `to_visit_` is a min-heap of + // entries still to visit (best on top); `candidates_` is a max-heap of the best valid + // entries (worst on top). `cleanup()` sorts `candidates_` once at the end. + // `seen_`: ids already queued in this search (exact duplicate check). + bool use_heap_ = false; + bool heap_done_ = false; + vector_type to_visit_{}; + std::unordered_set seen_{}; + value_type popped_{}; public: MutableBuffer() = default; @@ -197,6 +206,9 @@ template > class MutableBuffer { /// @brief Prepare the buffer for a new search operation. void clear() { candidates_.clear(); + to_visit_.clear(); + seen_.clear(); + heap_done_ = false; best_unvisited_ = 0; roi_end_ = 0; valid_ = 0; @@ -223,6 +235,9 @@ template > class MutableBuffer { size_t capacity() const { return candidates_.capacity(); } + /// @brief Use two heaps during search (search window = capacity). + void enable_heap(bool enable = true) { use_heap_ = enable; } + /// @brief Return the number of valid elements currently contained in the buffer. size_t valid() const { return valid_; } @@ -267,6 +282,12 @@ template > class MutableBuffer { /// N.B.: If `done()` evaluates to `true`, do not try to extract further candidates /// from it using `next()`. bool done() const { + if (heap_mode()) { + // Nothing left to visit, or the best one is worse than the worst result. + return to_visit_.empty() || + (full() && + compare_(candidates_.front().distance(), to_visit_.front().distance())); + } // Until we've reached the target number of valid elements, we have to ignore // the state of `roi_end_`. return best_unvisited_ == ((valid() < target_valid_) ? size() : roi_end_); @@ -282,6 +303,13 @@ template > class MutableBuffer { /// Post-conditions: /// * The returned neighbor will be marked as visited. const_reference next() { + if (heap_mode()) { + std::pop_heap(to_visit_.begin(), to_visit_.end(), to_visit_compare()); + popped_ = to_visit_.back(); + to_visit_.pop_back(); + popped_.set_visited(); + return popped_; + } // Get the best unvisited node reference node = getindex(candidates_, best_unvisited_); node.set_visited(); @@ -302,6 +330,10 @@ template > class MutableBuffer { /// /// Therefore, sequences of `push_back` must always be followed by a call to `sort`. void push_back(value_type neighbor) { + if (heap_mode()) { + heap_insert(neighbor); + return; + } // Follow the contract of not appending more valid elements than the capacity // allows. bool valid = neighbor.valid(); @@ -453,6 +485,9 @@ template > class MutableBuffer { /// Insert the given neighbor into the search buffer. size_t insert(value_type neighbor) { + if (heap_mode()) { + return heap_insert(neighbor); + } if (can_skip(neighbor.distance())) { return size(); } @@ -505,6 +540,9 @@ template > class MutableBuffer { // Sort all stored elements in the buffer. void sort() { + if (heap_mode()) { + return; // Both heaps are kept in order on every insert. + } assert(valid() <= valid_capacity_); // Put the elements in ascending order. @@ -554,6 +592,14 @@ template > class MutableBuffer { // TODO: Switch over to using iterators for the return values to avoid this. void cleanup() { + if (heap_mode()) { + // `candidates_` holds only valid entries: sort once, then continue as a sorted + // list (e.g. for the top-up). + std::sort(begin(), end(), compare_); + to_visit_.clear(); + seen_.clear(); + heap_done_ = true; + } auto new_end = std::remove_if(begin(), end(), [](const auto& x) { return !x.valid(); }); candidates_.resize(new_end - begin()); @@ -619,6 +665,44 @@ template > class MutableBuffer { } private: + bool heap_mode() const { return use_heap_ && !heap_done_; } + + // `to_visit_`: the best entry on top. + auto to_visit_compare() const { + return [this](const value_type& a, const value_type& b) { + return compare_(b.distance(), a.distance()); + }; + } + + // `candidates_`: the worst entry on top. + auto result_compare() const { + return [this](const value_type& a, const value_type& b) { + return compare_(a.distance(), b.distance()); + }; + } + + size_t heap_insert(value_type neighbor) { + // Skip entries worse than the worst result when full, and ids already queued. + if ((full() && compare_(candidates_.front().distance(), neighbor.distance())) || + !seen_.insert(neighbor.id()).second) { + return size(); + } + to_visit_.push_back(neighbor); + std::push_heap(to_visit_.begin(), to_visit_.end(), to_visit_compare()); + if (!neighbor.valid()) { + return size(); + } + // Valid: keep it among the best `valid_capacity_` results (drop the worst). + candidates_.push_back(neighbor); + std::push_heap(begin(), end(), result_compare()); + if (candidates_.size() > valid_capacity_) { + std::pop_heap(begin(), end(), result_compare()); + candidates_.pop_back(); + } + valid_ = candidates_.size(); + return size(); + } + /// Return how many more valid candidates exist than required. /// If the number of valid candidates is *less* than the target, a negative number /// is returned. diff --git a/tests/svs/index/vamana/search_buffer.cpp b/tests/svs/index/vamana/search_buffer.cpp index 8ecb698a5..f5a89b63a 100644 --- a/tests/svs/index/vamana/search_buffer.cpp +++ b/tests/svs/index/vamana/search_buffer.cpp @@ -906,6 +906,46 @@ CATCH_TEST_CASE("MutableBuffer", "[core][search_buffer]") { CATCH_REQUIRE(b.done()); } } + + // With `enable_heap`, entries are visited from a min-heap and the best valid entries + // are kept in a max-heap; `cleanup()` returns them sorted. + CATCH_SECTION("Heap") { + buffer.enable_heap(); + buffer.insert({10, 5, false}); + buffer.insert({0, 100, true}); + buffer.insert({11, 1, false}); + buffer.insert({1, 50, true}); + buffer.insert({11, 1, false}); // duplicate: ignored + CATCH_REQUIRE(buffer.size() == 2); + CATCH_REQUIRE(buffer.valid() == 2); + CATCH_REQUIRE(buffer[0].id() == 0); // worst result on top + + CATCH_REQUIRE(buffer.next().id() == 11); + CATCH_REQUIRE(buffer.next().id() == 10); + buffer.insert({2, 40, true}); + buffer.insert({3, 30, true}); + CATCH_REQUIRE(buffer.full()); + // Full: worse than the worst valid entry (100) is skipped. + buffer.insert({4, 200, true}); + buffer.insert({12, 150, false}); + // A better valid entry pushes out the worst valid one (id 0). + buffer.insert({5, 20, true}); + CATCH_REQUIRE(buffer.size() == 4); + + CATCH_REQUIRE(buffer.next().id() == 5); + CATCH_REQUIRE(buffer.next().id() == 3); + CATCH_REQUIRE(buffer.next().id() == 2); + CATCH_REQUIRE(buffer.next().id() == 1); + // Only id 0 (100) is left in the heap; it is worse than the worst result (50). + CATCH_REQUIRE(buffer.done()); + + buffer.cleanup(); + CATCH_REQUIRE(buffer.size() == 4); + CATCH_REQUIRE(eq(buffer[0], {5, 20, true})); + CATCH_REQUIRE(eq(buffer[1], {3, 30, true})); + CATCH_REQUIRE(eq(buffer[2], {2, 40, true})); + CATCH_REQUIRE(eq(buffer[3], {1, 50, true})); + } } CATCH_TEST_CASE("Fuzzing Mutable", "[core][search_buffer]") { From dd7c723bde316647f33c26b4a0fc4654c9c964f1 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Mon, 26 Oct 2026 23:27:28 -0700 Subject: [PATCH 07/33] Filtered search: skip the filter for entries the buffer would drop ValidBuilder can be given the search buffer. With a filter, it first asks the buffer whether the entry would be dropped anyway (can_skip) and, if so, returns it as invalid without looking up the external id or calling the filter. Results are unchanged. --- include/svs/index/vamana/dynamic_index.h | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 24d0acce3..37f78c957 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -105,8 +105,19 @@ class ValidBuilder { , translator_{translator} , pred_{pred} {} + /// With a predicate: skip the predicate for entries `buffer` would drop anyway. + template void skip_if_dropped_by(const Buffer& buffer) { + skip_buffer_ = &buffer; + skip_ = [](const void* b, float d) { + return static_cast(b)->can_skip(d); + }; + } + template constexpr PredicatedSearchNeighbor operator()(I i, float distance) const { + if (pred_ != nullptr && skip_ != nullptr && skip_(skip_buffer_, distance)) { + return PredicatedSearchNeighbor(i, distance, false); + } bool invalid = getindex(status_, i) == SlotMetadata::Deleted; // This neighbor should be skipped if the metadata corresponding to the given index // marks this slot as deleted. @@ -120,6 +131,8 @@ class ValidBuilder { const std::vector& status_; const IDTranslator* translator_; const std::function* pred_; + const void* skip_buffer_ = nullptr; + bool (*skip_)(const void*, float) = nullptr; }; template @@ -602,6 +615,10 @@ class MutableVamanaIndex { return [&, prefetch_parameters, pred]( const auto& query, auto& accessor, auto& distance, auto& buffer ) { + auto builder = internal_search_builder(pred); + if constexpr (requires { buffer.can_skip(0.0f); }) { + builder.skip_if_dropped_by(buffer); + } // Perform the greedy search using the provided resources. greedy_search( graph_, @@ -611,7 +628,7 @@ class MutableVamanaIndex { distance, buffer, vamana::EntryPointInitializer{lib::as_const_span(entry_point_)}, - internal_search_builder(pred), + builder, prefetch_parameters, cancel ); From 71e93e9871cb4011bfc06d24ab29d0e366545fbc Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Mon, 26 Oct 2026 23:27:43 -0700 Subject: [PATCH 08/33] Make can_skip use the heap's worst entry while searching The filter skip in ValidBuilder asks can_skip; in heap mode the worst kept entry is not back(). --- include/svs/index/vamana/dynamic_search_buffer.h | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index 94953f03b..6c1c78a7c 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -478,6 +478,10 @@ template > class MutableBuffer { } bool can_skip(float distance) const { + if (heap_mode()) { + // While searching, the worst result is the top of the result max-heap. + return full() && compare_(candidates_.front().distance(), distance); + } // If not full - we cannot skip appending this item. // If we are full, then rely on invariant 6 to compare with the last valid element. return full() && compare_(back().distance(), distance); From f43745ac14e3eb077ae61d2151de704c03b5b5ed Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Tue, 27 Oct 2026 10:48:31 -0700 Subject: [PATCH 09/33] Revert "Make can_skip use the heap's worst entry while searching" This reverts commit 71e93e9871cb4011bfc06d24ab29d0e366545fbc. --- include/svs/index/vamana/dynamic_search_buffer.h | 4 ---- 1 file changed, 4 deletions(-) diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index 6c1c78a7c..94953f03b 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -478,10 +478,6 @@ template > class MutableBuffer { } bool can_skip(float distance) const { - if (heap_mode()) { - // While searching, the worst result is the top of the result max-heap. - return full() && compare_(candidates_.front().distance(), distance); - } // If not full - we cannot skip appending this item. // If we are full, then rely on invariant 6 to compare with the last valid element. return full() && compare_(back().distance(), distance); From 572b4ec0553d22efafddd5e71bec6e1828db5f89 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Tue, 27 Oct 2026 10:48:31 -0700 Subject: [PATCH 10/33] Revert "Filtered search: skip the filter for entries the buffer would drop" This reverts commit dd7c723bde316647f33c26b4a0fc4654c9c964f1. --- include/svs/index/vamana/dynamic_index.h | 19 +------------------ 1 file changed, 1 insertion(+), 18 deletions(-) diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 37f78c957..24d0acce3 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -105,19 +105,8 @@ class ValidBuilder { , translator_{translator} , pred_{pred} {} - /// With a predicate: skip the predicate for entries `buffer` would drop anyway. - template void skip_if_dropped_by(const Buffer& buffer) { - skip_buffer_ = &buffer; - skip_ = [](const void* b, float d) { - return static_cast(b)->can_skip(d); - }; - } - template constexpr PredicatedSearchNeighbor operator()(I i, float distance) const { - if (pred_ != nullptr && skip_ != nullptr && skip_(skip_buffer_, distance)) { - return PredicatedSearchNeighbor(i, distance, false); - } bool invalid = getindex(status_, i) == SlotMetadata::Deleted; // This neighbor should be skipped if the metadata corresponding to the given index // marks this slot as deleted. @@ -131,8 +120,6 @@ class ValidBuilder { const std::vector& status_; const IDTranslator* translator_; const std::function* pred_; - const void* skip_buffer_ = nullptr; - bool (*skip_)(const void*, float) = nullptr; }; template @@ -615,10 +602,6 @@ class MutableVamanaIndex { return [&, prefetch_parameters, pred]( const auto& query, auto& accessor, auto& distance, auto& buffer ) { - auto builder = internal_search_builder(pred); - if constexpr (requires { buffer.can_skip(0.0f); }) { - builder.skip_if_dropped_by(buffer); - } // Perform the greedy search using the provided resources. greedy_search( graph_, @@ -628,7 +611,7 @@ class MutableVamanaIndex { distance, buffer, vamana::EntryPointInitializer{lib::as_const_span(entry_point_)}, - builder, + internal_search_builder(pred), prefetch_parameters, cancel ); From 03aabeacd5a4308510148f35d7a7a337a1aec1ac Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Tue, 27 Oct 2026 10:48:41 -0700 Subject: [PATCH 11/33] Reuse existing tests for the heap search buffer The existing [filtered] dynamic index test drives heap mode end to end. --- tests/svs/index/vamana/search_buffer.cpp | 40 ------------------------ 1 file changed, 40 deletions(-) diff --git a/tests/svs/index/vamana/search_buffer.cpp b/tests/svs/index/vamana/search_buffer.cpp index f5a89b63a..8ecb698a5 100644 --- a/tests/svs/index/vamana/search_buffer.cpp +++ b/tests/svs/index/vamana/search_buffer.cpp @@ -906,46 +906,6 @@ CATCH_TEST_CASE("MutableBuffer", "[core][search_buffer]") { CATCH_REQUIRE(b.done()); } } - - // With `enable_heap`, entries are visited from a min-heap and the best valid entries - // are kept in a max-heap; `cleanup()` returns them sorted. - CATCH_SECTION("Heap") { - buffer.enable_heap(); - buffer.insert({10, 5, false}); - buffer.insert({0, 100, true}); - buffer.insert({11, 1, false}); - buffer.insert({1, 50, true}); - buffer.insert({11, 1, false}); // duplicate: ignored - CATCH_REQUIRE(buffer.size() == 2); - CATCH_REQUIRE(buffer.valid() == 2); - CATCH_REQUIRE(buffer[0].id() == 0); // worst result on top - - CATCH_REQUIRE(buffer.next().id() == 11); - CATCH_REQUIRE(buffer.next().id() == 10); - buffer.insert({2, 40, true}); - buffer.insert({3, 30, true}); - CATCH_REQUIRE(buffer.full()); - // Full: worse than the worst valid entry (100) is skipped. - buffer.insert({4, 200, true}); - buffer.insert({12, 150, false}); - // A better valid entry pushes out the worst valid one (id 0). - buffer.insert({5, 20, true}); - CATCH_REQUIRE(buffer.size() == 4); - - CATCH_REQUIRE(buffer.next().id() == 5); - CATCH_REQUIRE(buffer.next().id() == 3); - CATCH_REQUIRE(buffer.next().id() == 2); - CATCH_REQUIRE(buffer.next().id() == 1); - // Only id 0 (100) is left in the heap; it is worse than the worst result (50). - CATCH_REQUIRE(buffer.done()); - - buffer.cleanup(); - CATCH_REQUIRE(buffer.size() == 4); - CATCH_REQUIRE(eq(buffer[0], {5, 20, true})); - CATCH_REQUIRE(eq(buffer[1], {3, 30, true})); - CATCH_REQUIRE(eq(buffer[2], {2, 40, true})); - CATCH_REQUIRE(eq(buffer[3], {1, 50, true})); - } } CATCH_TEST_CASE("Fuzzing Mutable", "[core][search_buffer]") { From 3974527aec59394894cf80bb923d40c81ca18785 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Mon, 26 Oct 2026 23:27:28 -0700 Subject: [PATCH 12/33] Filtered search: skip the filter for entries the buffer would drop ValidBuilder can be given the search buffer. With a filter, it first asks the buffer whether the entry would be dropped anyway (can_skip) and, if so, returns it as invalid without looking up the external id or calling the filter. Results are unchanged. --- include/svs/index/vamana/dynamic_index.h | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 24d0acce3..37f78c957 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -105,8 +105,19 @@ class ValidBuilder { , translator_{translator} , pred_{pred} {} + /// With a predicate: skip the predicate for entries `buffer` would drop anyway. + template void skip_if_dropped_by(const Buffer& buffer) { + skip_buffer_ = &buffer; + skip_ = [](const void* b, float d) { + return static_cast(b)->can_skip(d); + }; + } + template constexpr PredicatedSearchNeighbor operator()(I i, float distance) const { + if (pred_ != nullptr && skip_ != nullptr && skip_(skip_buffer_, distance)) { + return PredicatedSearchNeighbor(i, distance, false); + } bool invalid = getindex(status_, i) == SlotMetadata::Deleted; // This neighbor should be skipped if the metadata corresponding to the given index // marks this slot as deleted. @@ -120,6 +131,8 @@ class ValidBuilder { const std::vector& status_; const IDTranslator* translator_; const std::function* pred_; + const void* skip_buffer_ = nullptr; + bool (*skip_)(const void*, float) = nullptr; }; template @@ -602,6 +615,10 @@ class MutableVamanaIndex { return [&, prefetch_parameters, pred]( const auto& query, auto& accessor, auto& distance, auto& buffer ) { + auto builder = internal_search_builder(pred); + if constexpr (requires { buffer.can_skip(0.0f); }) { + builder.skip_if_dropped_by(buffer); + } // Perform the greedy search using the provided resources. greedy_search( graph_, @@ -611,7 +628,7 @@ class MutableVamanaIndex { distance, buffer, vamana::EntryPointInitializer{lib::as_const_span(entry_point_)}, - internal_search_builder(pred), + builder, prefetch_parameters, cancel ); From 3fd4de0d6cc02b75501197b3d726568396b7b323 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Mon, 26 Oct 2026 23:27:43 -0700 Subject: [PATCH 13/33] Make can_skip use the heap's worst entry while searching The filter skip in ValidBuilder asks can_skip; in heap mode the worst kept entry is not back(). --- include/svs/index/vamana/dynamic_search_buffer.h | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index 94953f03b..6c1c78a7c 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -478,6 +478,10 @@ template > class MutableBuffer { } bool can_skip(float distance) const { + if (heap_mode()) { + // While searching, the worst result is the top of the result max-heap. + return full() && compare_(candidates_.front().distance(), distance); + } // If not full - we cannot skip appending this item. // If we are full, then rely on invariant 6 to compare with the last valid element. return full() && compare_(back().distance(), distance); From 1aa2516b2464c69ff9b5f7c20356122ef2d23a4e Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Tue, 27 Oct 2026 11:07:00 -0700 Subject: [PATCH 14/33] Remove the added filtered-search tests --- tests/svs/index/vamana/dynamic_index.cpp | 222 ----------------------- 1 file changed, 222 deletions(-) diff --git a/tests/svs/index/vamana/dynamic_index.cpp b/tests/svs/index/vamana/dynamic_index.cpp index 44ed215d7..94daddcf1 100644 --- a/tests/svs/index/vamana/dynamic_index.cpp +++ b/tests/svs/index/vamana/dynamic_index.cpp @@ -23,14 +23,11 @@ #include #include #include -#include #include #include -#include #include #include #include -#include // svs #include "svs/core/recall.h" @@ -467,222 +464,3 @@ CATCH_TEST_CASE("MutableVamana Index Memory Usage", "[graph_index][dynamic_index const size_t usage = index.get_memory_breakdown().total(); CATCH_REQUIRE(usage == expected_total_bytes); } - -CATCH_TEST_CASE( - "MutableVamana Index Filtered Search", "[graph_index][dynamic_index][filtered]" -) { - const size_t num_threads = 2; - const size_t n = 3000; - const size_t dim = 32; - const size_t num_queries = 50; - const size_t k = 10; - using Distance = svs::distance::DistanceL2; - - std::mt19937 rng(42); - std::normal_distribution normal; - auto data = svs::data::BlockedData(n, dim); - std::vector> raw(n, std::vector(dim)); - for (size_t i = 0; i < n; ++i) { - for (auto& x : raw[i]) { - x = normal(rng); - } - data.set_datum(i, raw[i]); - } - auto queries = svs::data::SimpleData(num_queries, dim); - for (size_t i = 0; i < num_queries; ++i) { - std::vector q(dim); - for (auto& x : q) { - x = normal(rng); - } - queries.set_datum(i, q); - } - - // External ids differ from internal ids. - std::vector ids(n); - std::iota(ids.begin(), ids.end(), 7); - svs::index::vamana::VamanaBuildParameters parameters{1.2, 32, 64, 64, 32, true}; - auto index = svs::index::vamana::MutableVamanaIndex( - parameters, std::move(data), ids, Distance(), num_threads - ); - - auto sp = svs::index::vamana::VamanaSearchParameters{}; - sp.buffer_config_ = svs::index::vamana::SearchBufferConfig{50}; - - // Runs the filtered search; `found` = number of non-padding ids per query. - auto search_filtered = [&](auto& results, - const std::function& pred, - std::vector& found) { - index.search( - results.view(), - queries.cview(), - sp, - svs::lib::Returns(svs::lib::Const()), - &pred - ); - found.assign(num_queries, 0); - for (size_t q = 0; q < num_queries; ++q) { - for (size_t j = 0; j < k; ++j) { - found[q] += results.index(q, j) != static_cast(-1); - } - } - }; - - // Brute force over live ids passing `pred`. - std::unordered_set deleted; - auto brute_force = [&](size_t q, auto&& pred) { - std::vector> all; - for (size_t i = 0; i < n; ++i) { - size_t e = ids[i]; - if (deleted.contains(e) || !pred(e)) { - continue; - } - float d = 0; - for (size_t j = 0; j < dim; ++j) { - float diff = queries.get_datum(q)[j] - raw[i][j]; - d += diff * diff; - } - all.emplace_back(d, e); - } - std::sort(all.begin(), all.end()); - std::vector out; - for (size_t i = 0; i < std::min(k, all.size()); ++i) { - out.push_back(all[i].second); - } - return out; - }; - - auto check = [&](auto&& pred, bool expect_full) { - auto results = svs::QueryResult(num_queries, k); - std::vector found; - search_filtered(results, pred, found); - CATCH_REQUIRE(found.size() == num_queries); - size_t hits = 0, total = 0; - for (size_t q = 0; q < num_queries; ++q) { - auto gt = brute_force(q, pred); - CATCH_REQUIRE(found[q] <= k); - if (expect_full) { - CATCH_REQUIRE(found[q] == k); - } - for (size_t j = 0; j < found[q]; ++j) { - size_t id = results.index(q, j); - CATCH_REQUIRE(pred(id)); - CATCH_REQUIRE(!deleted.contains(id)); - hits += std::count(gt.begin(), gt.end(), id); - } - for (size_t j = found[q]; j < k; ++j) { - CATCH_REQUIRE(results.index(q, j) == static_cast(-1)); - CATCH_REQUIRE( - results.distance(q, j) == std::numeric_limits::infinity() - ); - } - total += gt.size(); - } - return double(hits) / double(total); - }; - - auto mod10 = [](size_t e) { return e % 10 == 0; }; - - CATCH_SECTION("recall against brute force") { - CATCH_REQUIRE(check(mod10, true) >= 0.95); - } - - CATCH_SECTION("deleted ids are never returned") { - std::vector to_delete; - for (size_t i = 0; i < n; i += 3) { - to_delete.push_back(ids[i]); - } - index.delete_entries(to_delete); - deleted.insert(to_delete.begin(), to_delete.end()); - CATCH_REQUIRE(check(mod10, true) >= 0.9); - CATCH_REQUIRE(check([](size_t) { return true; }, true) >= 0.9); - } - - CATCH_SECTION("fewer than k passing ids are padded") { - auto few = [](size_t e) { return e == 100 || e == 200 || e == 300; }; - check(few, false); - auto results = svs::QueryResult(num_queries, k); - std::vector found; - search_filtered(results, few, found); - for (size_t q = 0; q < num_queries; ++q) { - CATCH_REQUIRE(found[q] == 3); - } - auto none = [](size_t) { return false; }; - search_filtered(results, none, found); - for (size_t q = 0; q < num_queries; ++q) { - CATCH_REQUIRE(found[q] == 0); - } - } - - CATCH_SECTION("unfiltered search is unchanged by the optional argument") { - index.delete_entries(std::vector{ids[0], ids[1], ids[2]}); - auto plain = svs::QueryResult(num_queries, k); - auto with_null = svs::QueryResult(num_queries, k); - index.search(plain.view(), queries.cview(), sp); - index.search( - with_null.view(), - queries.cview(), - sp, - svs::lib::Returns(svs::lib::Const()), - nullptr - ); - for (size_t q = 0; q < num_queries; ++q) { - for (size_t j = 0; j < k; ++j) { - CATCH_REQUIRE(plain.index(q, j) == with_null.index(q, j)); - CATCH_REQUIRE(plain.distance(q, j) == with_null.distance(q, j)); - } - } - } - - CATCH_SECTION("inner product pads with the worst score (-inf)") { - auto data_ip = svs::data::BlockedData(n, dim); - for (size_t i = 0; i < n; ++i) { - data_ip.set_datum(i, raw[i]); - } - svs::index::vamana::VamanaBuildParameters parameters_ip{0.95, 32, 64, 64, 32, true}; - auto index_ip = svs::index::vamana::MutableVamanaIndex( - parameters_ip, std::move(data_ip), ids, svs::distance::DistanceIP(), num_threads - ); - std::function few = [](size_t e) { return e == 100 || e == 200; }; - auto results = svs::QueryResult(num_queries, k); - index_ip.search( - results.view(), - queries.cview(), - sp, - svs::lib::Returns(svs::lib::Const()), - &few - ); - for (size_t q = 0; q < num_queries; ++q) { - for (size_t j = 2; j < k; ++j) { - CATCH_REQUIRE(results.index(q, j) == static_cast(-1)); - CATCH_REQUIRE( - results.distance(q, j) == -std::numeric_limits::infinity() - ); - } - } - } - - CATCH_SECTION("orchestrator entry point") { - auto dyn = svs::DynamicVamana::build( - parameters, - [&] { - auto d = svs::data::SimpleData(n, dim); - for (size_t i = 0; i < n; ++i) { - d.set_datum(i, raw[i]); - } - return d; - }(), - ids, - Distance(), - num_threads - ); - auto results = svs::QueryResult(num_queries, k); - CATCH_REQUIRE(dyn.filtered_search( - results.view(), queries.cview(), sp, std::function(mod10) - )); - for (size_t q = 0; q < num_queries; ++q) { - for (size_t j = 0; j < k; ++j) { - CATCH_REQUIRE(mod10(results.index(q, j))); - } - } - } -} From 0325ca49134ca0d576fc9e62d949e4374234a01f Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Tue, 27 Oct 2026 11:58:05 -0700 Subject: [PATCH 15/33] Check the filter inside the search buffer, not in ValidBuilder In heap mode, MutableBuffer now calls the filter itself: a neighbor goes to the results heap only if it could improve the results and passes the filter. ValidBuilder is back to the main behaviour (deleted check only). --- include/svs/index/vamana/dynamic_index.h | 56 +++++-------------- .../svs/index/vamana/dynamic_search_buffer.h | 24 ++++---- 2 files changed, 27 insertions(+), 53 deletions(-) diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 37f78c957..f35d57ff8 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -94,45 +94,19 @@ inline constexpr std::string_view name(SlotMetadata metadata) { class ValidBuilder { public: - // Optional `pred` (with `translator`): an entry is also invalid if `pred(external_id)` - // is false. Invalid entries are navigated through but never returned. - ValidBuilder( - const std::vector& status, - const IDTranslator* translator = nullptr, - const std::function* pred = nullptr - ) - : status_{status} - , translator_{translator} - , pred_{pred} {} - - /// With a predicate: skip the predicate for entries `buffer` would drop anyway. - template void skip_if_dropped_by(const Buffer& buffer) { - skip_buffer_ = &buffer; - skip_ = [](const void* b, float d) { - return static_cast(b)->can_skip(d); - }; - } + ValidBuilder(const std::vector& status) + : status_{status} {} template constexpr PredicatedSearchNeighbor operator()(I i, float distance) const { - if (pred_ != nullptr && skip_ != nullptr && skip_(skip_buffer_, distance)) { - return PredicatedSearchNeighbor(i, distance, false); - } bool invalid = getindex(status_, i) == SlotMetadata::Deleted; // This neighbor should be skipped if the metadata corresponding to the given index // marks this slot as deleted. - if (!invalid && pred_ != nullptr) { - invalid = !(*pred_)(translator_->get_external(i)); - } return PredicatedSearchNeighbor(i, distance, !invalid); } private: const std::vector& status_; - const IDTranslator* translator_; - const std::function* pred_; - const void* skip_buffer_ = nullptr; - bool (*skip_)(const void*, float) = nullptr; }; template @@ -602,23 +576,15 @@ class MutableVamanaIndex { // Return a `greedy_search` compatible builder for this index. // This is an internal method, mostly used to help implement the batch iterator. - ValidBuilder internal_search_builder(const std::function* pred = nullptr - ) const { - return ValidBuilder{status_, &translator_, pred}; - } + ValidBuilder internal_search_builder() const { return ValidBuilder{status_}; } auto greedy_search_closure( GreedySearchPrefetchParameters prefetch_parameters, - const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()), - const std::function* pred = nullptr + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) ) const { - return [&, prefetch_parameters, pred]( + return [&, prefetch_parameters]( const auto& query, auto& accessor, auto& distance, auto& buffer ) { - auto builder = internal_search_builder(pred); - if constexpr (requires { buffer.can_skip(0.0f); }) { - builder.skip_if_dropped_by(buffer); - } // Perform the greedy search using the provided resources. greedy_search( graph_, @@ -628,7 +594,7 @@ class MutableVamanaIndex { distance, buffer, vamana::EntryPointInitializer{lib::as_const_span(entry_point_)}, - builder, + internal_search_builder(), prefetch_parameters, cancel ); @@ -678,8 +644,12 @@ class MutableVamanaIndex { : sp.buffer_config_, distance::comparator(distance_), sp.search_buffer_visited_set_}; - // Filtered: use two heaps in the buffer during the search. - buffer.enable_heap(pred != nullptr); + // Filtered: the buffer checks the filter itself (two heaps). + std::function passes; + if (pred != nullptr) { + passes = [&](Idx i) { return (*pred)(translator_.get_external(i)); }; + buffer.enable_heap(&passes); + } auto prefetch_parameters = GreedySearchPrefetchParameters{ sp.prefetch_lookahead_, sp.prefetch_step_}; @@ -699,7 +669,7 @@ class MutableVamanaIndex { cancel ); }; - auto search = greedy_search_closure(prefetch_parameters, cancel, pred); + auto search = greedy_search_closure(prefetch_parameters, cancel); if (pred == nullptr) { run(search); } else { diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index 6c1c78a7c..b38b47366 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -96,10 +96,13 @@ template > class MutableBuffer { // An optional visited filter. std::optional visited_{std::nullopt}; // Optional (`enable_heap`): two heaps while searching. `to_visit_` is a min-heap of - // entries still to visit (best on top); `candidates_` is a max-heap of the best valid - // entries (worst on top). `cleanup()` sorts `candidates_` once at the end. + // entries still to visit (best on top). During heap mode `candidates_` holds the + // results: a max-heap of the best valid entries that pass the filter (worst on top). + // `cleanup()` sorts `candidates_` once at the end. // `seen_`: ids already queued in this search (exact duplicate check). + // `passes_`: the filter, called with an internal id. bool use_heap_ = false; + const std::function* passes_ = nullptr; bool heap_done_ = false; vector_type to_visit_{}; std::unordered_set seen_{}; @@ -235,8 +238,12 @@ template > class MutableBuffer { size_t capacity() const { return candidates_.capacity(); } - /// @brief Use two heaps during search (search window = capacity). - void enable_heap(bool enable = true) { use_heap_ = enable; } + /// @brief Use two heaps during search (search window = capacity). A neighbor enters the + /// results only if `*passes` accepts its internal id. + void enable_heap(const std::function* passes) { + use_heap_ = passes != nullptr; + passes_ = passes; + } /// @brief Return the number of valid elements currently contained in the buffer. size_t valid() const { return valid_; } @@ -478,10 +485,6 @@ template > class MutableBuffer { } bool can_skip(float distance) const { - if (heap_mode()) { - // While searching, the worst result is the top of the result max-heap. - return full() && compare_(candidates_.front().distance(), distance); - } // If not full - we cannot skip appending this item. // If we are full, then rely on invariant 6 to compare with the last valid element. return full() && compare_(back().distance(), distance); @@ -693,10 +696,11 @@ template > class MutableBuffer { } to_visit_.push_back(neighbor); std::push_heap(to_visit_.begin(), to_visit_.end(), to_visit_compare()); - if (!neighbor.valid()) { + // Only valid entries that pass the filter become results; the rest are only visited. + if (!neighbor.valid() || !(*passes_)(neighbor.id())) { return size(); } - // Valid: keep it among the best `valid_capacity_` results (drop the worst). + // Keep it among the best `valid_capacity_` results (drop the worst). candidates_.push_back(neighbor); std::push_heap(begin(), end(), result_compare()); if (candidates_.size() > valid_capacity_) { From b138a596cb41a507d48a268df36da3db84c70396 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Tue, 27 Oct 2026 13:47:57 -0700 Subject: [PATCH 16/33] C API: use the filtered search path for the dynamic index --- bindings/c/src/filtered_search.hpp | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/bindings/c/src/filtered_search.hpp b/bindings/c/src/filtered_search.hpp index 8e45ac691..76956ea25 100644 --- a/bindings/c/src/filtered_search.hpp +++ b/bindings/c/src/filtered_search.hpp @@ -23,6 +23,7 @@ #include #include #include +#include #include #include @@ -31,6 +32,7 @@ #include #include #include +#include #include namespace svs::c_runtime { @@ -168,6 +170,29 @@ std::vector filtered_topk_search( assert(results.n_queries() == queries.size()); + // The dynamic index filters inside the search itself when its data type supports it. + // The search window is max(k, initial_batch_hint), the same as the C API's unfiltered + // search (hint = max(k, window)); other parameters come from the index. + if constexpr (std::is_same_v) { + auto params = index.get_search_parameters(); + params.buffer_config_ = svs::index::vamana::SearchBufferConfig{ + std::max(results.n_neighbors(), initial_batch_hint)}; + if (index.filtered_search( + results.view(), + queries, + params, + [id_filter](size_t id) { return id_filter->is_member(id); } + )) { + std::vector found(queries.size(), 0); + for (size_t i = 0; i < queries.size(); ++i) { + for (size_t j = 0; j < results.n_neighbors(); ++j) { + found[i] += (results.index(i, j) != static_cast(-1)); + } + } + return found; + } + } + // Filtered search: we need to estimate the batch size based on the filter rate and // the number of hits const auto num_neighbors = results.n_neighbors(); From f89a1acc0dcc79a5d22dc33ae713c9575f1047bf Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Wed, 28 Oct 2026 16:21:39 -0700 Subject: [PATCH 17/33] Allow filtered search for data types that opt in --- include/svs/index/vamana/extensions.h | 17 +++++++++++++++++ include/svs/orchestrators/dynamic_vamana.h | 7 +------ 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/include/svs/index/vamana/extensions.h b/include/svs/index/vamana/extensions.h index 883286b27..c63e453d7 100644 --- a/include/svs/index/vamana/extensions.h +++ b/include/svs/index/vamana/extensions.h @@ -639,6 +639,23 @@ template struct UsesReranking { template inline constexpr UsesReranking calibration_uses_reranking{}; +/// Whether `DynamicVamana::filtered_search` is supported for `Data`. +/// Data types opt in by providing an `svs_invoke` overload. +template struct SupportsFilteredSearch { + constexpr bool operator()() const { + if constexpr (svs::svs_invocable) { + return svs::svs_invoke(*this); + } else { + return false; + } + } +}; + +template +constexpr bool svs_invoke(SupportsFilteredSearch>) { + return true; +} + ///// ///// Reconstruct Vector ///// diff --git a/include/svs/orchestrators/dynamic_vamana.h b/include/svs/orchestrators/dynamic_vamana.h index 9619bdc7d..4e78ac340 100644 --- a/include/svs/orchestrators/dynamic_vamana.h +++ b/include/svs/orchestrators/dynamic_vamana.h @@ -66,11 +66,6 @@ class DynamicVamanaInterface : public VamanaInterface { ) = 0; }; -// Filtered search is only enabled for uncompressed data. -template inline constexpr bool is_simple_data = false; -template -inline constexpr bool is_simple_data> = true; - template class DynamicVamanaImpl : public VamanaImpl { public: @@ -132,7 +127,7 @@ class DynamicVamanaImpl : public VamanaImpl && + index::vamana::extensions::SupportsFilteredSearch{}() && requires( Impl & idx, svs::QueryResultView r, From c06a25cb90b52edb6b7036ca28b0595d509ae1d7 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Thu, 29 Oct 2026 14:30:03 -0700 Subject: [PATCH 18/33] Apply clang-format --- include/svs/index/vamana/extensions.h | 7 +++---- include/svs/orchestrators/dynamic_vamana.h | 3 ++- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/include/svs/index/vamana/extensions.h b/include/svs/index/vamana/extensions.h index c63e453d7..2bf8c38e2 100644 --- a/include/svs/index/vamana/extensions.h +++ b/include/svs/index/vamana/extensions.h @@ -611,10 +611,9 @@ void svs_invoke( if constexpr (requires { search.pred; }) { if (j >= search_buffer.size()) { // Filtered: fewer than k passed. result.index(i, j) = static_cast(-1); - result.distance(i, j) = - distance::comparator(distance)(0.0f, 1.0f) - ? std::numeric_limits::infinity() - : -std::numeric_limits::infinity(); + result.distance(i, j) = distance::comparator(distance)(0.0f, 1.0f) + ? std::numeric_limits::infinity() + : -std::numeric_limits::infinity(); continue; } } diff --git a/include/svs/orchestrators/dynamic_vamana.h b/include/svs/orchestrators/dynamic_vamana.h index 4e78ac340..24f3c7f8d 100644 --- a/include/svs/orchestrators/dynamic_vamana.h +++ b/include/svs/orchestrators/dynamic_vamana.h @@ -127,7 +127,8 @@ class DynamicVamanaImpl : public VamanaImpl{}() && + index::vamana::extensions::SupportsFilteredSearch{}( + ) && requires( Impl & idx, svs::QueryResultView r, From e1762b7026504700b6e0d9810e5618e65b1fa954 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Fri, 30 Oct 2026 11:20:15 -0700 Subject: [PATCH 19/33] Enable filtered search for all data types --- include/svs/index/vamana/extensions.h | 17 ----------------- include/svs/orchestrators/dynamic_vamana.h | 6 ++---- 2 files changed, 2 insertions(+), 21 deletions(-) diff --git a/include/svs/index/vamana/extensions.h b/include/svs/index/vamana/extensions.h index 2bf8c38e2..74c347685 100644 --- a/include/svs/index/vamana/extensions.h +++ b/include/svs/index/vamana/extensions.h @@ -638,23 +638,6 @@ template struct UsesReranking { template inline constexpr UsesReranking calibration_uses_reranking{}; -/// Whether `DynamicVamana::filtered_search` is supported for `Data`. -/// Data types opt in by providing an `svs_invoke` overload. -template struct SupportsFilteredSearch { - constexpr bool operator()() const { - if constexpr (svs::svs_invocable) { - return svs::svs_invoke(*this); - } else { - return false; - } - } -}; - -template -constexpr bool svs_invoke(SupportsFilteredSearch>) { - return true; -} - ///// ///// Reconstruct Vector ///// diff --git a/include/svs/orchestrators/dynamic_vamana.h b/include/svs/orchestrators/dynamic_vamana.h index 24f3c7f8d..7d5ef3e36 100644 --- a/include/svs/orchestrators/dynamic_vamana.h +++ b/include/svs/orchestrators/dynamic_vamana.h @@ -124,11 +124,9 @@ class DynamicVamanaImpl : public VamanaImpl& filter ) override { - // Indexes whose `search` has no filter argument (e.g. the concurrent index) report - // "unsupported" like other data types. + // Filtered search is available for every data type. Indexes whose `search` has no + // filter argument (e.g. the concurrent index) report "unsupported". constexpr bool supported = - index::vamana::extensions::SupportsFilteredSearch{}( - ) && requires( Impl & idx, svs::QueryResultView r, From 2820f5bbcc9853b223f20b65da7223b78a23d3bc Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Fri, 30 Oct 2026 11:22:24 -0700 Subject: [PATCH 20/33] Apply clang-format --- include/svs/orchestrators/dynamic_vamana.h | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/include/svs/orchestrators/dynamic_vamana.h b/include/svs/orchestrators/dynamic_vamana.h index 7d5ef3e36..bc01a2828 100644 --- a/include/svs/orchestrators/dynamic_vamana.h +++ b/include/svs/orchestrators/dynamic_vamana.h @@ -126,14 +126,13 @@ class DynamicVamanaImpl : public VamanaImpl r, - data::ConstSimpleDataView q, - const index::vamana::VamanaSearchParameters& sp, - const std::function* pred - ) { idx.search(r, q, sp, lib::Returns(lib::Const()), pred); }; + constexpr bool supported = requires( + Impl & idx, + svs::QueryResultView r, + data::ConstSimpleDataView q, + const index::vamana::VamanaSearchParameters& sp, + const std::function* pred + ) { idx.search(r, q, sp, lib::Returns(lib::Const()), pred); }; if constexpr (supported) { lib::match( QueryTypes{}, From e194701a25c17f31af67416d333185443c1edcdb Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Fri, 30 Oct 2026 13:35:06 -0700 Subject: [PATCH 21/33] Name the filter argument `filter` everywhere --- include/svs/index/vamana/dynamic_index.h | 16 ++++++++-------- include/svs/index/vamana/dynamic_search_buffer.h | 14 +++++++------- include/svs/index/vamana/extensions.h | 16 ++++++++-------- include/svs/orchestrators/dynamic_vamana.h | 4 ++-- 4 files changed, 25 insertions(+), 25 deletions(-) diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index f35d57ff8..6eb97451a 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -627,9 +627,9 @@ class MutableVamanaIndex { const Queries& queries, const search_parameters_type& sp, const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()), - const std::function* pred = nullptr + const std::function* filter = nullptr ) { - // Optional `pred`: only entries whose external id passes are returned. It is called + // Optional `filter`: only entries whose external id passes are returned. It is called // from all search threads. Rows with fewer passing entries than requested are // padded with id `size_t(-1)` and the worst score. threads::parallel_for( @@ -645,10 +645,10 @@ class MutableVamanaIndex { distance::comparator(distance_), sp.search_buffer_visited_set_}; // Filtered: the buffer checks the filter itself (two heaps). - std::function passes; - if (pred != nullptr) { - passes = [&](Idx i) { return (*pred)(translator_.get_external(i)); }; - buffer.enable_heap(&passes); + std::function internal_filter; + if (filter != nullptr) { + internal_filter = [&](Idx i) { return (*filter)(translator_.get_external(i)); }; + buffer.enable_heap(&internal_filter); } auto prefetch_parameters = GreedySearchPrefetchParameters{ @@ -670,10 +670,10 @@ class MutableVamanaIndex { ); }; auto search = greedy_search_closure(prefetch_parameters, cancel); - if (pred == nullptr) { + if (filter == nullptr) { run(search); } else { - run(extensions::FilteredSearch{search, pred}); + run(extensions::FilteredSearch{search, filter}); } } ); diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index b38b47366..c5757b788 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -100,9 +100,9 @@ template > class MutableBuffer { // results: a max-heap of the best valid entries that pass the filter (worst on top). // `cleanup()` sorts `candidates_` once at the end. // `seen_`: ids already queued in this search (exact duplicate check). - // `passes_`: the filter, called with an internal id. + // `internal_filter_`: the filter, called with an internal id. bool use_heap_ = false; - const std::function* passes_ = nullptr; + const std::function* internal_filter_ = nullptr; bool heap_done_ = false; vector_type to_visit_{}; std::unordered_set seen_{}; @@ -239,10 +239,10 @@ template > class MutableBuffer { size_t capacity() const { return candidates_.capacity(); } /// @brief Use two heaps during search (search window = capacity). A neighbor enters the - /// results only if `*passes` accepts its internal id. - void enable_heap(const std::function* passes) { - use_heap_ = passes != nullptr; - passes_ = passes; + /// results only if `*internal_filter` accepts its internal id. + void enable_heap(const std::function* internal_filter) { + use_heap_ = internal_filter != nullptr; + internal_filter_ = internal_filter; } /// @brief Return the number of valid elements currently contained in the buffer. @@ -697,7 +697,7 @@ template > class MutableBuffer { to_visit_.push_back(neighbor); std::push_heap(to_visit_.begin(), to_visit_.end(), to_visit_compare()); // Only valid entries that pass the filter become results; the rest are only visited. - if (!neighbor.valid() || !(*passes_)(neighbor.id())) { + if (!neighbor.valid() || !(*internal_filter_)(neighbor.id())) { return size(); } // Keep it among the best `valid_capacity_` results (drop the worst). diff --git a/include/svs/index/vamana/extensions.h b/include/svs/index/vamana/extensions.h index 74c347685..75e352809 100644 --- a/include/svs/index/vamana/extensions.h +++ b/include/svs/index/vamana/extensions.h @@ -438,13 +438,13 @@ void check_and_supplement_search_buffer( const Index& index, SearchBuffer& search_buffer, const Query& query, - const std::function* pred = nullptr + const std::function* filter = nullptr ) { if (search_buffer.valid() < search_buffer.target_window() && search_buffer.valid() < index.size()) { search_buffer.sort(); for (auto external_id : index.external_ids()) { - if (pred != nullptr && !(*pred)(external_id)) { + if (filter != nullptr && !(*filter)(external_id)) { continue; } auto internal_id = index.translate_external_id(external_id); @@ -461,7 +461,7 @@ void check_and_supplement_search_buffer( /// A search closure that also carries the filter of a filtered search. template struct FilteredSearch { Search search; - const std::function* pred; + const std::function* filter; template void operator()(Args&&... args) const { search(std::forward(args)...); } @@ -499,11 +499,11 @@ SVS_FORCE_INLINE void svs_invoke( search(query, accessor, distance, search_buffer); if constexpr (Index::needs_id_translation) { - const std::function* pred = nullptr; - if constexpr (requires { search.pred; }) { - pred = search.pred; // Filtered search: keep filtered-out ids out of the top-up. + const std::function* filter = nullptr; + if constexpr (requires { search.filter; }) { + filter = search.filter; // Filtered search: keep filtered-out ids out of the top-up. } - check_and_supplement_search_buffer(index, search_buffer, query, pred); + check_and_supplement_search_buffer(index, search_buffer, query, filter); } } @@ -608,7 +608,7 @@ void svs_invoke( // Copy back results. for (size_t j = 0; j < num_neighbors; ++j) { - if constexpr (requires { search.pred; }) { + if constexpr (requires { search.filter; }) { if (j >= search_buffer.size()) { // Filtered: fewer than k passed. result.index(i, j) = static_cast(-1); result.distance(i, j) = distance::comparator(distance)(0.0f, 1.0f) diff --git a/include/svs/orchestrators/dynamic_vamana.h b/include/svs/orchestrators/dynamic_vamana.h index bc01a2828..1fb30f041 100644 --- a/include/svs/orchestrators/dynamic_vamana.h +++ b/include/svs/orchestrators/dynamic_vamana.h @@ -131,8 +131,8 @@ class DynamicVamanaImpl : public VamanaImpl r, data::ConstSimpleDataView q, const index::vamana::VamanaSearchParameters& sp, - const std::function* pred - ) { idx.search(r, q, sp, lib::Returns(lib::Const()), pred); }; + const std::function* filter + ) { idx.search(r, q, sp, lib::Returns(lib::Const()), filter); }; if constexpr (supported) { lib::match( QueryTypes{}, From b9391ed98e9e8d6c88b50cc35b49d8d623a18c38 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Fri, 30 Oct 2026 13:35:15 -0700 Subject: [PATCH 22/33] Apply clang-format --- include/svs/index/vamana/dynamic_index.h | 10 ++++++---- include/svs/index/vamana/dynamic_search_buffer.h | 3 ++- include/svs/index/vamana/extensions.h | 3 ++- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 6eb97451a..8da0aa318 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -629,9 +629,9 @@ class MutableVamanaIndex { const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()), const std::function* filter = nullptr ) { - // Optional `filter`: only entries whose external id passes are returned. It is called - // from all search threads. Rows with fewer passing entries than requested are - // padded with id `size_t(-1)` and the worst score. + // Optional `filter`: only entries whose external id passes are returned. It is + // called from all search threads. Rows with fewer passing entries than requested + // are padded with id `size_t(-1)` and the worst score. threads::parallel_for( threadpool_, threads::StaticPartition{queries.size()}, @@ -647,7 +647,9 @@ class MutableVamanaIndex { // Filtered: the buffer checks the filter itself (two heaps). std::function internal_filter; if (filter != nullptr) { - internal_filter = [&](Idx i) { return (*filter)(translator_.get_external(i)); }; + internal_filter = [&](Idx i) { + return (*filter)(translator_.get_external(i)); + }; buffer.enable_heap(&internal_filter); } diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index c5757b788..8d001c257 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -696,7 +696,8 @@ template > class MutableBuffer { } to_visit_.push_back(neighbor); std::push_heap(to_visit_.begin(), to_visit_.end(), to_visit_compare()); - // Only valid entries that pass the filter become results; the rest are only visited. + // Only valid entries that pass the filter become results; the rest are only + // visited. if (!neighbor.valid() || !(*internal_filter_)(neighbor.id())) { return size(); } diff --git a/include/svs/index/vamana/extensions.h b/include/svs/index/vamana/extensions.h index 75e352809..9ec594486 100644 --- a/include/svs/index/vamana/extensions.h +++ b/include/svs/index/vamana/extensions.h @@ -501,7 +501,8 @@ SVS_FORCE_INLINE void svs_invoke( if constexpr (Index::needs_id_translation) { const std::function* filter = nullptr; if constexpr (requires { search.filter; }) { - filter = search.filter; // Filtered search: keep filtered-out ids out of the top-up. + filter = + search.filter; // Filtered search: keep filtered-out ids out of the top-up. } check_and_supplement_search_buffer(index, search_buffer, query, filter); } From 3dbbd81e66a72f5783ea9cd7f91cd71bb0def358 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Fri, 30 Oct 2026 23:26:00 -0700 Subject: [PATCH 23/33] Two heaps: keep up to the buffer capacity for the re-rank --- .../svs/index/vamana/dynamic_search_buffer.h | 75 ++++++++++++++----- 1 file changed, 56 insertions(+), 19 deletions(-) diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index 8d001c257..e0bfa627e 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -97,14 +97,16 @@ template > class MutableBuffer { std::optional visited_{std::nullopt}; // Optional (`enable_heap`): two heaps while searching. `to_visit_` is a min-heap of // entries still to visit (best on top). During heap mode `candidates_` holds the - // results: a max-heap of the best valid entries that pass the filter (worst on top). - // `cleanup()` sorts `candidates_` once at the end. + // results: a max-heap of the best `target_valid_` valid entries that pass the filter + // (worst on top). `extras_` is a max-heap of the next best passing entries, up to + // `valid_capacity_` in total. `cleanup()` merges both and sorts once at the end. // `seen_`: ids already queued in this search (exact duplicate check). // `internal_filter_`: the filter, called with an internal id. bool use_heap_ = false; const std::function* internal_filter_ = nullptr; bool heap_done_ = false; vector_type to_visit_{}; + vector_type extras_{}; std::unordered_set seen_{}; value_type popped_{}; @@ -210,6 +212,7 @@ template > class MutableBuffer { void clear() { candidates_.clear(); to_visit_.clear(); + extras_.clear(); seen_.clear(); heap_done_ = false; best_unvisited_ = 0; @@ -238,11 +241,16 @@ template > class MutableBuffer { size_t capacity() const { return candidates_.capacity(); } - /// @brief Use two heaps during search (search window = capacity). A neighbor enters the - /// results only if `*internal_filter` accepts its internal id. + /// @brief Use two heaps during search. The results heap holds the best + /// `target_window()` entries and bounds the search; up to `target_capacity()` entries + /// are kept in total. A neighbor enters the results only if `*internal_filter` accepts + /// its internal id. void enable_heap(const std::function* internal_filter) { use_heap_ = internal_filter != nullptr; internal_filter_ = internal_filter; + if (use_heap_) { + extras_.reserve(valid_capacity_ - target_valid_); + } } /// @brief Return the number of valid elements currently contained in the buffer. @@ -292,7 +300,7 @@ template > class MutableBuffer { if (heap_mode()) { // Nothing left to visit, or the best one is worse than the worst result. return to_visit_.empty() || - (full() && + (candidates_.size() == target_valid_ && compare_(candidates_.front().distance(), to_visit_.front().distance())); } // Until we've reached the target number of valid elements, we have to ignore @@ -600,9 +608,11 @@ template > class MutableBuffer { // TODO: Switch over to using iterators for the return values to avoid this. void cleanup() { if (heap_mode()) { - // `candidates_` holds only valid entries: sort once, then continue as a sorted - // list (e.g. for the top-up). + // `candidates_` and `extras_` hold only valid entries: merge, sort once, then + // continue as a sorted list (e.g. for the re-rank and the top-up). + candidates_.insert(candidates_.end(), extras_.begin(), extras_.end()); std::sort(begin(), end(), compare_); + extras_.clear(); to_visit_.clear(); seen_.clear(); heap_done_ = true; @@ -688,27 +698,54 @@ template > class MutableBuffer { }; } + // Keep `neighbor` among the extras (the passing entries after the best `target_valid_`) + // unless they already hold `valid_capacity_ - target_valid_` better ones. + void heap_keep_extra(const value_type& neighbor) { + extras_.push_back(neighbor); + std::push_heap(extras_.begin(), extras_.end(), result_compare()); + if (extras_.size() > valid_capacity_ - target_valid_) { + std::pop_heap(extras_.begin(), extras_.end(), result_compare()); + extras_.pop_back(); + } + } + size_t heap_insert(value_type neighbor) { - // Skip entries worse than the worst result when full, and ids already queued. - if ((full() && compare_(candidates_.front().distance(), neighbor.distance())) || - !seen_.insert(neighbor.id()).second) { + // Entries worse than the `target_valid_`th result are not visited, but may still + // be kept among the extras. + bool beyond_window = candidates_.size() == target_valid_ && + compare_(candidates_.front().distance(), neighbor.distance()); + bool extras_wanted = + extras_.size() < valid_capacity_ - target_valid_ || + (!extras_.empty() && compare_(neighbor.distance(), extras_.front().distance())); + // Skip entries that are not wanted, and ids already queued. + if ((beyond_window && !extras_wanted) || !seen_.insert(neighbor.id()).second) { return size(); } - to_visit_.push_back(neighbor); - std::push_heap(to_visit_.begin(), to_visit_.end(), to_visit_compare()); + if (!beyond_window) { + to_visit_.push_back(neighbor); + std::push_heap(to_visit_.begin(), to_visit_.end(), to_visit_compare()); + } // Only valid entries that pass the filter become results; the rest are only // visited. if (!neighbor.valid() || !(*internal_filter_)(neighbor.id())) { return size(); } - // Keep it among the best `valid_capacity_` results (drop the worst). - candidates_.push_back(neighbor); - std::push_heap(begin(), end(), result_compare()); - if (candidates_.size() > valid_capacity_) { - std::pop_heap(begin(), end(), result_compare()); - candidates_.pop_back(); + if (beyond_window) { + heap_keep_extra(neighbor); + } else { + // Keep it among the best `target_valid_` results (the worst becomes an extra). + candidates_.push_back(neighbor); + std::push_heap(begin(), end(), result_compare()); + if (candidates_.size() > target_valid_) { + std::pop_heap(begin(), end(), result_compare()); + value_type worst = candidates_.back(); + candidates_.pop_back(); + if (valid_capacity_ > target_valid_) { + heap_keep_extra(worst); + } + } } - valid_ = candidates_.size(); + valid_ = candidates_.size() + extras_.size(); return size(); } From 203af29ec6b3e44135df3086baa700f9a4cad1b3 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Fri, 30 Oct 2026 23:26:00 -0700 Subject: [PATCH 24/33] Keep the unfiltered search path inlined --- include/svs/index/vamana/dynamic_index.h | 2 +- .../svs/index/vamana/dynamic_search_buffer.h | 23 +++++++++++-------- include/svs/index/vamana/greedy_search.h | 4 ++-- 3 files changed, 16 insertions(+), 13 deletions(-) diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 8da0aa318..4c637e6e0 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -582,7 +582,7 @@ class MutableVamanaIndex { GreedySearchPrefetchParameters prefetch_parameters, const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) ) const { - return [&, prefetch_parameters]( + return [&, prefetch_parameters] [[gnu::always_inline]] ( const auto& query, auto& accessor, auto& distance, auto& buffer ) { // Perform the greedy search using the provided resources. diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index e0bfa627e..273cb7c76 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -104,7 +104,9 @@ template > class MutableBuffer { // `internal_filter_`: the filter, called with an internal id. bool use_heap_ = false; const std::function* internal_filter_ = nullptr; - bool heap_done_ = false; + // `heap_mode_` is `use_heap_` until `cleanup()`; it is a single flag so the sorted-list + // path pays one load per check. + bool heap_mode_ = false; vector_type to_visit_{}; vector_type extras_{}; std::unordered_set seen_{}; @@ -214,7 +216,7 @@ template > class MutableBuffer { to_visit_.clear(); extras_.clear(); seen_.clear(); - heap_done_ = false; + heap_mode_ = use_heap_; best_unvisited_ = 0; roi_end_ = 0; valid_ = 0; @@ -247,6 +249,7 @@ template > class MutableBuffer { /// its internal id. void enable_heap(const std::function* internal_filter) { use_heap_ = internal_filter != nullptr; + heap_mode_ = use_heap_; internal_filter_ = internal_filter; if (use_heap_) { extras_.reserve(valid_capacity_ - target_valid_); @@ -297,7 +300,7 @@ template > class MutableBuffer { /// N.B.: If `done()` evaluates to `true`, do not try to extract further candidates /// from it using `next()`. bool done() const { - if (heap_mode()) { + if (heap_mode()) [[unlikely]] { // Nothing left to visit, or the best one is worse than the worst result. return to_visit_.empty() || (candidates_.size() == target_valid_ && @@ -318,7 +321,7 @@ template > class MutableBuffer { /// Post-conditions: /// * The returned neighbor will be marked as visited. const_reference next() { - if (heap_mode()) { + if (heap_mode()) [[unlikely]] { std::pop_heap(to_visit_.begin(), to_visit_.end(), to_visit_compare()); popped_ = to_visit_.back(); to_visit_.pop_back(); @@ -345,7 +348,7 @@ template > class MutableBuffer { /// /// Therefore, sequences of `push_back` must always be followed by a call to `sort`. void push_back(value_type neighbor) { - if (heap_mode()) { + if (heap_mode()) [[unlikely]] { heap_insert(neighbor); return; } @@ -500,7 +503,7 @@ template > class MutableBuffer { /// Insert the given neighbor into the search buffer. size_t insert(value_type neighbor) { - if (heap_mode()) { + if (heap_mode()) [[unlikely]] { return heap_insert(neighbor); } if (can_skip(neighbor.distance())) { @@ -555,7 +558,7 @@ template > class MutableBuffer { // Sort all stored elements in the buffer. void sort() { - if (heap_mode()) { + if (heap_mode()) [[unlikely]] { return; // Both heaps are kept in order on every insert. } assert(valid() <= valid_capacity_); @@ -607,7 +610,7 @@ template > class MutableBuffer { // TODO: Switch over to using iterators for the return values to avoid this. void cleanup() { - if (heap_mode()) { + if (heap_mode()) [[unlikely]] { // `candidates_` and `extras_` hold only valid entries: merge, sort once, then // continue as a sorted list (e.g. for the re-rank and the top-up). candidates_.insert(candidates_.end(), extras_.begin(), extras_.end()); @@ -615,7 +618,7 @@ template > class MutableBuffer { extras_.clear(); to_visit_.clear(); seen_.clear(); - heap_done_ = true; + heap_mode_ = false; } auto new_end = std::remove_if(begin(), end(), [](const auto& x) { return !x.valid(); }); @@ -682,7 +685,7 @@ template > class MutableBuffer { } private: - bool heap_mode() const { return use_heap_ && !heap_done_; } + bool heap_mode() const { return heap_mode_; } // `to_visit_`: the best entry on top. auto to_visit_compare() const { diff --git a/include/svs/index/vamana/greedy_search.h b/include/svs/index/vamana/greedy_search.h index f12c0129f..714315eb2 100644 --- a/include/svs/index/vamana/greedy_search.h +++ b/include/svs/index/vamana/greedy_search.h @@ -121,7 +121,7 @@ template < typename Initializer, typename Builder, GreedySearchTracker Tracker> -void greedy_search( +SVS_FORCE_INLINE void greedy_search( const Graph& graph, const Dataset& dataset, Accessor& accessor, @@ -213,7 +213,7 @@ template < typename Buffer, typename Initializer, typename Builder = NeighborBuilder> -void greedy_search( +SVS_FORCE_INLINE void greedy_search( const Graph& graph, const Dataset& dataset, Accessor& accessor, From f5f5f5004301fb562e5a933bc903e987ef039bcb Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Mon, 2 Nov 2026 11:14:12 -0800 Subject: [PATCH 25/33] Revert "Keep the unfiltered search path inlined" This reverts commit 203af29ec6b3e44135df3086baa700f9a4cad1b3. --- include/svs/index/vamana/dynamic_index.h | 2 +- .../svs/index/vamana/dynamic_search_buffer.h | 23 ++++++++----------- include/svs/index/vamana/greedy_search.h | 4 ++-- 3 files changed, 13 insertions(+), 16 deletions(-) diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 4c637e6e0..8da0aa318 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -582,7 +582,7 @@ class MutableVamanaIndex { GreedySearchPrefetchParameters prefetch_parameters, const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) ) const { - return [&, prefetch_parameters] [[gnu::always_inline]] ( + return [&, prefetch_parameters]( const auto& query, auto& accessor, auto& distance, auto& buffer ) { // Perform the greedy search using the provided resources. diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index 273cb7c76..e0bfa627e 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -104,9 +104,7 @@ template > class MutableBuffer { // `internal_filter_`: the filter, called with an internal id. bool use_heap_ = false; const std::function* internal_filter_ = nullptr; - // `heap_mode_` is `use_heap_` until `cleanup()`; it is a single flag so the sorted-list - // path pays one load per check. - bool heap_mode_ = false; + bool heap_done_ = false; vector_type to_visit_{}; vector_type extras_{}; std::unordered_set seen_{}; @@ -216,7 +214,7 @@ template > class MutableBuffer { to_visit_.clear(); extras_.clear(); seen_.clear(); - heap_mode_ = use_heap_; + heap_done_ = false; best_unvisited_ = 0; roi_end_ = 0; valid_ = 0; @@ -249,7 +247,6 @@ template > class MutableBuffer { /// its internal id. void enable_heap(const std::function* internal_filter) { use_heap_ = internal_filter != nullptr; - heap_mode_ = use_heap_; internal_filter_ = internal_filter; if (use_heap_) { extras_.reserve(valid_capacity_ - target_valid_); @@ -300,7 +297,7 @@ template > class MutableBuffer { /// N.B.: If `done()` evaluates to `true`, do not try to extract further candidates /// from it using `next()`. bool done() const { - if (heap_mode()) [[unlikely]] { + if (heap_mode()) { // Nothing left to visit, or the best one is worse than the worst result. return to_visit_.empty() || (candidates_.size() == target_valid_ && @@ -321,7 +318,7 @@ template > class MutableBuffer { /// Post-conditions: /// * The returned neighbor will be marked as visited. const_reference next() { - if (heap_mode()) [[unlikely]] { + if (heap_mode()) { std::pop_heap(to_visit_.begin(), to_visit_.end(), to_visit_compare()); popped_ = to_visit_.back(); to_visit_.pop_back(); @@ -348,7 +345,7 @@ template > class MutableBuffer { /// /// Therefore, sequences of `push_back` must always be followed by a call to `sort`. void push_back(value_type neighbor) { - if (heap_mode()) [[unlikely]] { + if (heap_mode()) { heap_insert(neighbor); return; } @@ -503,7 +500,7 @@ template > class MutableBuffer { /// Insert the given neighbor into the search buffer. size_t insert(value_type neighbor) { - if (heap_mode()) [[unlikely]] { + if (heap_mode()) { return heap_insert(neighbor); } if (can_skip(neighbor.distance())) { @@ -558,7 +555,7 @@ template > class MutableBuffer { // Sort all stored elements in the buffer. void sort() { - if (heap_mode()) [[unlikely]] { + if (heap_mode()) { return; // Both heaps are kept in order on every insert. } assert(valid() <= valid_capacity_); @@ -610,7 +607,7 @@ template > class MutableBuffer { // TODO: Switch over to using iterators for the return values to avoid this. void cleanup() { - if (heap_mode()) [[unlikely]] { + if (heap_mode()) { // `candidates_` and `extras_` hold only valid entries: merge, sort once, then // continue as a sorted list (e.g. for the re-rank and the top-up). candidates_.insert(candidates_.end(), extras_.begin(), extras_.end()); @@ -618,7 +615,7 @@ template > class MutableBuffer { extras_.clear(); to_visit_.clear(); seen_.clear(); - heap_mode_ = false; + heap_done_ = true; } auto new_end = std::remove_if(begin(), end(), [](const auto& x) { return !x.valid(); }); @@ -685,7 +682,7 @@ template > class MutableBuffer { } private: - bool heap_mode() const { return heap_mode_; } + bool heap_mode() const { return use_heap_ && !heap_done_; } // `to_visit_`: the best entry on top. auto to_visit_compare() const { diff --git a/include/svs/index/vamana/greedy_search.h b/include/svs/index/vamana/greedy_search.h index 714315eb2..f12c0129f 100644 --- a/include/svs/index/vamana/greedy_search.h +++ b/include/svs/index/vamana/greedy_search.h @@ -121,7 +121,7 @@ template < typename Initializer, typename Builder, GreedySearchTracker Tracker> -SVS_FORCE_INLINE void greedy_search( +void greedy_search( const Graph& graph, const Dataset& dataset, Accessor& accessor, @@ -213,7 +213,7 @@ template < typename Buffer, typename Initializer, typename Builder = NeighborBuilder> -SVS_FORCE_INLINE void greedy_search( +void greedy_search( const Graph& graph, const Dataset& dataset, Accessor& accessor, From 34bfa34299ff3bc220365ee9241d78bf3542ab4c Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Mon, 2 Nov 2026 11:16:42 -0800 Subject: [PATCH 26/33] Choose heap mode at compile time --- include/svs/index/vamana/dynamic_index.h | 87 ++++++--- .../svs/index/vamana/dynamic_search_buffer.h | 170 +++++++++++------- 2 files changed, 164 insertions(+), 93 deletions(-) diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 8da0aa318..fcfe5a98e 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -130,6 +130,7 @@ class MutableVamanaIndex { using distance_type = Dist; using search_buffer_type = MutableBuffer>; + using filtered_search_buffer_type = MutableBuffer, true>; using graph_type = Graph; using data_type = Data; @@ -636,6 +637,10 @@ class MutableVamanaIndex { threadpool_, threads::StaticPartition{queries.size()}, [&](const auto is, uint64_t SVS_UNUSED(tid)) { + if (filter != nullptr) { + filtered_search_thread(results, queries, sp, cancel, *filter, is); + return; + } size_t num_neighbors = results.n_neighbors(); auto buffer = search_buffer_type{ // Legalize search buffer for this search. @@ -644,39 +649,23 @@ class MutableVamanaIndex { : sp.buffer_config_, distance::comparator(distance_), sp.search_buffer_visited_set_}; - // Filtered: the buffer checks the filter itself (two heaps). - std::function internal_filter; - if (filter != nullptr) { - internal_filter = [&](Idx i) { - return (*filter)(translator_.get_external(i)); - }; - buffer.enable_heap(&internal_filter); - } auto prefetch_parameters = GreedySearchPrefetchParameters{ sp.prefetch_lookahead_, sp.prefetch_step_}; auto scratch = extensions::per_thread_batch_search_setup(data_, distance_); - auto run = [&](const auto& search) { - extensions::per_thread_batch_search( - data_, - buffer, - scratch, - queries, - results, - threads::UnitRange{is}, - search, - *this, - cancel - ); - }; - auto search = greedy_search_closure(prefetch_parameters, cancel); - if (filter == nullptr) { - run(search); - } else { - run(extensions::FilteredSearch{search, filter}); - } + extensions::per_thread_batch_search( + data_, + buffer, + scratch, + queries, + results, + threads::UnitRange{is}, + greedy_search_closure(prefetch_parameters, cancel), + *this, + cancel + ); } ); @@ -690,6 +679,50 @@ class MutableVamanaIndex { translate_to_external(results.indices()); } + private: + // The part of `search` that handles a filter. Kept apart so that the unfiltered search + // above is compiled with the regular buffer only. + template + SVS_NOINLINE void filtered_search_thread( + QueryResultView& results, + const Queries& queries, + const search_parameters_type& sp, + const lib::DefaultPredicate& cancel, + const std::function& filter, + const Range& is + ) { + size_t num_neighbors = results.n_neighbors(); + auto buffer = filtered_search_buffer_type{ + // Legalize search buffer for this search. + sp.buffer_config_.get_total_capacity() < num_neighbors + ? SearchBufferConfig{num_neighbors} + : sp.buffer_config_, + distance::comparator(distance_), + sp.search_buffer_visited_set_}; + // The buffer checks the filter itself (two heaps). + std::function internal_filter = [&](Idx i) { + return filter(translator_.get_external(i)); + }; + buffer.enable_heap(&internal_filter); + + auto prefetch_parameters = + GreedySearchPrefetchParameters{sp.prefetch_lookahead_, sp.prefetch_step_}; + auto scratch = extensions::per_thread_batch_search_setup(data_, distance_); + auto search = greedy_search_closure(prefetch_parameters, cancel); + extensions::per_thread_batch_search( + data_, + buffer, + scratch, + queries, + results, + threads::UnitRange{is}, + extensions::FilteredSearch{search, &filter}, + *this, + cancel + ); + } + + public: /// /// @brief Return a unique instance of the distance function. /// diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index e0bfa627e..df30cec33 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -33,12 +33,30 @@ namespace svs::index::vamana { +namespace detail { +// State used only by the two-heap mode of `MutableBuffer` (see below). +template struct MutableBufferHeapState { + const std::function* internal_filter_ = nullptr; + bool done_ = false; + Vector to_visit_{}; + Vector extras_{}; + std::unordered_set seen_{}; + PredicatedSearchNeighbor popped_{}; +}; +struct NoMutableBufferHeapState {}; +} // namespace detail + /// /// A search buffer that allows entries to be predicated out. /// The search buffer will still navigate through those entries, but it won't include /// them in its final result. /// -template > class MutableBuffer { +/// With `HeapMode == true` the buffer uses two heaps while searching and takes the filter +/// through `enable_heap()`. With `HeapMode == false` (the default) none of the heap code or +/// state exists. +/// +template , bool HeapMode = false> +class MutableBuffer { public: // Type Aliases using index_type = Idx; @@ -95,20 +113,18 @@ template > class MutableBuffer { vector_type candidates_{}; // An optional visited filter. std::optional visited_{std::nullopt}; - // Optional (`enable_heap`): two heaps while searching. `to_visit_` is a min-heap of - // entries still to visit (best on top). During heap mode `candidates_` holds the - // results: a max-heap of the best `target_valid_` valid entries that pass the filter - // (worst on top). `extras_` is a max-heap of the next best passing entries, up to - // `valid_capacity_` in total. `cleanup()` merges both and sorts once at the end. + // `HeapMode` only (`enable_heap`): two heaps while searching. `to_visit_` is a + // min-heap of entries still to visit (best on top). Until `cleanup()`, `candidates_` + // holds the results: a max-heap of the best `target_valid_` valid entries that pass + // the filter (worst on top). `extras_` is a max-heap of the next best passing entries, + // up to `valid_capacity_` in total. `cleanup()` merges both and sorts once at the end. // `seen_`: ids already queued in this search (exact duplicate check). // `internal_filter_`: the filter, called with an internal id. - bool use_heap_ = false; - const std::function* internal_filter_ = nullptr; - bool heap_done_ = false; - vector_type to_visit_{}; - vector_type extras_{}; - std::unordered_set seen_{}; - value_type popped_{}; + [[no_unique_address]] std::conditional_t< + HeapMode, + detail::MutableBufferHeapState, + detail::NoMutableBufferHeapState> + heap_{}; public: MutableBuffer() = default; @@ -211,10 +227,12 @@ template > class MutableBuffer { /// @brief Prepare the buffer for a new search operation. void clear() { candidates_.clear(); - to_visit_.clear(); - extras_.clear(); - seen_.clear(); - heap_done_ = false; + if constexpr (HeapMode) { + heap_.to_visit_.clear(); + heap_.extras_.clear(); + heap_.seen_.clear(); + heap_.done_ = false; + } best_unvisited_ = 0; roi_end_ = 0; valid_ = 0; @@ -245,12 +263,11 @@ template > class MutableBuffer { /// `target_window()` entries and bounds the search; up to `target_capacity()` entries /// are kept in total. A neighbor enters the results only if `*internal_filter` accepts /// its internal id. - void enable_heap(const std::function* internal_filter) { - use_heap_ = internal_filter != nullptr; - internal_filter_ = internal_filter; - if (use_heap_) { - extras_.reserve(valid_capacity_ - target_valid_); - } + void enable_heap(const std::function* internal_filter) + requires HeapMode + { + heap_.internal_filter_ = internal_filter; + heap_.extras_.reserve(valid_capacity_ - target_valid_); } /// @brief Return the number of valid elements currently contained in the buffer. @@ -297,11 +314,15 @@ template > class MutableBuffer { /// N.B.: If `done()` evaluates to `true`, do not try to extract further candidates /// from it using `next()`. bool done() const { - if (heap_mode()) { - // Nothing left to visit, or the best one is worse than the worst result. - return to_visit_.empty() || - (candidates_.size() == target_valid_ && - compare_(candidates_.front().distance(), to_visit_.front().distance())); + if constexpr (HeapMode) { + if (!heap_.done_) { + // Nothing left to visit, or the best one is worse than the worst result. + return heap_.to_visit_.empty() || (candidates_.size() == target_valid_ && + compare_( + candidates_.front().distance(), + heap_.to_visit_.front().distance() + )); + } } // Until we've reached the target number of valid elements, we have to ignore // the state of `roi_end_`. @@ -318,12 +339,16 @@ template > class MutableBuffer { /// Post-conditions: /// * The returned neighbor will be marked as visited. const_reference next() { - if (heap_mode()) { - std::pop_heap(to_visit_.begin(), to_visit_.end(), to_visit_compare()); - popped_ = to_visit_.back(); - to_visit_.pop_back(); - popped_.set_visited(); - return popped_; + if constexpr (HeapMode) { + if (!heap_.done_) { + std::pop_heap( + heap_.to_visit_.begin(), heap_.to_visit_.end(), to_visit_compare() + ); + heap_.popped_ = heap_.to_visit_.back(); + heap_.to_visit_.pop_back(); + heap_.popped_.set_visited(); + return heap_.popped_; + } } // Get the best unvisited node reference node = getindex(candidates_, best_unvisited_); @@ -345,9 +370,11 @@ template > class MutableBuffer { /// /// Therefore, sequences of `push_back` must always be followed by a call to `sort`. void push_back(value_type neighbor) { - if (heap_mode()) { - heap_insert(neighbor); - return; + if constexpr (HeapMode) { + if (!heap_.done_) { + heap_insert(neighbor); + return; + } } // Follow the contract of not appending more valid elements than the capacity // allows. @@ -500,8 +527,10 @@ template > class MutableBuffer { /// Insert the given neighbor into the search buffer. size_t insert(value_type neighbor) { - if (heap_mode()) { - return heap_insert(neighbor); + if constexpr (HeapMode) { + if (!heap_.done_) { + return heap_insert(neighbor); + } } if (can_skip(neighbor.distance())) { return size(); @@ -555,8 +584,10 @@ template > class MutableBuffer { // Sort all stored elements in the buffer. void sort() { - if (heap_mode()) { - return; // Both heaps are kept in order on every insert. + if constexpr (HeapMode) { + if (!heap_.done_) { + return; // Both heaps are kept in order on every insert. + } } assert(valid() <= valid_capacity_); @@ -607,15 +638,19 @@ template > class MutableBuffer { // TODO: Switch over to using iterators for the return values to avoid this. void cleanup() { - if (heap_mode()) { - // `candidates_` and `extras_` hold only valid entries: merge, sort once, then - // continue as a sorted list (e.g. for the re-rank and the top-up). - candidates_.insert(candidates_.end(), extras_.begin(), extras_.end()); - std::sort(begin(), end(), compare_); - extras_.clear(); - to_visit_.clear(); - seen_.clear(); - heap_done_ = true; + if constexpr (HeapMode) { + if (!heap_.done_) { + // `candidates_` and `extras_` hold only valid entries: merge, sort once, + // then continue as a sorted list (e.g. for the re-rank and the top-up). + candidates_.insert( + candidates_.end(), heap_.extras_.begin(), heap_.extras_.end() + ); + std::sort(begin(), end(), compare_); + heap_.extras_.clear(); + heap_.to_visit_.clear(); + heap_.seen_.clear(); + heap_.done_ = true; + } } auto new_end = std::remove_if(begin(), end(), [](const auto& x) { return !x.valid(); }); @@ -682,8 +717,6 @@ template > class MutableBuffer { } private: - bool heap_mode() const { return use_heap_ && !heap_done_; } - // `to_visit_`: the best entry on top. auto to_visit_compare() const { return [this](const value_type& a, const value_type& b) { @@ -701,11 +734,11 @@ template > class MutableBuffer { // Keep `neighbor` among the extras (the passing entries after the best `target_valid_`) // unless they already hold `valid_capacity_ - target_valid_` better ones. void heap_keep_extra(const value_type& neighbor) { - extras_.push_back(neighbor); - std::push_heap(extras_.begin(), extras_.end(), result_compare()); - if (extras_.size() > valid_capacity_ - target_valid_) { - std::pop_heap(extras_.begin(), extras_.end(), result_compare()); - extras_.pop_back(); + heap_.extras_.push_back(neighbor); + std::push_heap(heap_.extras_.begin(), heap_.extras_.end(), result_compare()); + if (heap_.extras_.size() > valid_capacity_ - target_valid_) { + std::pop_heap(heap_.extras_.begin(), heap_.extras_.end(), result_compare()); + heap_.extras_.pop_back(); } } @@ -715,19 +748,23 @@ template > class MutableBuffer { bool beyond_window = candidates_.size() == target_valid_ && compare_(candidates_.front().distance(), neighbor.distance()); bool extras_wanted = - extras_.size() < valid_capacity_ - target_valid_ || - (!extras_.empty() && compare_(neighbor.distance(), extras_.front().distance())); + heap_.extras_.size() < valid_capacity_ - target_valid_ || + (!heap_.extras_.empty() && + compare_(neighbor.distance(), heap_.extras_.front().distance())); // Skip entries that are not wanted, and ids already queued. - if ((beyond_window && !extras_wanted) || !seen_.insert(neighbor.id()).second) { + if ((beyond_window && !extras_wanted) || + !heap_.seen_.insert(neighbor.id()).second) { return size(); } if (!beyond_window) { - to_visit_.push_back(neighbor); - std::push_heap(to_visit_.begin(), to_visit_.end(), to_visit_compare()); + heap_.to_visit_.push_back(neighbor); + std::push_heap( + heap_.to_visit_.begin(), heap_.to_visit_.end(), to_visit_compare() + ); } // Only valid entries that pass the filter become results; the rest are only // visited. - if (!neighbor.valid() || !(*internal_filter_)(neighbor.id())) { + if (!neighbor.valid() || !(*heap_.internal_filter_)(neighbor.id())) { return size(); } if (beyond_window) { @@ -745,7 +782,7 @@ template > class MutableBuffer { } } } - valid_ = candidates_.size() + extras_.size(); + valid_ = candidates_.size() + heap_.extras_.size(); return size(); } @@ -771,8 +808,9 @@ template > class MutableBuffer { } }; -template -std::ostream& operator<<(std::ostream& io, const MutableBuffer& buffer) { +template +std::ostream& +operator<<(std::ostream& io, const MutableBuffer& buffer) { return io << "MutableBuffer<" << datatype_v << ">(" << "target_valid = " << buffer.target_capacity() << ", best_unvisited = " << buffer.best_unvisited() From a1beeca52569d844215527a8ed0bb2275d5c430d Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Mon, 2 Nov 2026 13:09:27 -0800 Subject: [PATCH 27/33] Apply clang-format --- bindings/c/src/filtered_search.hpp | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/bindings/c/src/filtered_search.hpp b/bindings/c/src/filtered_search.hpp index 76956ea25..50a81608d 100644 --- a/bindings/c/src/filtered_search.hpp +++ b/bindings/c/src/filtered_search.hpp @@ -177,12 +177,9 @@ std::vector filtered_topk_search( auto params = index.get_search_parameters(); params.buffer_config_ = svs::index::vamana::SearchBufferConfig{ std::max(results.n_neighbors(), initial_batch_hint)}; - if (index.filtered_search( - results.view(), - queries, - params, - [id_filter](size_t id) { return id_filter->is_member(id); } - )) { + if (index.filtered_search(results.view(), queries, params, [id_filter](size_t id) { + return id_filter->is_member(id); + })) { std::vector found(queries.size(), 0); for (size_t i = 0; i < queries.size(); ++i) { for (size_t j = 0; j < results.n_neighbors(); ++j) { From 6c2e199fd3dbfe1e357bb850993ba53e6b9f4b86 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Mon, 2 Nov 2026 15:18:21 -0800 Subject: [PATCH 28/33] Keep the filter in the search buffer, chosen at compile time --- include/svs/index/vamana/dynamic_index.h | 118 +++++++----------- .../svs/index/vamana/dynamic_search_buffer.h | 39 +++--- include/svs/index/vamana/extensions.h | 35 ++---- 3 files changed, 76 insertions(+), 116 deletions(-) diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index fcfe5a98e..9f6646e7d 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -18,8 +18,10 @@ // stdlib #include +#include #include #include +#include // Include the flat index to spin-up exhaustive searches on demand. #include "svs/index/flat/flat.h" @@ -130,7 +132,6 @@ class MutableVamanaIndex { using distance_type = Dist; using search_buffer_type = MutableBuffer>; - using filtered_search_buffer_type = MutableBuffer, true>; using graph_type = Graph; using data_type = Data; @@ -630,42 +631,53 @@ class MutableVamanaIndex { const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()), const std::function* filter = nullptr ) { - // Optional `filter`: only entries whose external id passes are returned. It is - // called from all search threads. Rows with fewer passing entries than requested - // are padded with id `size_t(-1)` and the worst score. threads::parallel_for( threadpool_, threads::StaticPartition{queries.size()}, [&](const auto is, uint64_t SVS_UNUSED(tid)) { + auto body = [&](std::bool_constant) { + size_t num_neighbors = results.n_neighbors(); + auto buffer = MutableBuffer, Filtered>{ + // Legalize search buffer for this search. + sp.buffer_config_.get_total_capacity() < num_neighbors + ? SearchBufferConfig{num_neighbors} + : sp.buffer_config_, + distance::comparator(distance_), + sp.search_buffer_visited_set_}; + // Alive for the whole search; only the filtered buffer uses it. + std::conditional_t, std::monostate> + internal_filter{}; + if constexpr (Filtered) { + internal_filter = [&](Idx i) { + return (*filter)(translator_.get_external(i)); + }; + buffer.set_filter(&internal_filter); + } + + auto prefetch_parameters = GreedySearchPrefetchParameters{ + sp.prefetch_lookahead_, sp.prefetch_step_}; + + auto scratch = + extensions::per_thread_batch_search_setup(data_, distance_); + + extensions::per_thread_batch_search( + data_, + buffer, + scratch, + queries, + results, + threads::UnitRange{is}, + greedy_search_closure(prefetch_parameters, cancel), + *this, + cancel + ); + }; if (filter != nullptr) { - filtered_search_thread(results, queries, sp, cancel, *filter, is); - return; + [&]() __attribute__((noinline)) { body(std::true_type{}); } + (); + } else { + body(std::false_type{}); } - size_t num_neighbors = results.n_neighbors(); - auto buffer = search_buffer_type{ - // Legalize search buffer for this search. - sp.buffer_config_.get_total_capacity() < num_neighbors - ? SearchBufferConfig{num_neighbors} - : sp.buffer_config_, - distance::comparator(distance_), - sp.search_buffer_visited_set_}; - - auto prefetch_parameters = GreedySearchPrefetchParameters{ - sp.prefetch_lookahead_, sp.prefetch_step_}; - - auto scratch = extensions::per_thread_batch_search_setup(data_, distance_); - - extensions::per_thread_batch_search( - data_, - buffer, - scratch, - queries, - results, - threads::UnitRange{is}, - greedy_search_closure(prefetch_parameters, cancel), - *this, - cancel - ); } ); @@ -679,50 +691,6 @@ class MutableVamanaIndex { translate_to_external(results.indices()); } - private: - // The part of `search` that handles a filter. Kept apart so that the unfiltered search - // above is compiled with the regular buffer only. - template - SVS_NOINLINE void filtered_search_thread( - QueryResultView& results, - const Queries& queries, - const search_parameters_type& sp, - const lib::DefaultPredicate& cancel, - const std::function& filter, - const Range& is - ) { - size_t num_neighbors = results.n_neighbors(); - auto buffer = filtered_search_buffer_type{ - // Legalize search buffer for this search. - sp.buffer_config_.get_total_capacity() < num_neighbors - ? SearchBufferConfig{num_neighbors} - : sp.buffer_config_, - distance::comparator(distance_), - sp.search_buffer_visited_set_}; - // The buffer checks the filter itself (two heaps). - std::function internal_filter = [&](Idx i) { - return filter(translator_.get_external(i)); - }; - buffer.enable_heap(&internal_filter); - - auto prefetch_parameters = - GreedySearchPrefetchParameters{sp.prefetch_lookahead_, sp.prefetch_step_}; - auto scratch = extensions::per_thread_batch_search_setup(data_, distance_); - auto search = greedy_search_closure(prefetch_parameters, cancel); - extensions::per_thread_batch_search( - data_, - buffer, - scratch, - queries, - results, - threads::UnitRange{is}, - extensions::FilteredSearch{search, &filter}, - *this, - cancel - ); - } - - public: /// /// @brief Return a unique instance of the distance function. /// diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index df30cec33..100b8054d 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -51,11 +51,11 @@ struct NoMutableBufferHeapState {}; /// The search buffer will still navigate through those entries, but it won't include /// them in its final result. /// -/// With `HeapMode == true` the buffer uses two heaps while searching and takes the filter -/// through `enable_heap()`. With `HeapMode == false` (the default) none of the heap code or +/// With `Filtered == true` the buffer takes a filter through `set_filter()` and uses two +/// heaps while searching. With `Filtered == false` (the default) none of the heap code or /// state exists. /// -template , bool HeapMode = false> +template , bool Filtered = false> class MutableBuffer { public: // Type Aliases @@ -113,7 +113,7 @@ class MutableBuffer { vector_type candidates_{}; // An optional visited filter. std::optional visited_{std::nullopt}; - // `HeapMode` only (`enable_heap`): two heaps while searching. `to_visit_` is a + // `Filtered` only (`set_filter`): two heaps while searching. `to_visit_` is a // min-heap of entries still to visit (best on top). Until `cleanup()`, `candidates_` // holds the results: a max-heap of the best `target_valid_` valid entries that pass // the filter (worst on top). `extras_` is a max-heap of the next best passing entries, @@ -121,7 +121,7 @@ class MutableBuffer { // `seen_`: ids already queued in this search (exact duplicate check). // `internal_filter_`: the filter, called with an internal id. [[no_unique_address]] std::conditional_t< - HeapMode, + Filtered, detail::MutableBufferHeapState, detail::NoMutableBufferHeapState> heap_{}; @@ -227,7 +227,7 @@ class MutableBuffer { /// @brief Prepare the buffer for a new search operation. void clear() { candidates_.clear(); - if constexpr (HeapMode) { + if constexpr (Filtered) { heap_.to_visit_.clear(); heap_.extras_.clear(); heap_.seen_.clear(); @@ -263,13 +263,20 @@ class MutableBuffer { /// `target_window()` entries and bounds the search; up to `target_capacity()` entries /// are kept in total. A neighbor enters the results only if `*internal_filter` accepts /// its internal id. - void enable_heap(const std::function* internal_filter) - requires HeapMode + void set_filter(const std::function* internal_filter) + requires Filtered { heap_.internal_filter_ = internal_filter; heap_.extras_.reserve(valid_capacity_ - target_valid_); } + /// @brief The filter given to `set_filter`. + const std::function* filter() const + requires Filtered + { + return heap_.internal_filter_; + } + /// @brief Return the number of valid elements currently contained in the buffer. size_t valid() const { return valid_; } @@ -314,7 +321,7 @@ class MutableBuffer { /// N.B.: If `done()` evaluates to `true`, do not try to extract further candidates /// from it using `next()`. bool done() const { - if constexpr (HeapMode) { + if constexpr (Filtered) { if (!heap_.done_) { // Nothing left to visit, or the best one is worse than the worst result. return heap_.to_visit_.empty() || (candidates_.size() == target_valid_ && @@ -339,7 +346,7 @@ class MutableBuffer { /// Post-conditions: /// * The returned neighbor will be marked as visited. const_reference next() { - if constexpr (HeapMode) { + if constexpr (Filtered) { if (!heap_.done_) { std::pop_heap( heap_.to_visit_.begin(), heap_.to_visit_.end(), to_visit_compare() @@ -370,7 +377,7 @@ class MutableBuffer { /// /// Therefore, sequences of `push_back` must always be followed by a call to `sort`. void push_back(value_type neighbor) { - if constexpr (HeapMode) { + if constexpr (Filtered) { if (!heap_.done_) { heap_insert(neighbor); return; @@ -527,7 +534,7 @@ class MutableBuffer { /// Insert the given neighbor into the search buffer. size_t insert(value_type neighbor) { - if constexpr (HeapMode) { + if constexpr (Filtered) { if (!heap_.done_) { return heap_insert(neighbor); } @@ -584,7 +591,7 @@ class MutableBuffer { // Sort all stored elements in the buffer. void sort() { - if constexpr (HeapMode) { + if constexpr (Filtered) { if (!heap_.done_) { return; // Both heaps are kept in order on every insert. } @@ -638,7 +645,7 @@ class MutableBuffer { // TODO: Switch over to using iterators for the return values to avoid this. void cleanup() { - if constexpr (HeapMode) { + if constexpr (Filtered) { if (!heap_.done_) { // `candidates_` and `extras_` hold only valid entries: merge, sort once, // then continue as a sorted list (e.g. for the re-rank and the top-up). @@ -808,9 +815,9 @@ class MutableBuffer { } }; -template +template std::ostream& -operator<<(std::ostream& io, const MutableBuffer& buffer) { +operator<<(std::ostream& io, const MutableBuffer& buffer) { return io << "MutableBuffer<" << datatype_v << ">(" << "target_valid = " << buffer.target_capacity() << ", best_unvisited = " << buffer.best_unvisited() diff --git a/include/svs/index/vamana/extensions.h b/include/svs/index/vamana/extensions.h index 9ec594486..1d6b69053 100644 --- a/include/svs/index/vamana/extensions.h +++ b/include/svs/index/vamana/extensions.h @@ -27,7 +27,6 @@ #include "svs/lib/preprocessor.h" #include "svs/lib/threads.h" -#include #include namespace svs::index::vamana::extensions { @@ -435,19 +434,18 @@ struct VamanaSingleSearchType { /// with additional results. template void check_and_supplement_search_buffer( - const Index& index, - SearchBuffer& search_buffer, - const Query& query, - const std::function* filter = nullptr + const Index& index, SearchBuffer& search_buffer, const Query& query ) { if (search_buffer.valid() < search_buffer.target_window() && search_buffer.valid() < index.size()) { search_buffer.sort(); for (auto external_id : index.external_ids()) { - if (filter != nullptr && !(*filter)(external_id)) { - continue; - } auto internal_id = index.translate_external_id(external_id); + if constexpr (requires { search_buffer.filter(); }) { + if (!(*search_buffer.filter())(internal_id)) { + continue; // Filtered search: keep filtered-out ids out of the top-up. + } + } auto dist = index.get_distance(external_id, query); auto builder = index.internal_search_builder(); search_buffer.insert(builder(internal_id, dist)); @@ -458,15 +456,6 @@ void check_and_supplement_search_buffer( } } -/// A search closure that also carries the filter of a filtered search. -template struct FilteredSearch { - Search search; - const std::function* filter; - template void operator()(Args&&... args) const { - search(std::forward(args)...); - } -}; - /// Customization point object for processing single queries. inline constexpr VamanaSingleSearchType single_search{}; @@ -499,12 +488,7 @@ SVS_FORCE_INLINE void svs_invoke( search(query, accessor, distance, search_buffer); if constexpr (Index::needs_id_translation) { - const std::function* filter = nullptr; - if constexpr (requires { search.filter; }) { - filter = - search.filter; // Filtered search: keep filtered-out ids out of the top-up. - } - check_and_supplement_search_buffer(index, search_buffer, query, filter); + check_and_supplement_search_buffer(index, search_buffer, query); } } @@ -609,8 +593,9 @@ void svs_invoke( // Copy back results. for (size_t j = 0; j < num_neighbors; ++j) { - if constexpr (requires { search.filter; }) { - if (j >= search_buffer.size()) { // Filtered: fewer than k passed. + if constexpr (requires { search_buffer.filter(); }) { + if (j >= search_buffer.size()) { + // Filtered search with fewer than k passing entries: pad. result.index(i, j) = static_cast(-1); result.distance(i, j) = distance::comparator(distance)(0.0f, 1.0f) ? std::numeric_limits::infinity() From 2add52f93fca61a22ff2c3300266ab8b9dfed4e4 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Mon, 2 Nov 2026 16:25:33 -0800 Subject: [PATCH 29/33] Create the search buffer in one helper --- include/svs/index/vamana/dynamic_index.h | 40 +++++++++++-------- .../svs/index/vamana/dynamic_search_buffer.h | 18 +++++---- 2 files changed, 34 insertions(+), 24 deletions(-) diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 9f6646e7d..cf08f0164 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -21,7 +21,6 @@ #include #include #include -#include // Include the flat index to spin-up exhaustive searches on demand. #include "svs/index/flat/flat.h" @@ -623,6 +622,28 @@ class MutableVamanaIndex { ); } + private: + // Build the search buffer for one thread of a batch search. + template + auto make_search_buffer( + const search_parameters_type& sp, + size_t num_neighbors, + const std::function* filter + ) const { + MutableBuffer, Filtered> buffer{ + // Legalize search buffer for this search. + sp.buffer_config_.get_total_capacity() < num_neighbors + ? SearchBufferConfig{num_neighbors} + : sp.buffer_config_, + distance::comparator(distance_), + sp.search_buffer_visited_set_}; + if constexpr (Filtered) { + buffer.set_filter(*filter, translator_); + } + return buffer; + } + + public: template void search( QueryResultView results, @@ -637,22 +658,7 @@ class MutableVamanaIndex { [&](const auto is, uint64_t SVS_UNUSED(tid)) { auto body = [&](std::bool_constant) { size_t num_neighbors = results.n_neighbors(); - auto buffer = MutableBuffer, Filtered>{ - // Legalize search buffer for this search. - sp.buffer_config_.get_total_capacity() < num_neighbors - ? SearchBufferConfig{num_neighbors} - : sp.buffer_config_, - distance::comparator(distance_), - sp.search_buffer_visited_set_}; - // Alive for the whole search; only the filtered buffer uses it. - std::conditional_t, std::monostate> - internal_filter{}; - if constexpr (Filtered) { - internal_filter = [&](Idx i) { - return (*filter)(translator_.get_external(i)); - }; - buffer.set_filter(&internal_filter); - } + auto buffer = make_search_buffer(sp, num_neighbors, filter); auto prefetch_parameters = GreedySearchPrefetchParameters{ sp.prefetch_lookahead_, sp.prefetch_step_}; diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index 100b8054d..d410924e2 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -36,7 +36,7 @@ namespace svs::index::vamana { namespace detail { // State used only by the two-heap mode of `MutableBuffer` (see below). template struct MutableBufferHeapState { - const std::function* internal_filter_ = nullptr; + std::function internal_filter_{}; bool done_ = false; Vector to_visit_{}; Vector extras_{}; @@ -261,12 +261,16 @@ class MutableBuffer { /// @brief Use two heaps during search. The results heap holds the best /// `target_window()` entries and bounds the search; up to `target_capacity()` entries - /// are kept in total. A neighbor enters the results only if `*internal_filter` accepts - /// its internal id. - void set_filter(const std::function* internal_filter) + /// are kept in total. A neighbor enters the results only if `filter` accepts its + /// external id (`translator` maps internal to external ids). Both must outlive the + /// use of this buffer. + template + void set_filter(const std::function& filter, const Translator& translator) requires Filtered { - heap_.internal_filter_ = internal_filter; + heap_.internal_filter_ = [f = &filter, t = &translator](Idx i) { + return (*f)(t->get_external(i)); + }; heap_.extras_.reserve(valid_capacity_ - target_valid_); } @@ -274,7 +278,7 @@ class MutableBuffer { const std::function* filter() const requires Filtered { - return heap_.internal_filter_; + return &heap_.internal_filter_; } /// @brief Return the number of valid elements currently contained in the buffer. @@ -771,7 +775,7 @@ class MutableBuffer { } // Only valid entries that pass the filter become results; the rest are only // visited. - if (!neighbor.valid() || !(*heap_.internal_filter_)(neighbor.id())) { + if (!neighbor.valid() || !heap_.internal_filter_(neighbor.id())) { return size(); } if (beyond_window) { From 0296fec200fb17140737ef779394a3bf5c5da450 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Mon, 2 Nov 2026 16:58:36 -0800 Subject: [PATCH 30/33] Keep the extra candidates in the result vector --- .../svs/index/vamana/dynamic_search_buffer.h | 85 ++++++++++--------- 1 file changed, 45 insertions(+), 40 deletions(-) diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index d410924e2..00deb17a5 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -39,7 +39,6 @@ template struct MutableBufferHeapState { std::function internal_filter_{}; bool done_ = false; Vector to_visit_{}; - Vector extras_{}; std::unordered_set seen_{}; PredicatedSearchNeighbor popped_{}; }; @@ -116,8 +115,10 @@ class MutableBuffer { // `Filtered` only (`set_filter`): two heaps while searching. `to_visit_` is a // min-heap of entries still to visit (best on top). Until `cleanup()`, `candidates_` // holds the results: a max-heap of the best `target_valid_` valid entries that pass - // the filter (worst on top). `extras_` is a max-heap of the next best passing entries, - // up to `valid_capacity_` in total. `cleanup()` merges both and sorts once at the end. + // the filter (worst on top) in `[0, target_valid_)`. Passing entries beyond that are + // appended unsorted after it, trimmed to the best `valid_capacity_ - target_valid_` + // when the tail gets too long. `cleanup()` sorts once and keeps the best + // `valid_capacity_`. // `seen_`: ids already queued in this search (exact duplicate check). // `internal_filter_`: the filter, called with an internal id. [[no_unique_address]] std::conditional_t< @@ -229,7 +230,6 @@ class MutableBuffer { candidates_.clear(); if constexpr (Filtered) { heap_.to_visit_.clear(); - heap_.extras_.clear(); heap_.seen_.clear(); heap_.done_ = false; } @@ -271,7 +271,6 @@ class MutableBuffer { heap_.internal_filter_ = [f = &filter, t = &translator](Idx i) { return (*f)(t->get_external(i)); }; - heap_.extras_.reserve(valid_capacity_ - target_valid_); } /// @brief The filter given to `set_filter`. @@ -328,7 +327,7 @@ class MutableBuffer { if constexpr (Filtered) { if (!heap_.done_) { // Nothing left to visit, or the best one is worse than the worst result. - return heap_.to_visit_.empty() || (candidates_.size() == target_valid_ && + return heap_.to_visit_.empty() || (candidates_.size() >= target_valid_ && compare_( candidates_.front().distance(), heap_.to_visit_.front().distance() @@ -651,13 +650,14 @@ class MutableBuffer { void cleanup() { if constexpr (Filtered) { if (!heap_.done_) { - // `candidates_` and `extras_` hold only valid entries: merge, sort once, - // then continue as a sorted list (e.g. for the re-rank and the top-up). - candidates_.insert( - candidates_.end(), heap_.extras_.begin(), heap_.extras_.end() - ); + // `candidates_` holds only valid entries: sort once and keep the best + // `valid_capacity_`, then continue as a sorted list (e.g. for the re-rank + // and the top-up). std::sort(begin(), end(), compare_); - heap_.extras_.clear(); + if (candidates_.size() > valid_capacity_) { + candidates_.resize(valid_capacity_); + } + valid_ = candidates_.size(); heap_.to_visit_.clear(); heap_.seen_.clear(); heap_.done_ = true; @@ -742,29 +742,35 @@ class MutableBuffer { }; } - // Keep `neighbor` among the extras (the passing entries after the best `target_valid_`) - // unless they already hold `valid_capacity_ - target_valid_` better ones. + // Append `neighbor` after the results heap. When there are twice as many appended + // entries as extras, keep the best `valid_capacity_ - target_valid_` of them plus the + // next one, `candidates_[valid_capacity_]`, which then bounds what is worth appending. void heap_keep_extra(const value_type& neighbor) { - heap_.extras_.push_back(neighbor); - std::push_heap(heap_.extras_.begin(), heap_.extras_.end(), result_compare()); - if (heap_.extras_.size() > valid_capacity_ - target_valid_) { - std::pop_heap(heap_.extras_.begin(), heap_.extras_.end(), result_compare()); - heap_.extras_.pop_back(); + const size_t extra = valid_capacity_ - target_valid_; + if (extra == 0) { + return; + } + candidates_.push_back(neighbor); + if (candidates_.size() - target_valid_ > 2 * extra) { + std::nth_element( + begin() + target_valid_, begin() + valid_capacity_, end(), compare_ + ); + candidates_.resize(valid_capacity_ + 1); } } size_t heap_insert(value_type neighbor) { // Entries worse than the `target_valid_`th result are not visited, but may still - // be kept among the extras. - bool beyond_window = candidates_.size() == target_valid_ && - compare_(candidates_.front().distance(), neighbor.distance()); - bool extras_wanted = - heap_.extras_.size() < valid_capacity_ - target_valid_ || - (!heap_.extras_.empty() && - compare_(neighbor.distance(), heap_.extras_.front().distance())); + // be kept after the results heap. + bool full = candidates_.size() >= target_valid_; + bool beyond_window = + full && compare_(candidates_.front().distance(), neighbor.distance()); // Skip entries that are not wanted, and ids already queued. - if ((beyond_window && !extras_wanted) || - !heap_.seen_.insert(neighbor.id()).second) { + bool extra_wanted = + valid_capacity_ > target_valid_ && + (candidates_.size() <= valid_capacity_ || + compare_(neighbor.distance(), candidates_[valid_capacity_].distance())); + if ((beyond_window && !extra_wanted) || !heap_.seen_.insert(neighbor.id()).second) { return size(); } if (!beyond_window) { @@ -780,20 +786,19 @@ class MutableBuffer { } if (beyond_window) { heap_keep_extra(neighbor); - } else { - // Keep it among the best `target_valid_` results (the worst becomes an extra). + } else if (!full) { candidates_.push_back(neighbor); std::push_heap(begin(), end(), result_compare()); - if (candidates_.size() > target_valid_) { - std::pop_heap(begin(), end(), result_compare()); - value_type worst = candidates_.back(); - candidates_.pop_back(); - if (valid_capacity_ > target_valid_) { - heap_keep_extra(worst); - } - } - } - valid_ = candidates_.size() + heap_.extras_.size(); + } else { + // Replace the worst of the best `target_valid_` results; it moves to the tail. + auto heap_end = begin() + target_valid_; + std::pop_heap(begin(), heap_end, result_compare()); + value_type worst = *(heap_end - 1); + *(heap_end - 1) = neighbor; + std::push_heap(begin(), heap_end, result_compare()); + heap_keep_extra(worst); + } + valid_ = std::min(candidates_.size(), valid_capacity_); return size(); } From 375da1d2dd17567f517a600c6450f1c1dcd4cc02 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Mon, 2 Nov 2026 18:55:00 -0800 Subject: [PATCH 31/33] Trim the extras before the tail outgrows its bound --- include/svs/index/vamana/dynamic_search_buffer.h | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index 00deb17a5..df5ba1941 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -742,16 +742,18 @@ class MutableBuffer { }; } - // Append `neighbor` after the results heap. When there are twice as many appended - // entries as extras, keep the best `valid_capacity_ - target_valid_` of them plus the - // next one, `candidates_[valid_capacity_]`, which then bounds what is worth appending. + // Append `neighbor` after the results heap. When the appended entries first exceed + // the extras, and whenever they are twice as many, keep the best + // `valid_capacity_ - target_valid_` of them plus the next one, + // `candidates_[valid_capacity_]`, which then bounds what is worth appending. void heap_keep_extra(const value_type& neighbor) { const size_t extra = valid_capacity_ - target_valid_; if (extra == 0) { return; } candidates_.push_back(neighbor); - if (candidates_.size() - target_valid_ > 2 * extra) { + if (candidates_.size() == valid_capacity_ + 1 || + candidates_.size() - target_valid_ > 2 * extra) { std::nth_element( begin() + target_valid_, begin() + valid_capacity_, end(), compare_ ); From 1af97ae06452fdbc06cf079e2fee8b97e286bf4c Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Mon, 2 Nov 2026 19:12:41 -0800 Subject: [PATCH 32/33] Return to the two-heap design before the capacity change --- bindings/c/src/filtered_search.hpp | 9 +- include/svs/index/vamana/dynamic_index.h | 63 +++-- .../svs/index/vamana/dynamic_search_buffer.h | 215 +++++------------- include/svs/index/vamana/extensions.h | 35 ++- 4 files changed, 120 insertions(+), 202 deletions(-) diff --git a/bindings/c/src/filtered_search.hpp b/bindings/c/src/filtered_search.hpp index 50a81608d..76956ea25 100644 --- a/bindings/c/src/filtered_search.hpp +++ b/bindings/c/src/filtered_search.hpp @@ -177,9 +177,12 @@ std::vector filtered_topk_search( auto params = index.get_search_parameters(); params.buffer_config_ = svs::index::vamana::SearchBufferConfig{ std::max(results.n_neighbors(), initial_batch_hint)}; - if (index.filtered_search(results.view(), queries, params, [id_filter](size_t id) { - return id_filter->is_member(id); - })) { + if (index.filtered_search( + results.view(), + queries, + params, + [id_filter](size_t id) { return id_filter->is_member(id); } + )) { std::vector found(queries.size(), 0); for (size_t i = 0; i < queries.size(); ++i) { for (size_t j = 0; j < results.n_neighbors(); ++j) { diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index cf08f0164..8da0aa318 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -18,7 +18,6 @@ // stdlib #include -#include #include #include @@ -622,28 +621,6 @@ class MutableVamanaIndex { ); } - private: - // Build the search buffer for one thread of a batch search. - template - auto make_search_buffer( - const search_parameters_type& sp, - size_t num_neighbors, - const std::function* filter - ) const { - MutableBuffer, Filtered> buffer{ - // Legalize search buffer for this search. - sp.buffer_config_.get_total_capacity() < num_neighbors - ? SearchBufferConfig{num_neighbors} - : sp.buffer_config_, - distance::comparator(distance_), - sp.search_buffer_visited_set_}; - if constexpr (Filtered) { - buffer.set_filter(*filter, translator_); - } - return buffer; - } - - public: template void search( QueryResultView results, @@ -652,20 +629,36 @@ class MutableVamanaIndex { const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()), const std::function* filter = nullptr ) { + // Optional `filter`: only entries whose external id passes are returned. It is + // called from all search threads. Rows with fewer passing entries than requested + // are padded with id `size_t(-1)` and the worst score. threads::parallel_for( threadpool_, threads::StaticPartition{queries.size()}, [&](const auto is, uint64_t SVS_UNUSED(tid)) { - auto body = [&](std::bool_constant) { - size_t num_neighbors = results.n_neighbors(); - auto buffer = make_search_buffer(sp, num_neighbors, filter); + size_t num_neighbors = results.n_neighbors(); + auto buffer = search_buffer_type{ + // Legalize search buffer for this search. + sp.buffer_config_.get_total_capacity() < num_neighbors + ? SearchBufferConfig{num_neighbors} + : sp.buffer_config_, + distance::comparator(distance_), + sp.search_buffer_visited_set_}; + // Filtered: the buffer checks the filter itself (two heaps). + std::function internal_filter; + if (filter != nullptr) { + internal_filter = [&](Idx i) { + return (*filter)(translator_.get_external(i)); + }; + buffer.enable_heap(&internal_filter); + } - auto prefetch_parameters = GreedySearchPrefetchParameters{ - sp.prefetch_lookahead_, sp.prefetch_step_}; + auto prefetch_parameters = GreedySearchPrefetchParameters{ + sp.prefetch_lookahead_, sp.prefetch_step_}; - auto scratch = - extensions::per_thread_batch_search_setup(data_, distance_); + auto scratch = extensions::per_thread_batch_search_setup(data_, distance_); + auto run = [&](const auto& search) { extensions::per_thread_batch_search( data_, buffer, @@ -673,16 +666,16 @@ class MutableVamanaIndex { queries, results, threads::UnitRange{is}, - greedy_search_closure(prefetch_parameters, cancel), + search, *this, cancel ); }; - if (filter != nullptr) { - [&]() __attribute__((noinline)) { body(std::true_type{}); } - (); + auto search = greedy_search_closure(prefetch_parameters, cancel); + if (filter == nullptr) { + run(search); } else { - body(std::false_type{}); + run(extensions::FilteredSearch{search, filter}); } } ); diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index df5ba1941..8d001c257 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -33,29 +33,12 @@ namespace svs::index::vamana { -namespace detail { -// State used only by the two-heap mode of `MutableBuffer` (see below). -template struct MutableBufferHeapState { - std::function internal_filter_{}; - bool done_ = false; - Vector to_visit_{}; - std::unordered_set seen_{}; - PredicatedSearchNeighbor popped_{}; -}; -struct NoMutableBufferHeapState {}; -} // namespace detail - /// /// A search buffer that allows entries to be predicated out. /// The search buffer will still navigate through those entries, but it won't include /// them in its final result. /// -/// With `Filtered == true` the buffer takes a filter through `set_filter()` and uses two -/// heaps while searching. With `Filtered == false` (the default) none of the heap code or -/// state exists. -/// -template , bool Filtered = false> -class MutableBuffer { +template > class MutableBuffer { public: // Type Aliases using index_type = Idx; @@ -112,20 +95,18 @@ class MutableBuffer { vector_type candidates_{}; // An optional visited filter. std::optional visited_{std::nullopt}; - // `Filtered` only (`set_filter`): two heaps while searching. `to_visit_` is a - // min-heap of entries still to visit (best on top). Until `cleanup()`, `candidates_` - // holds the results: a max-heap of the best `target_valid_` valid entries that pass - // the filter (worst on top) in `[0, target_valid_)`. Passing entries beyond that are - // appended unsorted after it, trimmed to the best `valid_capacity_ - target_valid_` - // when the tail gets too long. `cleanup()` sorts once and keeps the best - // `valid_capacity_`. + // Optional (`enable_heap`): two heaps while searching. `to_visit_` is a min-heap of + // entries still to visit (best on top). During heap mode `candidates_` holds the + // results: a max-heap of the best valid entries that pass the filter (worst on top). + // `cleanup()` sorts `candidates_` once at the end. // `seen_`: ids already queued in this search (exact duplicate check). // `internal_filter_`: the filter, called with an internal id. - [[no_unique_address]] std::conditional_t< - Filtered, - detail::MutableBufferHeapState, - detail::NoMutableBufferHeapState> - heap_{}; + bool use_heap_ = false; + const std::function* internal_filter_ = nullptr; + bool heap_done_ = false; + vector_type to_visit_{}; + std::unordered_set seen_{}; + value_type popped_{}; public: MutableBuffer() = default; @@ -228,11 +209,9 @@ class MutableBuffer { /// @brief Prepare the buffer for a new search operation. void clear() { candidates_.clear(); - if constexpr (Filtered) { - heap_.to_visit_.clear(); - heap_.seen_.clear(); - heap_.done_ = false; - } + to_visit_.clear(); + seen_.clear(); + heap_done_ = false; best_unvisited_ = 0; roi_end_ = 0; valid_ = 0; @@ -259,25 +238,11 @@ class MutableBuffer { size_t capacity() const { return candidates_.capacity(); } - /// @brief Use two heaps during search. The results heap holds the best - /// `target_window()` entries and bounds the search; up to `target_capacity()` entries - /// are kept in total. A neighbor enters the results only if `filter` accepts its - /// external id (`translator` maps internal to external ids). Both must outlive the - /// use of this buffer. - template - void set_filter(const std::function& filter, const Translator& translator) - requires Filtered - { - heap_.internal_filter_ = [f = &filter, t = &translator](Idx i) { - return (*f)(t->get_external(i)); - }; - } - - /// @brief The filter given to `set_filter`. - const std::function* filter() const - requires Filtered - { - return &heap_.internal_filter_; + /// @brief Use two heaps during search (search window = capacity). A neighbor enters the + /// results only if `*internal_filter` accepts its internal id. + void enable_heap(const std::function* internal_filter) { + use_heap_ = internal_filter != nullptr; + internal_filter_ = internal_filter; } /// @brief Return the number of valid elements currently contained in the buffer. @@ -324,15 +289,11 @@ class MutableBuffer { /// N.B.: If `done()` evaluates to `true`, do not try to extract further candidates /// from it using `next()`. bool done() const { - if constexpr (Filtered) { - if (!heap_.done_) { - // Nothing left to visit, or the best one is worse than the worst result. - return heap_.to_visit_.empty() || (candidates_.size() >= target_valid_ && - compare_( - candidates_.front().distance(), - heap_.to_visit_.front().distance() - )); - } + if (heap_mode()) { + // Nothing left to visit, or the best one is worse than the worst result. + return to_visit_.empty() || + (full() && + compare_(candidates_.front().distance(), to_visit_.front().distance())); } // Until we've reached the target number of valid elements, we have to ignore // the state of `roi_end_`. @@ -349,16 +310,12 @@ class MutableBuffer { /// Post-conditions: /// * The returned neighbor will be marked as visited. const_reference next() { - if constexpr (Filtered) { - if (!heap_.done_) { - std::pop_heap( - heap_.to_visit_.begin(), heap_.to_visit_.end(), to_visit_compare() - ); - heap_.popped_ = heap_.to_visit_.back(); - heap_.to_visit_.pop_back(); - heap_.popped_.set_visited(); - return heap_.popped_; - } + if (heap_mode()) { + std::pop_heap(to_visit_.begin(), to_visit_.end(), to_visit_compare()); + popped_ = to_visit_.back(); + to_visit_.pop_back(); + popped_.set_visited(); + return popped_; } // Get the best unvisited node reference node = getindex(candidates_, best_unvisited_); @@ -380,11 +337,9 @@ class MutableBuffer { /// /// Therefore, sequences of `push_back` must always be followed by a call to `sort`. void push_back(value_type neighbor) { - if constexpr (Filtered) { - if (!heap_.done_) { - heap_insert(neighbor); - return; - } + if (heap_mode()) { + heap_insert(neighbor); + return; } // Follow the contract of not appending more valid elements than the capacity // allows. @@ -537,10 +492,8 @@ class MutableBuffer { /// Insert the given neighbor into the search buffer. size_t insert(value_type neighbor) { - if constexpr (Filtered) { - if (!heap_.done_) { - return heap_insert(neighbor); - } + if (heap_mode()) { + return heap_insert(neighbor); } if (can_skip(neighbor.distance())) { return size(); @@ -594,10 +547,8 @@ class MutableBuffer { // Sort all stored elements in the buffer. void sort() { - if constexpr (Filtered) { - if (!heap_.done_) { - return; // Both heaps are kept in order on every insert. - } + if (heap_mode()) { + return; // Both heaps are kept in order on every insert. } assert(valid() <= valid_capacity_); @@ -648,20 +599,13 @@ class MutableBuffer { // TODO: Switch over to using iterators for the return values to avoid this. void cleanup() { - if constexpr (Filtered) { - if (!heap_.done_) { - // `candidates_` holds only valid entries: sort once and keep the best - // `valid_capacity_`, then continue as a sorted list (e.g. for the re-rank - // and the top-up). - std::sort(begin(), end(), compare_); - if (candidates_.size() > valid_capacity_) { - candidates_.resize(valid_capacity_); - } - valid_ = candidates_.size(); - heap_.to_visit_.clear(); - heap_.seen_.clear(); - heap_.done_ = true; - } + if (heap_mode()) { + // `candidates_` holds only valid entries: sort once, then continue as a sorted + // list (e.g. for the top-up). + std::sort(begin(), end(), compare_); + to_visit_.clear(); + seen_.clear(); + heap_done_ = true; } auto new_end = std::remove_if(begin(), end(), [](const auto& x) { return !x.valid(); }); @@ -728,6 +672,8 @@ class MutableBuffer { } private: + bool heap_mode() const { return use_heap_ && !heap_done_; } + // `to_visit_`: the best entry on top. auto to_visit_compare() const { return [this](const value_type& a, const value_type& b) { @@ -742,65 +688,27 @@ class MutableBuffer { }; } - // Append `neighbor` after the results heap. When the appended entries first exceed - // the extras, and whenever they are twice as many, keep the best - // `valid_capacity_ - target_valid_` of them plus the next one, - // `candidates_[valid_capacity_]`, which then bounds what is worth appending. - void heap_keep_extra(const value_type& neighbor) { - const size_t extra = valid_capacity_ - target_valid_; - if (extra == 0) { - return; - } - candidates_.push_back(neighbor); - if (candidates_.size() == valid_capacity_ + 1 || - candidates_.size() - target_valid_ > 2 * extra) { - std::nth_element( - begin() + target_valid_, begin() + valid_capacity_, end(), compare_ - ); - candidates_.resize(valid_capacity_ + 1); - } - } - size_t heap_insert(value_type neighbor) { - // Entries worse than the `target_valid_`th result are not visited, but may still - // be kept after the results heap. - bool full = candidates_.size() >= target_valid_; - bool beyond_window = - full && compare_(candidates_.front().distance(), neighbor.distance()); - // Skip entries that are not wanted, and ids already queued. - bool extra_wanted = - valid_capacity_ > target_valid_ && - (candidates_.size() <= valid_capacity_ || - compare_(neighbor.distance(), candidates_[valid_capacity_].distance())); - if ((beyond_window && !extra_wanted) || !heap_.seen_.insert(neighbor.id()).second) { + // Skip entries worse than the worst result when full, and ids already queued. + if ((full() && compare_(candidates_.front().distance(), neighbor.distance())) || + !seen_.insert(neighbor.id()).second) { return size(); } - if (!beyond_window) { - heap_.to_visit_.push_back(neighbor); - std::push_heap( - heap_.to_visit_.begin(), heap_.to_visit_.end(), to_visit_compare() - ); - } + to_visit_.push_back(neighbor); + std::push_heap(to_visit_.begin(), to_visit_.end(), to_visit_compare()); // Only valid entries that pass the filter become results; the rest are only // visited. - if (!neighbor.valid() || !heap_.internal_filter_(neighbor.id())) { + if (!neighbor.valid() || !(*internal_filter_)(neighbor.id())) { return size(); } - if (beyond_window) { - heap_keep_extra(neighbor); - } else if (!full) { - candidates_.push_back(neighbor); - std::push_heap(begin(), end(), result_compare()); - } else { - // Replace the worst of the best `target_valid_` results; it moves to the tail. - auto heap_end = begin() + target_valid_; - std::pop_heap(begin(), heap_end, result_compare()); - value_type worst = *(heap_end - 1); - *(heap_end - 1) = neighbor; - std::push_heap(begin(), heap_end, result_compare()); - heap_keep_extra(worst); + // Keep it among the best `valid_capacity_` results (drop the worst). + candidates_.push_back(neighbor); + std::push_heap(begin(), end(), result_compare()); + if (candidates_.size() > valid_capacity_) { + std::pop_heap(begin(), end(), result_compare()); + candidates_.pop_back(); } - valid_ = std::min(candidates_.size(), valid_capacity_); + valid_ = candidates_.size(); return size(); } @@ -826,9 +734,8 @@ class MutableBuffer { } }; -template -std::ostream& -operator<<(std::ostream& io, const MutableBuffer& buffer) { +template +std::ostream& operator<<(std::ostream& io, const MutableBuffer& buffer) { return io << "MutableBuffer<" << datatype_v << ">(" << "target_valid = " << buffer.target_capacity() << ", best_unvisited = " << buffer.best_unvisited() diff --git a/include/svs/index/vamana/extensions.h b/include/svs/index/vamana/extensions.h index 1d6b69053..9ec594486 100644 --- a/include/svs/index/vamana/extensions.h +++ b/include/svs/index/vamana/extensions.h @@ -27,6 +27,7 @@ #include "svs/lib/preprocessor.h" #include "svs/lib/threads.h" +#include #include namespace svs::index::vamana::extensions { @@ -434,18 +435,19 @@ struct VamanaSingleSearchType { /// with additional results. template void check_and_supplement_search_buffer( - const Index& index, SearchBuffer& search_buffer, const Query& query + const Index& index, + SearchBuffer& search_buffer, + const Query& query, + const std::function* filter = nullptr ) { if (search_buffer.valid() < search_buffer.target_window() && search_buffer.valid() < index.size()) { search_buffer.sort(); for (auto external_id : index.external_ids()) { - auto internal_id = index.translate_external_id(external_id); - if constexpr (requires { search_buffer.filter(); }) { - if (!(*search_buffer.filter())(internal_id)) { - continue; // Filtered search: keep filtered-out ids out of the top-up. - } + if (filter != nullptr && !(*filter)(external_id)) { + continue; } + auto internal_id = index.translate_external_id(external_id); auto dist = index.get_distance(external_id, query); auto builder = index.internal_search_builder(); search_buffer.insert(builder(internal_id, dist)); @@ -456,6 +458,15 @@ void check_and_supplement_search_buffer( } } +/// A search closure that also carries the filter of a filtered search. +template struct FilteredSearch { + Search search; + const std::function* filter; + template void operator()(Args&&... args) const { + search(std::forward(args)...); + } +}; + /// Customization point object for processing single queries. inline constexpr VamanaSingleSearchType single_search{}; @@ -488,7 +499,12 @@ SVS_FORCE_INLINE void svs_invoke( search(query, accessor, distance, search_buffer); if constexpr (Index::needs_id_translation) { - check_and_supplement_search_buffer(index, search_buffer, query); + const std::function* filter = nullptr; + if constexpr (requires { search.filter; }) { + filter = + search.filter; // Filtered search: keep filtered-out ids out of the top-up. + } + check_and_supplement_search_buffer(index, search_buffer, query, filter); } } @@ -593,9 +609,8 @@ void svs_invoke( // Copy back results. for (size_t j = 0; j < num_neighbors; ++j) { - if constexpr (requires { search_buffer.filter(); }) { - if (j >= search_buffer.size()) { - // Filtered search with fewer than k passing entries: pad. + if constexpr (requires { search.filter; }) { + if (j >= search_buffer.size()) { // Filtered: fewer than k passed. result.index(i, j) = static_cast(-1); result.distance(i, j) = distance::comparator(distance)(0.0f, 1.0f) ? std::numeric_limits::infinity() From 8a3e7b07ab52603f5fb0abd75f8482a11e76f332 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Mon, 2 Nov 2026 19:20:53 -0800 Subject: [PATCH 33/33] Choose filtered or regular search at compile time --- include/svs/index/vamana/dynamic_index.h | 83 ++++++++++--------- .../svs/index/vamana/dynamic_search_buffer.h | 16 +++- 2 files changed, 58 insertions(+), 41 deletions(-) diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 8da0aa318..534154e26 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -636,46 +636,55 @@ class MutableVamanaIndex { threadpool_, threads::StaticPartition{queries.size()}, [&](const auto is, uint64_t SVS_UNUSED(tid)) { - size_t num_neighbors = results.n_neighbors(); - auto buffer = search_buffer_type{ - // Legalize search buffer for this search. - sp.buffer_config_.get_total_capacity() < num_neighbors - ? SearchBufferConfig{num_neighbors} - : sp.buffer_config_, - distance::comparator(distance_), - sp.search_buffer_visited_set_}; - // Filtered: the buffer checks the filter itself (two heaps). - std::function internal_filter; - if (filter != nullptr) { - internal_filter = [&](Idx i) { - return (*filter)(translator_.get_external(i)); - }; - buffer.enable_heap(&internal_filter); - } + auto body = [&](std::bool_constant) { + size_t num_neighbors = results.n_neighbors(); + auto buffer = MutableBuffer, Filtered>{ + // Legalize search buffer for this search. + sp.buffer_config_.get_total_capacity() < num_neighbors + ? SearchBufferConfig{num_neighbors} + : sp.buffer_config_, + distance::comparator(distance_), + sp.search_buffer_visited_set_}; + // Filtered: the buffer checks the filter itself (two heaps). + std::function internal_filter; + if constexpr (Filtered) { + internal_filter = [&](Idx i) { + return (*filter)(translator_.get_external(i)); + }; + buffer.enable_heap(&internal_filter); + } - auto prefetch_parameters = GreedySearchPrefetchParameters{ - sp.prefetch_lookahead_, sp.prefetch_step_}; - - auto scratch = extensions::per_thread_batch_search_setup(data_, distance_); - - auto run = [&](const auto& search) { - extensions::per_thread_batch_search( - data_, - buffer, - scratch, - queries, - results, - threads::UnitRange{is}, - search, - *this, - cancel - ); + auto prefetch_parameters = GreedySearchPrefetchParameters{ + sp.prefetch_lookahead_, sp.prefetch_step_}; + + auto scratch = + extensions::per_thread_batch_search_setup(data_, distance_); + + auto run = [&](const auto& search) { + extensions::per_thread_batch_search( + data_, + buffer, + scratch, + queries, + results, + threads::UnitRange{is}, + search, + *this, + cancel + ); + }; + auto search = greedy_search_closure(prefetch_parameters, cancel); + if constexpr (Filtered) { + run(extensions::FilteredSearch{search, filter}); + } else { + run(search); + } }; - auto search = greedy_search_closure(prefetch_parameters, cancel); - if (filter == nullptr) { - run(search); + if (filter != nullptr) { + [&]() __attribute__((noinline)) { body(std::true_type{}); } + (); } else { - run(extensions::FilteredSearch{search, filter}); + body(std::false_type{}); } } ); diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index 8d001c257..769a68e00 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -38,7 +38,8 @@ namespace svs::index::vamana { /// The search buffer will still navigate through those entries, but it won't include /// them in its final result. /// -template > class MutableBuffer { +template , bool Filtered = false> +class MutableBuffer { public: // Type Aliases using index_type = Idx; @@ -672,7 +673,13 @@ template > class MutableBuffer { } private: - bool heap_mode() const { return use_heap_ && !heap_done_; } + bool heap_mode() const { + if constexpr (!Filtered) { + return false; + } else { + return use_heap_ && !heap_done_; + } + } // `to_visit_`: the best entry on top. auto to_visit_compare() const { @@ -734,8 +741,9 @@ template > class MutableBuffer { } }; -template -std::ostream& operator<<(std::ostream& io, const MutableBuffer& buffer) { +template +std::ostream& +operator<<(std::ostream& io, const MutableBuffer& buffer) { return io << "MutableBuffer<" << datatype_v << ">(" << "target_valid = " << buffer.target_capacity() << ", best_unvisited = " << buffer.best_unvisited()