From 4fc740c69f4ec3df74478785126ba291cd6328e1 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Wed, 21 Oct 2026 12:38:46 -0700 Subject: [PATCH 1/7] 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 2/7] 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 3/7] 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 4/7] 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 5/7] 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 cc912ec9dcee4ec7e13881b503ebb5c1a29ee0b0 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Fri, 23 Oct 2026 11:45:24 -0700 Subject: [PATCH 6/7] Keep invalid entries in a min-heap for filtered dynamic search With a filter, MutableBuffer keeps entries that fail it in a min-heap instead of the sorted buffer. The sorted buffer then holds only valid entries, so inserts no longer shift long runs of invalid entries. Invalid entries are still visited in distance order, using the same rule as before. --- include/svs/index/vamana/dynamic_index.h | 3 + .../svs/index/vamana/dynamic_search_buffer.h | 70 ++++++++++++++++++- tests/svs/index/vamana/search_buffer.cpp | 49 +++++++++++++ 3 files changed, 121 insertions(+), 1 deletion(-) diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index a78b212e7..7e1bff7a3 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -661,6 +661,9 @@ class MutableVamanaIndex { : sp.buffer_config_, distance::comparator(distance_), sp.search_buffer_visited_set_}; + // Filtered: entries that fail go into the buffer's min-heap, not its + // sorted list. + buffer.enable_invalid_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..a2bf25a58 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -95,6 +95,13 @@ template > class MutableBuffer { vector_type candidates_{}; // An optional visited filter. std::optional visited_{std::nullopt}; + // Optional (`enable_invalid_heap`): invalid entries go into this min-heap (best on + // top) instead of `candidates_`, so `candidates_` holds only valid entries and stays + // short. `invalid_seen_` replaces the duplicate check `candidates_` would have done. + bool use_invalid_heap_ = false; + vector_type invalid_heap_{}; + std::unordered_set invalid_seen_{}; + value_type popped_{}; public: MutableBuffer() = default; @@ -197,6 +204,8 @@ template > class MutableBuffer { /// @brief Prepare the buffer for a new search operation. void clear() { candidates_.clear(); + invalid_heap_.clear(); + invalid_seen_.clear(); best_unvisited_ = 0; roi_end_ = 0; valid_ = 0; @@ -223,6 +232,9 @@ template > class MutableBuffer { size_t capacity() const { return candidates_.capacity(); } + /// @brief Keep invalid entries in a separate min-heap instead of the sorted buffer. + void enable_invalid_heap(bool enable = true) { use_invalid_heap_ = enable; } + /// @brief Return the number of valid elements currently contained in the buffer. size_t valid() const { return valid_; } @@ -269,7 +281,8 @@ template > class MutableBuffer { bool done() const { // 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_); + return best_unvisited_ == ((valid() < target_valid_) ? size() : roi_end_) && + !invalid_ready(); } /// @brief Return the best unvisited neighbor in the buffer. @@ -282,6 +295,21 @@ template > class MutableBuffer { /// Post-conditions: /// * The returned neighbor will be marked as visited. const_reference next() { + // The best invalid entry goes first if it beats the best unvisited valid one. + if (invalid_ready()) { + size_t end = (valid() < target_valid_) ? size() : roi_end_; + if (best_unvisited_ == end || compare_( + invalid_heap_.front().distance(), + candidates_[best_unvisited_].distance() + )) { + std::pop_heap(invalid_heap_.begin(), invalid_heap_.end(), heap_compare()); + popped_ = invalid_heap_.back(); + invalid_heap_.pop_back(); + popped_.set_visited(); + return popped_; + } + } + // Get the best unvisited node reference node = getindex(candidates_, best_unvisited_); node.set_visited(); @@ -308,6 +336,10 @@ template > class MutableBuffer { if (full() && valid) { return; } + if (use_invalid_heap_ && !valid) { + push_invalid(neighbor); + return; + } candidates_.push_back(neighbor); if (valid) { ++valid_; @@ -460,6 +492,10 @@ template > class MutableBuffer { } size_t insert_inner(value_type neighbor) { + if (use_invalid_heap_ && !neighbor.valid()) { + push_invalid(neighbor); + return size(); + } const auto start = begin(); // Binary search to the first location where `distance` is less than the stored @@ -557,6 +593,8 @@ template > class MutableBuffer { auto new_end = std::remove_if(begin(), end(), [](const auto& x) { return !x.valid(); }); candidates_.resize(new_end - begin()); + invalid_heap_.clear(); + invalid_seen_.clear(); } /// @brief Return the size of the underlying vector. @@ -619,6 +657,36 @@ template > class MutableBuffer { } private: + // Heap order for `invalid_heap_`: the best (closest) entry on top. + auto heap_compare() const { + return [this](const value_type& a, const value_type& b) { + return compare_(b.distance(), a.distance()); + }; + } + + // Queue an invalid entry once per search. + void push_invalid(value_type neighbor) { + if (invalid_seen_.insert(neighbor.id()).second) { + invalid_heap_.push_back(neighbor); + std::push_heap(invalid_heap_.begin(), invalid_heap_.end(), heap_compare()); + } + } + + // Is the best invalid entry worth expanding? Same rule as an invalid entry in the + // sorted buffer: always while fewer than `target_valid_` valid entries are found, + // afterwards only if it is closer than the `target_valid_`th valid entry. + bool invalid_ready() const { + if (invalid_heap_.empty()) { + return false; + } + if (valid() < target_valid_) { + return true; + } + return compare_( + invalid_heap_.front().distance(), candidates_[roi_end_ - 1].distance() + ); + } + /// 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..c73f760ea 100644 --- a/tests/svs/index/vamana/search_buffer.cpp +++ b/tests/svs/index/vamana/search_buffer.cpp @@ -906,6 +906,55 @@ CATCH_TEST_CASE("MutableBuffer", "[core][search_buffer]") { CATCH_REQUIRE(b.done()); } } + + // With `enable_invalid_heap`, invalid entries are kept out of the sorted buffer but are + // still visited in distance order, together with the valid entries. + CATCH_SECTION("Invalid heap") { + using T = svs::PredicatedSearchNeighbor; + auto make_visited = [](uint32_t id, float distance, bool valid) { + auto n = T{id, distance, valid}; + n.set_visited(); + return n; + }; + buffer.enable_invalid_heap(); + + CATCH_SECTION("Not full") { + 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(eq(buffer.next(), make_visited(11, 1, false))); + CATCH_REQUIRE(eq(buffer.next(), make_visited(10, 5, false))); + CATCH_REQUIRE(eq(buffer.next(), make_visited(1, 50, true))); + CATCH_REQUIRE(eq(buffer.next(), make_visited(0, 100, true))); + CATCH_REQUIRE(buffer.done()); + } + + CATCH_SECTION("Full") { + // Invalid entries beyond the 4th valid entry are not visited. + for (uint32_t i = 0; i < 4; ++i) { + buffer.insert({i, svs::lib::narrow_cast(10 * (i + 1)), true}); + } + buffer.insert({10, 35, false}); + buffer.insert({11, 50, false}); + CATCH_REQUIRE(buffer.size() == 4); + CATCH_REQUIRE(buffer.full()); + + CATCH_REQUIRE(eq(buffer.next(), make_visited(0, 10, true))); + CATCH_REQUIRE(eq(buffer.next(), make_visited(1, 20, true))); + CATCH_REQUIRE(eq(buffer.next(), make_visited(2, 30, true))); + CATCH_REQUIRE(eq(buffer.next(), make_visited(10, 35, false))); + CATCH_REQUIRE(eq(buffer.next(), make_visited(3, 40, true))); + CATCH_REQUIRE(buffer.done()); + + buffer.cleanup(); + CATCH_REQUIRE(buffer.size() == 4); + } + } } CATCH_TEST_CASE("Fuzzing Mutable", "[core][search_buffer]") { From 3145c7349a7996c365d292c14feca2fb9923d506 Mon Sep 17 00:00:00 2001 From: yuejiaointel Date: Fri, 23 Oct 2026 11:47:13 -0700 Subject: [PATCH 7/7] Test the stop rule for queued invalid entries --- tests/svs/index/vamana/search_buffer.cpp | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/svs/index/vamana/search_buffer.cpp b/tests/svs/index/vamana/search_buffer.cpp index c73f760ea..282f1fdc0 100644 --- a/tests/svs/index/vamana/search_buffer.cpp +++ b/tests/svs/index/vamana/search_buffer.cpp @@ -923,7 +923,7 @@ CATCH_TEST_CASE("MutableBuffer", "[core][search_buffer]") { buffer.insert({0, 100, true}); buffer.insert({11, 1, false}); buffer.insert({1, 50, true}); - buffer.insert({11, 1, false}); // duplicate: ignored + buffer.insert({11, 1, false}); // duplicate: visited once (checked by done()) CATCH_REQUIRE(buffer.size() == 2); CATCH_REQUIRE(buffer.valid() == 2); @@ -935,12 +935,15 @@ CATCH_TEST_CASE("MutableBuffer", "[core][search_buffer]") { } CATCH_SECTION("Full") { - // Invalid entries beyond the 4th valid entry are not visited. - for (uint32_t i = 0; i < 4; ++i) { + // Invalid entries beyond the 4th valid entry are not visited. Both invalid + // entries are queued before the buffer is full, so neither is skipped on + // insert. + for (uint32_t i = 0; i < 3; ++i) { buffer.insert({i, svs::lib::narrow_cast(10 * (i + 1)), true}); } buffer.insert({10, 35, false}); buffer.insert({11, 50, false}); + buffer.insert({3, 40, true}); CATCH_REQUIRE(buffer.size() == 4); CATCH_REQUIRE(buffer.full());