From 11dd6442a420df112bfa9d07d0d73a04d6e3da6e Mon Sep 17 00:00:00 2001 From: Rafik Saliev Date: Fri, 13 Nov 2026 06:37:53 -0700 Subject: [PATCH 1/5] fix(test): add scope guard for logger sink cleanup in dynamic index tests --- tests/svs/index/vamana/dynamic_index_2.cpp | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/svs/index/vamana/dynamic_index_2.cpp b/tests/svs/index/vamana/dynamic_index_2.cpp index 2daeda29e..da85a6903 100644 --- a/tests/svs/index/vamana/dynamic_index_2.cpp +++ b/tests/svs/index/vamana/dynamic_index_2.cpp @@ -21,6 +21,7 @@ #include "svs/index/vamana/dynamic_index.h" #include "svs/lib/float16.h" #include "svs/lib/preprocessor.h" +#include "svs/lib/scopeguard.h" #include "svs/lib/timing.h" #include "svs/misc/dynamic_helper.h" @@ -283,6 +284,13 @@ CATCH_TEST_CASE("Testing Graph Index", "[graph_index][dynamic_index]") { global_callback_sink->set_level(spdlog::level::trace); auto original_logger = svs::logging::get(); original_logger->sinks().push_back(global_callback_sink); + // The sink captures a local vector: remove it before the vector goes out of scope. + auto remove_sink = svs::lib::make_scope_guard([&]() noexcept { + auto& sinks = original_logger->sinks(); + sinks.erase( + std::remove(sinks.begin(), sinks.end(), global_callback_sink), sinks.end() + ); + }); // Load the base dataset and queries. auto data = svs::data::SimpleData::load(test_dataset::data_svs_file()); @@ -495,6 +503,13 @@ CATCH_TEST_CASE("Dynamic MutableVamanaIndex Per-Index Logging Test", "[logging]" auto original_logger = svs::logging::get(); original_logger->sinks().push_back(global_callback_sink); + // The sink captures a local vector: remove it before the vector goes out of scope. + auto remove_sink = svs::lib::make_scope_guard([&]() noexcept { + auto& sinks = original_logger->sinks(); + sinks.erase( + std::remove(sinks.begin(), sinks.end(), global_callback_sink), sinks.end() + ); + }); // Setup index auto data = svs::data::SimpleData::load(test_dataset::data_svs_file()); From b5afca39a9f7417e7af22c198aaa4011b45550df Mon Sep 17 00:00:00 2001 From: Rafik Saliev Date: Fri, 13 Nov 2026 03:58:25 -0700 Subject: [PATCH 2/5] Implement predicated search and range search in MutableBuffer - Removed the sort() method from concurrent::MutableBuffer to streamline the search process. - Introduced predicated_buffer_type and predicated_scratchspace_type for handling searches with stop predicates. - Added DefaultStop and RangeStop predicates to control search termination based on conditions. - Enhanced MutableBuffer to maintain state information for valid candidates and total candidates. - Updated the search logic in MutableVamanaIndex to utilize the new predicated search functionality. - Added tests for predicated and range search to ensure correctness and performance. - Improved handling of candidate insertion and buffer state management. --- .../svs/concurrent/dynamic_search_buffer.h | 59 ---- include/svs/index/vamana/dynamic_index.h | 177 ++++++++++ .../svs/index/vamana/dynamic_search_buffer.h | 320 +++++++++--------- tests/svs/index/vamana/dynamic_index_2.cpp | 139 ++++++++ tests/svs/index/vamana/search_buffer.cpp | 195 ++++++++++- 5 files changed, 662 insertions(+), 228 deletions(-) diff --git a/include/svs/concurrent/dynamic_search_buffer.h b/include/svs/concurrent/dynamic_search_buffer.h index 0484f228d..c51a685aa 100644 --- a/include/svs/concurrent/dynamic_search_buffer.h +++ b/include/svs/concurrent/dynamic_search_buffer.h @@ -44,65 +44,6 @@ class MutableBuffer : public vamana::MutableBuffer { return MutableBuffer{ base_type::config(), this->compare_, base_type::visited_set_enabled()}; } - - // Sort all stored elements in the buffer. - void sort() { - assert(base_type::valid() <= this->valid_capacity_); - - // Put the elements in ascending order. - std::sort(base_type::begin(), base_type::end(), this->compare_); - - // Now - fix our invariants. - if (this->valid_ < this->target_valid_) { - // Maintain 5A - // Invariant 6 has not activated. - this->roi_end_ = base_type::size(); - - // Maintain best_unvisited_ invariant. - // Recompute it from scratch, mirroring the 5B branch below. - for (this->best_unvisited_ = 0; this->best_unvisited_ < this->roi_end_; - ++this->best_unvisited_) { - if (!this->candidates_[this->best_unvisited_].visited()) { - break; - } - } - return; - } - - // Maintain 5B - size_t valid_count = 0; - for (size_t i = 0; i < base_type::size(); ++i) { - if (this->candidates_[i].valid()) { - ++valid_count; - if (valid_count == this->target_valid_) { - // One past the last valid element. - this->roi_end_ = i + 1; - break; - } - } - } - - // TODO: Switch over to using iterators for the return values to avoid this. - // Maintain best_unvisited_ invariant - // cleanup() and std::sort() may invalidate best_unvisited_. - // so, have to lookup the next unvisited from scratch. - for (this->best_unvisited_ = 0; this->best_unvisited_ < this->roi_end_; - ++this->best_unvisited_) { - if (!this->candidates_[this->best_unvisited_].visited()) { - break; - } - } - - // Check if invariant 6 is active. - // If so, drop invalid elements off the end until a valid element is found. - if (base_type::slack() == 0) { - assert(!this->candidates_.empty()); - while (!base_type::back().valid()) { - this->candidates_.pop_back(); - assert(!this->candidates_.empty()); - } - } - } }; } // namespace svs::index::vamana::concurrent diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 8a9a64dd1..9fe811036 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -771,6 +771,183 @@ class MutableVamanaIndex { translate_to_external(results.indices()); } + ///// Predicated Search + + /// The search buffer type terminating with the stop predicate `Stop`. + template + using predicated_buffer_type = MutableBuffer, Stop>; + + /// The scratchspace type for ``predicated_search`` with the stop predicate `Stop`. + template + using predicated_scratchspace_type = + SearchScratchspace, inner_scratch_type>; + + /// + /// @brief Return a scratchspace for ``predicated_search``. + /// + /// @param sp The search parameters. A zero total buffer capacity makes the search + /// buffer unbounded, otherwise the capacity bounds the number of results. + /// @param stop The stop predicate invoked by the search buffer with a + /// ``MutableBufferState``. The search terminates when it returns ``true`` or when + /// no unvisited candidate is left. + /// + template + predicated_scratchspace_type + predicated_scratchspace(const search_parameters_type& sp, Stop stop) const { + return predicated_scratchspace_type{ + predicated_buffer_type{ + predicated_buffer_config(sp.buffer_config_, 0), + distance::comparator(distance_), + sp.search_buffer_visited_set_, + std::move(stop)}, + extensions::single_search_setup(data_, distance_), + {sp.prefetch_lookahead_, sp.prefetch_step_}}; + } + + /// + /// @brief Search for a single query, terminating with the scratchspace stop predicate. + /// + /// After the call, ``scratch.buffer`` holds the valid results with internal IDs, sorted + /// from best to worst. + /// + template + void predicated_search( + const Query& query, + predicated_scratchspace_type& scratch, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + ) const { + extensions::single_search( + data_, + scratch.buffer, + scratch.scratch, + query, + greedy_search_closure(scratch.prefetch_parameters, cancel), + *this, + cancel + ); + } + + /// + /// @brief Batch search terminating with the stop predicate `stop`. + /// + /// Up to ``results.n_neighbors()`` results with external IDs are stored per query. + /// Missing results are padded with sentinel neighbors. + /// + template + void predicated_search( + QueryResultView results, + const Queries& queries, + const search_parameters_type& sp, + const Stop& stop, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + ) { + using compare_type = distance::compare_t; + constexpr auto sentinel = type_traits::sentinel_v, compare_type>; + const size_t num_neighbors = results.n_neighbors(); + threads::parallel_for( + threadpool_, + threads::StaticPartition{queries.size()}, + [&](const auto is, uint64_t SVS_UNUSED(tid)) { + auto scratch = predicated_scratchspace_type{ + predicated_buffer_type{ + predicated_buffer_config(sp.buffer_config_, num_neighbors), + distance::comparator(distance_), + sp.search_buffer_visited_set_, + stop}, + extensions::single_search_setup(data_, distance_), + {sp.prefetch_lookahead_, sp.prefetch_step_}}; + + for (auto i : is) { + if (cancel()) { + return; + } + predicated_search(queries.get_datum(i), scratch, cancel); + + const auto& buffer = scratch.buffer; + const size_t found = std::min(buffer.size(), num_neighbors); + for (size_t j = 0; j < found; ++j) { + const auto& neighbor = buffer[j]; + results.set( + Neighbor{ + lib::narrow_cast(translate_internal_id(neighbor.id())), + neighbor.distance()}, + i, + j + ); + } + for (size_t j = found; j < num_neighbors; ++j) { + results.set(sentinel, i, j); + } + } + } + ); + } + + /// + /// @brief Return all neighbors within ``radius`` of each query. + /// + /// A neighbor at distance ``d`` is within the radius if ``compare(d, radius)`` for the + /// index comparator (e.g., ``d < radius`` for L2, ``d > radius`` for inner product). + /// The search visits at least the search window, then continues while the best + /// unvisited candidate is within the radius. + /// + /// @param queries The queries. + /// @param radius The search radius. + /// @param sp The search parameters. A zero total buffer capacity makes the number of + /// results per query unbounded, otherwise the capacity bounds it. + /// @param cancel Cancellation predicate. + /// + /// @returns The neighbors with external IDs for each query, sorted from best to worst. + /// + template + std::vector>> range_search( + const Queries& queries, + float radius, + const search_parameters_type& sp, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + ) { + using stop_type = RangeStop>; + const auto compare = distance::comparator(distance_); + auto results = std::vector>>(queries.size()); + threads::parallel_for( + threadpool_, + threads::StaticPartition{queries.size()}, + [&](const auto is, uint64_t SVS_UNUSED(tid)) { + auto scratch = predicated_scratchspace(sp, stop_type{radius, compare}); + for (auto i : is) { + if (cancel()) { + return; + } + predicated_search(queries.get_datum(i), scratch, cancel); + + auto& result = results[i]; + for (const auto& neighbor : scratch.buffer) { + if (!compare(neighbor.distance(), radius)) { + break; + } + result.emplace_back( + translate_internal_id(neighbor.id()), neighbor.distance() + ); + } + } + } + ); + return results; + } + + private: + // Map search parameters to a buffer configuration for predicated search. + static SearchBufferConfig + predicated_buffer_config(const SearchBufferConfig& config, size_t min_capacity) { + const size_t window = config.get_search_window_size(); + const size_t capacity = config.get_total_capacity(); + if (capacity == 0) { + return SearchBufferConfig{window, search_buffer_type::unbounded_capacity}; + } + return SearchBufferConfig{window, std::max(capacity, min_capacity)}; + } + + 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 52b071d9a..2d52544b6 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -26,6 +26,7 @@ #include #include #include +#include #include #include #include @@ -33,12 +34,62 @@ namespace svs::index::vamana { +/// @brief Snapshot of the `MutableBuffer` search state passed to its stop predicate. +/// +/// Only constructed when an unvisited candidate exists, so `best_unvisited_distance` is +/// always the distance of a real candidate. +struct MutableBufferState { + /// Number of valid candidates in the buffer, i.e. the results available to the user. + size_t valid; + /// Number of candidates, valid or not, offered to the buffer by `push_back()` or + /// `insert()` since the last `clear()`, including dropped ones. + /// Counted only if the stop predicate declares `counts_total_candidates = true`. + size_t total_candidates; + /// Number of valid candidates preceding the best unvisited one, all of them visited. + size_t valid_visited; + /// Distance of the best unvisited candidate to the query. + float best_unvisited_distance; + /// Target number of valid candidates (search window size). + size_t search_window; +}; + +/// @brief Whether the buffer must count all candidates for the stop predicate. +template +inline constexpr bool counts_total_candidates_v = + requires { requires Stop::counts_total_candidates; }; + +/// @brief Default stop predicate: done once the best `search_window` valid candidates +/// have all been visited. +struct DefaultStop { + constexpr bool operator()(const MutableBufferState& state) const { + return state.valid_visited >= state.search_window; + } +}; + +/// @brief Stop predicate for range search. +/// +/// Visits at least the search window, then continues while the best unvisited candidate +/// is within `radius`. A distance `d` is within the radius if `compare(d, radius)`. +template > struct RangeStop { + float radius; + [[no_unique_address]] Cmp compare{}; + + constexpr bool operator()(const MutableBufferState& state) const { + return state.valid_visited >= state.search_window && + !compare(state.best_unvisited_distance, radius); + } +}; + /// /// 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 { +/// The search terminates when no unvisited candidate is left or when the `Stop` predicate, +/// invoked with a `MutableBufferState`, returns `true`. +/// +template , typename Stop = DefaultStop> +class MutableBuffer { public: // Type Aliases using index_type = Idx; @@ -46,6 +97,7 @@ template > class MutableBuffer { using reference = value_type&; using const_reference = const value_type&; using compare_type = Cmp; + using stop_type = Stop; using vector_type = std::vector>; using iterator = typename vector_type::iterator; @@ -54,6 +106,9 @@ template > class MutableBuffer { /// A visited filter with 65,536 entries with a memory footprint of 128 kiB. using set_type = VisitedFilter; + /// Total capacity value that lets the buffer grow without bound. + static constexpr size_t unbounded_capacity = std::numeric_limits::max(); + protected: ///// Invariants: // @@ -63,14 +118,10 @@ template > class MutableBuffer { // TODO: Enforce the first in `BufferConfig` as well. // // 2. best_unvisited_ <= size(). - // 3. roi_end_ <= size(). + // 3. `best_unvisited_` is the first unvisited entry (or `size()` if there is none). // 4. valid_ <= valid_capacity_; // - // 5. WHEN (5A): `valid_ < target_valid_` (haven't yet found enough valid neighbors) - // THEN: `roi_end_ == size()`: Region of interest points to the end of the buffer. - // ELSE (5B): `valid_ >= target_valid_` - // THEN: The number of valid between the start of the buffer and `roi_end_ - 1` is - // exactly `valid_`. + // 5. `valid_visited_` is the number of valid elements in `[0, best_unvisited_)`. // // 6. WHEN: `valid_ == valid_capacity_` // THEN: `back()` must be valid. @@ -85,8 +136,10 @@ template > class MutableBuffer { size_t valid_capacity_ = 0; // Index of the best unvisited candidate. size_t best_unvisited_ = 0; - // One past the "target_valid_"th entry. - size_t roi_end_ = 0; + // Number of valid neighbors before `best_unvisited_`. + size_t valid_visited_ = 0; + // Number of candidates offered to the buffer. + size_t total_candidates_ = 0; // number of valid neighbors. size_t valid_ = 0; // A buffer of candidates. @@ -95,27 +148,39 @@ template > class MutableBuffer { vector_type candidates_{}; // An optional visited filter. std::optional visited_{std::nullopt}; + [[no_unique_address]] Stop stop_{}; public: MutableBuffer() = default; /// Construct a new buffer with the given buffer configuration. explicit MutableBuffer( - SearchBufferConfig config, Cmp compare = Cmp{}, bool enable_visited = false + SearchBufferConfig config, + Cmp compare = Cmp{}, + bool enable_visited = false, + Stop stop = Stop{} ) : compare_{std::move(compare)} , target_valid_{config.get_search_window_size()} , valid_capacity_{config.get_total_capacity()} - , candidates_{valid_capacity_} { - candidates_.clear(); + , stop_{std::move(stop)} { + candidates_.reserve( + valid_capacity_ == unbounded_capacity ? target_valid_ : valid_capacity_ + ); if (enable_visited) { enable_visited_set(); } } /// Construct a new buffer with the given size and capacity. - explicit MutableBuffer(size_t size, Cmp compare = Cmp{}, bool enable_visited = false) - : MutableBuffer{SearchBufferConfig{size}, std::move(compare), enable_visited} {} + explicit MutableBuffer( + size_t size, Cmp compare = Cmp{}, bool enable_visited = false, Stop stop = Stop{} + ) + : MutableBuffer{ + SearchBufferConfig{size}, + std::move(compare), + enable_visited, + std::move(stop)} {} /// Estimate heap memory footprint of a search buffer with the given configuration. static constexpr size_t estimate_memory_footprint( @@ -156,7 +221,7 @@ template > class MutableBuffer { MutableBuffer shallow_copy() const { // We don't care about the contents of the buffer - just its size. // Therefore, we can construct a new buffer from scratch. - return MutableBuffer{config(), compare_, visited_set_enabled()}; + return MutableBuffer{config(), compare_, visited_set_enabled(), stop_}; } // TODO: Allow this construction to be noexcept. @@ -183,6 +248,7 @@ template > class MutableBuffer { // If the new capacity is lower then the current size, shrink the buffer to fit. if (valid_capacity_temp < candidates_.size()) { candidates_.resize(valid_capacity_temp); + restore_best_unvisited(); } // Commit the new sizes. Integer assignment is `noexcept`, so we don't need to worry @@ -198,7 +264,8 @@ template > class MutableBuffer { void clear() { candidates_.clear(); best_unvisited_ = 0; - roi_end_ = 0; + valid_visited_ = 0; + total_candidates_ = 0; valid_ = 0; if (visited_set_enabled()) { visited_->reset(); @@ -219,6 +286,7 @@ template > class MutableBuffer { } best_unvisited_ = 0; + valid_visited_ = 0; } size_t capacity() const { return candidates_.capacity(); } @@ -262,14 +330,28 @@ template > class MutableBuffer { /// @brief Return the index of the best unvisited candidate. size_t best_unvisited() const { return best_unvisited_; } + /// @brief Return the number of valid candidates preceding the best unvisited one. + size_t valid_visited() const { return valid_visited_; } + + /// @brief Return the number of candidates offered to the buffer since `clear()`. + /// Only counted for stop predicates that opt in. + size_t total_candidates() const { return total_candidates_; } + + /// @brief Return the stop predicate. + const Stop& stop_predicate() const { return stop_; } + /// @brief Return `true` if the search buffer has reached its terminating condition. /// /// N.B.: If `done()` evaluates to `true`, do not try to extract further candidates /// from it using `next()`. 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_ == size() || + stop_(MutableBufferState{ + .valid = valid_, + .total_candidates = total_candidates_, + .valid_visited = valid_visited_, + .best_unvisited_distance = candidates_[best_unvisited_].distance(), + .search_window = target_valid_}); } /// @brief Return the best unvisited neighbor in the buffer. @@ -285,10 +367,18 @@ template > class MutableBuffer { // Get the best unvisited node reference node = getindex(candidates_, best_unvisited_); node.set_visited(); + valid_visited_ += static_cast(node.valid()); // Increment `best_unvisited_` until it's equal to the size OR until we encounter // an unvisited node. - while (++best_unvisited_ != roi_end_ && candidates_[best_unvisited_].visited()) {} + const size_t end = size(); + while (++best_unvisited_ != end) { + const auto& candidate = candidates_[best_unvisited_]; + if (!candidate.visited()) { + break; + } + valid_visited_ += static_cast(candidate.valid()); + } return node; } @@ -305,6 +395,7 @@ template > class MutableBuffer { // Follow the contract of not appending more valid elements than the capacity // allows. bool valid = neighbor.valid(); + count_candidate(); if (full() && valid) { return; } @@ -312,7 +403,6 @@ template > class MutableBuffer { if (valid) { ++valid_; } - roi_end_ = candidates_.size(); } // Define iterators. @@ -327,121 +417,25 @@ template > class MutableBuffer { /// /// May invalidate any iterators for `candidates_`. /// - void unsafe_insert(value_type neighbor, iterator pos, size_t offset) { + void unsafe_insert(value_type neighbor, iterator pos) { // Insert the neighbor into the buffer. // The rest of this function fixes the invariants that got broken as a result // of this operation. candidates_.insert(pos, neighbor); - // Enough candidates for ROI to kick in (5B instead of 5A). - bool at_threshold = valid() == target_valid_; - bool this_valid = neighbor.valid(); - // Is this candidate begin inserted before the end of the ROI. - bool inserted_below = offset < roi_end_; - if (valid() < target_valid_) { - // Maintain invariant 5A - roi_end_ = size(); - } else if (at_threshold) { - // Several options: - // (1) This is the entry that causes us to cross the valid threshold. If so, - // then we need to move `roi_end_` to the last valid entry; - // (2) We've previously crossed the threshold and this entry is skipped. - // If this is the case, then we must bump `roi_end_` if this neighbor - // is inserted below. - if (this_valid) { - assert(roi_end_ == size() - 1); - // Case 1: This neighbor caused us to cross the threshold. - // - // In this case, the `roi_end_` USED to point to one past the end. - // After the insertion, it points to the end exactly. - // - // To maintain 5B, move `roi_end_` backwards until a valid neighbor is - // reached, then increment by 1 to point to one-past the end. - while (!svs::getindex(candidates_, roi_end_).valid()) { - --roi_end_; - } - ++roi_end_; - } else if (inserted_below) { - // Case 2: We've previously crossed the threshold and this neighbor is - // skipped. - // - // In this case, `roi_end_` used to point to one-past the last valid element - // (according to 5B). - // - // Since we inserted below (moving everything above by 1), we can restore - // 5B by incrementing `roi_end_`. - ++roi_end_; - } - } else { - // At this point - we are in split-buffer territory. - // We can assume that before the insertion, invariante 5B held. - // - // Therefore, we only need to fix 5B if we inserted below the previous - // `roi_end_`. - if (inserted_below) { - // Assert that 5B held prior to this invocation of `unsafe_insert`. - assert(candidates_.at(roi_end_).valid()); - - // If this neighbor is not valid, then the number of valid neighbors in - // `[0, roi_end_]` (with right inclusion) has not changed. - // We can restore 5B by incrementing `roi_end_`. - if (!this_valid) { - ++roi_end_; - } else { - // In this case, the number of neighbors in `[0, roi_end_`] (with - // right inclusion) has been incremented by 1. - // - // We restore 5B by walking back until the next previous valid neighbor - // is found, then step forward by 1 to point to one-past the end. - while (!svs::getindex(candidates_, --roi_end_).valid()) {} - ++roi_end_; - } - } - } - - // Restore Invariant 6. - bool no_split_buffer = target_valid_ == valid_capacity_; - if (no_split_buffer) { - // When using 5A: the ROI already points to the end - so resizing does nothing. - // Furthermore, when using 5A, `valid() > target_valid_` is always false, so - // the branch is never taken. - // - // When using 5B: `roi_end_` points to one-past the `target_valid_`th entry - // Shrinking `candidates_` to `roi_end_` is valid - // - // Prior to this call to, invariants 4 and 7 must have held. - // Therefore, the side-effect of adding this neighbor increased the number of - // valid elements by at most 1 (i.e., `slack() <= 1`). - // - // From this, we deduce that the extra valid neighbor used to be located at or - // above `roi_end_` and was therefore implicitly dropped when shrinking - // `candidates_`. - candidates_.resize(roi_end_); - if (slack() > 0) { - assert(slack() == 1); - valid_ -= 1; - } - } else if (slack() == 0) { - // From invariant 7, we know `candidates_` contains `valid_` number of elements. - // We can restore invariant 6 simply by finding the last valid neighbor and - // resizing appropriately. - candidates_.resize(walk_back(candidates_.size() - 1) + 1); - } else if (slack() >= 0) { - // We know that invariant 6 must have held prior to this invocation of - // `unsafe_insert`. - // - // Therefore the last neighbor in `candidates_` is valid and must be dropped. - // We can restore invariant 6 by starting at this neighbor and walking back - // until the previous valid neighbor is found. - assert(slack() == 1); + // Restore Invariants 4 and 6. + if (valid_ > valid_capacity_) { + // Invariant 6 held prior to the insertion, so the last element is valid and is + // the one that must be dropped. Then walk back to the previous valid element. + assert(valid_ == valid_capacity_ + 1); assert(back().valid()); candidates_.resize(walk_back(candidates_.size() - 2) + 1); // ^ // | // Not a typo. It should be "two". - - // We dropped a valid element off the back when shrinking. valid_ -= 1; + } else if (valid_ == valid_capacity_) { + candidates_.resize(walk_back(candidates_.size() - 1) + 1); } } @@ -453,6 +447,7 @@ template > class MutableBuffer { /// Insert the given neighbor into the search buffer. size_t insert(value_type neighbor) { + count_candidate(); if (can_skip(neighbor.distance())) { return size(); } @@ -497,9 +492,16 @@ template > class MutableBuffer { // Note: Iterators *may* be invalidated as a result of `unsafe_insert`. // Hoist out the insertion position as an index before inserting. size_t i = pos - start; + if (i < best_unvisited_) { + // Maintain Invariant 5: the elements in `[i, best_unvisited_)` will no longer + // precede the best unvisited candidate. + for (size_t j = i; j < best_unvisited_; ++j) { + valid_visited_ -= static_cast(candidates_[j].valid()); + } + best_unvisited_ = i; + } valid_ += static_cast(neighbor.valid()); - unsafe_insert(neighbor, pos, i); - best_unvisited_ = std::min(best_unvisited_, i); + unsafe_insert(neighbor, pos); return i; } @@ -510,46 +512,18 @@ template > class MutableBuffer { // Put the elements in ascending order. std::sort(begin(), end(), compare_); - // Now - fix our invariants. - if (valid_ < target_valid_) { - // Maintain 5A - // Invariant 6 has not activated. - roi_end_ = size(); - return; - } - - // Maintain 5B - size_t valid_count = 0; - for (size_t i = 0; i < size(); ++i) { - if (candidates_[i].valid()) { - ++valid_count; - if (valid_count == target_valid_) { - // One past the last valid element. - roi_end_ = i + 1; - break; - } - } - } - - // TODO: Switch over to using iterators for the return values to avoid this. - // Maintain best_unvisited_ invariant - // cleanup() and std::sort() may invalidate best_unvisited_. - // so, have to lookup the next unvisited from scratch. - for (best_unvisited_ = 0; best_unvisited_ < roi_end_; ++best_unvisited_) { - if (!candidates_[best_unvisited_].visited()) { - break; - } - } - // Check if invariant 6 is active. // If so, drop invalid elements off the end until a valid element is found. - if (slack() == 0) { + if (valid_ == valid_capacity_) { assert(!candidates_.empty()); while (!back().valid()) { candidates_.pop_back(); assert(!candidates_.empty()); } } + + // Sorting may move visited candidates anywhere, so recompute Invariants 3 and 5. + restore_best_unvisited(); } // TODO: Switch over to using iterators for the return values to avoid this. @@ -619,15 +593,27 @@ template > class MutableBuffer { } protected: - /// 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. - int64_t slack() const { - return lib::narrow_cast(valid()) - - lib::narrow_cast(target_capacity()); + /// Recompute `best_unvisited_` and `valid_visited_` from scratch. + void restore_best_unvisited() { + valid_visited_ = 0; + const size_t end = size(); + for (best_unvisited_ = 0; best_unvisited_ < end; ++best_unvisited_) { + const auto& candidate = candidates_[best_unvisited_]; + if (!candidate.visited()) { + break; + } + valid_visited_ += static_cast(candidate.valid()); + } } private: + // Opt-in: this update is measurable in the hot insertion path. + void count_candidate() { + if constexpr (counts_total_candidates_v) { + ++total_candidates_; + } + } + /// Return the index of the first preceding valid candidate beginning at the provided /// index. /// @@ -642,8 +628,8 @@ 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() diff --git a/tests/svs/index/vamana/dynamic_index_2.cpp b/tests/svs/index/vamana/dynamic_index_2.cpp index da85a6903..0e58e3327 100644 --- a/tests/svs/index/vamana/dynamic_index_2.cpp +++ b/tests/svs/index/vamana/dynamic_index_2.cpp @@ -38,8 +38,11 @@ // stl #include #include +#include +#include #include #include +#include using Idx = uint32_t; using Eltype = float; @@ -734,3 +737,139 @@ CATCH_TEST_CASE( CATCH_REQUIRE(index.size() == num_points + 1 - to_delete.size()); CATCH_REQUIRE(index.has_id(num_points)); } + +CATCH_TEST_CASE("Dynamic Vamana Index Predicated And Range Search", "[dynamic_index]") { + const size_t dims = 16; + const size_t num_points = 1000; + const size_t num_queries = 20; + std::mt19937 rng(4321); + std::uniform_real_distribution uniform(0, 1); + auto make_data = [&](size_t n) { + auto data = svs::data::SimpleData(n, dims); + std::vector datum(dims); + for (size_t i = 0; i < n; ++i) { + std::generate(datum.begin(), datum.end(), [&]() { return uniform(rng); }); + data.set_datum(i, datum); + } + return data; + }; + auto points = make_data(num_points); + auto queries = make_data(num_queries); + auto data = svs::data::BlockedData(num_points, dims); + for (size_t i = 0; i < num_points; ++i) { + data.set_datum(i, points.get_datum(i)); + } + std::vector ids(num_points); + std::iota(ids.begin(), ids.end(), 0); + + svs::index::vamana::VamanaBuildParameters parameters{1.2, 32, 64, 1000, 28, true}; + auto index = svs::index::vamana::MutableVamanaIndex( + parameters, std::move(data), ids, Distance(), 2 + ); + + // Brute-force distances indexed by [query][external id]. + auto distances = std::vector>(num_queries); + for (size_t q = 0; q < num_queries; ++q) { + auto& row = distances[q]; + for (size_t i = 0; i < num_points; ++i) { + row.push_back(svs::distance::compute( + Distance(), queries.get_datum(q), points.get_datum(i) + )); + } + } + + // A radius that holds about 30 neighbors per query. + float radius = 0; + for (const auto& row : distances) { + auto sorted = row; + std::nth_element(sorted.begin(), sorted.begin() + 30, sorted.end()); + radius += sorted[30]; + } + radius /= static_cast(num_queries); + + using buffer_type = typename decltype(index)::search_buffer_type; + auto sp = index.get_search_parameters(); + + CATCH_SECTION("Default stop matches search") { + sp.buffer_config({20, 40}); + auto expected = svs::QueryResult(num_queries, NUM_NEIGHBORS); + index.search(expected.view(), queries, sp); + auto actual = svs::QueryResult(num_queries, NUM_NEIGHBORS); + index.predicated_search( + actual.view(), queries, sp, svs::index::vamana::DefaultStop{} + ); + for (size_t q = 0; q < num_queries; ++q) { + for (size_t j = 0; j < NUM_NEIGHBORS; ++j) { + CATCH_REQUIRE(actual.index(q, j) == expected.index(q, j)); + CATCH_REQUIRE(actual.distance(q, j) == expected.distance(q, j)); + } + } + } + + auto check_range_results = [&](const auto& results, size_t max_results) { + CATCH_REQUIRE(results.size() == num_queries); + size_t num_expected = 0; + size_t num_found = 0; + for (size_t q = 0; q < num_queries; ++q) { + const auto& result = results[q]; + CATCH_REQUIRE(result.size() <= max_results); + auto seen = std::unordered_set{}; + float last = 0; + for (const auto& neighbor : result) { + CATCH_REQUIRE(neighbor.distance() < radius); + CATCH_REQUIRE(neighbor.distance() >= last); + CATCH_REQUIRE(index.has_id(neighbor.id())); + CATCH_REQUIRE(seen.insert(neighbor.id()).second); + last = neighbor.distance(); + } + for (size_t i = 0; i < num_points; ++i) { + if (index.has_id(i) && distances[q][i] < radius) { + ++num_expected; + num_found += seen.contains(i); + } + } + } + return static_cast(num_found) / static_cast(num_expected); + }; + + CATCH_SECTION("Unbounded range search") { + sp.buffer_config({10, buffer_type::unbounded_capacity}); + auto results = index.range_search(queries, radius, sp); + CATCH_REQUIRE(check_range_results(results, num_points) >= TARGET_RECALL); + + // A zero capacity is unbounded as well. With a zero search window the search + // stops at the radius boundary, trading recall for fewer visits. + sp.buffer_config({}); + auto zero_results = index.range_search(queries, radius, sp); + auto zero_recall = check_range_results(zero_results, num_points); + CATCH_REQUIRE(zero_recall > 0); + + // A radius beyond all points returns everything reachable. + auto all = index.range_search(queries, std::numeric_limits::max(), sp); + for (const auto& result : all) { + CATCH_REQUIRE(result.size() == num_points); + } + } + + CATCH_SECTION("Bounded range search") { + sp.buffer_config({10, 20}); + auto results = index.range_search(queries, radius, sp); + check_range_results(results, 20); + size_t num_full = 0; + for (const auto& result : results) { + num_full += result.size() == 20; + } + CATCH_REQUIRE(num_full > 0); + } + + CATCH_SECTION("Range search skips deleted points") { + std::vector to_delete; + for (size_t i = 0; i < num_points; i += 3) { + to_delete.push_back(i); + } + index.delete_entries(to_delete); + sp.buffer_config({10, buffer_type::unbounded_capacity}); + auto results = index.range_search(queries, radius, sp); + CATCH_REQUIRE(check_range_results(results, num_points) >= TARGET_RECALL); + } +} diff --git a/tests/svs/index/vamana/search_buffer.cpp b/tests/svs/index/vamana/search_buffer.cpp index 8ecb698a5..a2726e13f 100644 --- a/tests/svs/index/vamana/search_buffer.cpp +++ b/tests/svs/index/vamana/search_buffer.cpp @@ -30,6 +30,7 @@ // stdlib #include #include +#include #include namespace vamana = svs::index::vamana; @@ -118,7 +119,7 @@ template > struct SearchBufferReference { bool is_valid = neighbors_.back().valid_; if (!is_valid) { neighbors_.pop_back(); - } else if (num_valid >= valid_capacity_ + 1) { + } else if (num_valid > valid_capacity_) { neighbors_.pop_back(); --num_valid; } else { @@ -174,9 +175,24 @@ template > struct SearchBufferReference { } auto eq = svs::NeighborEqual(); CATCH_REQUIRE(valid_in_buffer.size() == valid()); - if (last) { + if (last && valid_capacity_ != std::numeric_limits::max()) { CATCH_REQUIRE(valid_in_buffer.size() == valid_capacity_); } + + // Check that the cached search state of mutable buffers is consistent. + if constexpr (requires { buffer.valid_visited(); }) { + size_t best = buffer.best_unvisited(); + CATCH_REQUIRE(best <= buffer.size()); + size_t valid_visited = 0; + for (size_t i = 0; i < best; ++i) { + CATCH_REQUIRE(buffer[i].visited()); + valid_visited += static_cast(detail::valid(buffer[i])); + } + CATCH_REQUIRE(buffer.valid_visited() == valid_visited); + if (best < buffer.size()) { + CATCH_REQUIRE(!buffer[best].visited()); + } + } CATCH_REQUIRE(valid_neighbors.size() == valid_in_buffer.size()); for (size_t i = 0; i < valid_in_buffer.size(); ++i) { @@ -930,3 +946,178 @@ CATCH_TEST_CASE("Fuzzing Mutable", "[core][search_buffer]") { CATCH_SECTION("Less") { run_test(std::less<>()); } CATCH_SECTION("Greater") { run_test(std::greater<>()); } } + +CATCH_TEST_CASE("Fuzzing Mutable Unbounded", "[core][search_buffer]") { + size_t num_trials = 5; + size_t dataset_size = 1000; + size_t seed = 0xc0ffee; + bool allow_invalid = true; + + auto run_test = [&](Cmp SVS_UNUSED(cmp)) { + using buffer_type = svs::index::vamana::MutableBuffer; + auto setup = FuzzSetup{ + num_trials, + dataset_size, + 32, + buffer_type::unbounded_capacity, + seed, + allow_invalid}; + auto buffer = buffer_type{ + svs::index::vamana::SearchBufferConfig{32, buffer_type::unbounded_capacity}}; + CATCH_REQUIRE(buffer.target_capacity() == buffer_type::unbounded_capacity); + fuzz_test(buffer, setup); + CATCH_REQUIRE(!buffer.full()); + }; + + CATCH_SECTION("Less") { run_test(std::less<>()); } + CATCH_SECTION("Greater") { run_test(std::greater<>()); } +} + +namespace { +// Stop like the default, or early once too few candidates pass the filter: +// `valid / total_candidates < min_rate`. +struct FilterRateStop { + static constexpr bool counts_total_candidates = true; + size_t min_samples; + double min_rate; + bool operator()(const svs::index::vamana::MutableBufferState& state) const { + if (svs::index::vamana::DefaultStop{}(state)) { + return true; + } + return state.total_candidates >= min_samples && + static_cast(state.valid) < + min_rate * static_cast(state.total_candidates); + } +}; +} // namespace + +CATCH_TEST_CASE("MutableBuffer Stop Predicate", "[core][search_buffer]") { + using svs::index::vamana::MutableBuffer; + using svs::index::vamana::SearchBufferConfig; + + CATCH_SECTION("Total Candidates Counter") { + static_assert(!svs::index::vamana::counts_total_candidates_v< + svs::index::vamana::DefaultStop>); + static_assert(svs::index::vamana::counts_total_candidates_v); + auto buffer = MutableBuffer, FilterRateStop>{ + SearchBufferConfig{2, 2}, std::less<>(), false, FilterRateStop{0, 0.0}}; + buffer.push_back({0, 5, true}); + buffer.push_back({1, 6, false}); + buffer.sort(); + CATCH_REQUIRE(buffer.total_candidates() == 2); + + // Candidates dropped by a full buffer count as well. + buffer.insert({2, 1, true}); + CATCH_REQUIRE(buffer.full()); + buffer.insert({3, 100, true}); + buffer.insert({4, 200, false}); + CATCH_REQUIRE(buffer.size() == 2); + CATCH_REQUIRE(buffer.valid() == 2); + CATCH_REQUIRE(buffer.total_candidates() == 5); + + // A soft clear restarts the search over the same candidates. + buffer.soft_clear(); + CATCH_REQUIRE(buffer.total_candidates() == 5); + buffer.clear(); + CATCH_REQUIRE(buffer.total_candidates() == 0); + + // Predicates that do not opt in pay nothing for the counter. + auto plain = MutableBuffer{SearchBufferConfig{2, 2}}; + plain.insert({0, 5, false}); + CATCH_REQUIRE(plain.total_candidates() == 0); + } + + CATCH_SECTION("Filter Rate Predicate") { + using buffer_type = MutableBuffer, FilterRateStop>; + auto make_buffer = [](size_t min_samples) { + return buffer_type{ + SearchBufferConfig{10, 10}, + std::less<>(), + false, + FilterRateStop{min_samples, 0.5}}; + }; + // One in five candidates passes the filter. + auto fill = [](buffer_type& buffer) { + for (uint32_t i = 0; i < 20; ++i) { + buffer.insert({i, static_cast(i), i % 5 == 0}); + } + }; + + auto buffer = make_buffer(10); + fill(buffer); + CATCH_REQUIRE(buffer.valid() == 4); + CATCH_REQUIRE(buffer.total_candidates() == 20); + CATCH_REQUIRE(buffer.done()); + + // Not enough samples to judge the filter rate yet. + auto patient = make_buffer(30); + fill(patient); + CATCH_REQUIRE(!patient.done()); + + // The predicate is preserved by shallow copies. + CATCH_REQUIRE(patient.shallow_copy().stop_predicate().min_samples == 30); + } + + CATCH_SECTION("Range Predicate") { + // Window of 2: visit at least the 2 best valid, then while within the radius. + using stop_type = svs::index::vamana::RangeStop>; + using buffer_type = MutableBuffer, stop_type>; + auto buffer = buffer_type{ + SearchBufferConfig{2, buffer_type::unbounded_capacity}, + std::less<>(), + false, + stop_type{5.5f}}; + for (uint32_t i = 0; i < 20; ++i) { + buffer.insert({i, static_cast(i), i % 2 == 0}); + } + CATCH_REQUIRE(buffer.size() == 20); + + // Candidates 0 to 5 lie within the radius. + for (uint32_t i = 0; i < 6; ++i) { + CATCH_REQUIRE(!buffer.done()); + CATCH_REQUIRE(buffer.next().id() == i); + } + CATCH_REQUIRE(buffer.done()); + CATCH_REQUIRE(buffer.valid_visited() == 3); + + // A new in-range candidate restarts the search. + buffer.insert({100, 0.5f, false}); + CATCH_REQUIRE(buffer.best_unvisited() == 1); + CATCH_REQUIRE(buffer.valid_visited() == 1); + CATCH_REQUIRE(!buffer.done()); + CATCH_REQUIRE(buffer.next().id() == 100); + CATCH_REQUIRE(buffer.valid_visited() == 3); + CATCH_REQUIRE(buffer.done()); + + // With a large radius, the window is not the limiting factor. + auto wide = buffer_type{ + SearchBufferConfig{2, buffer_type::unbounded_capacity}, + std::less<>(), + false, + stop_type{100.0f}}; + for (uint32_t i = 0; i < 20; ++i) { + wide.insert({i, static_cast(i), true}); + } + size_t count = 0; + while (!wide.done()) { + wide.next(); + ++count; + } + CATCH_REQUIRE(count == 20); + } + + CATCH_SECTION("Range Predicate Greater") { + using stop_type = svs::index::vamana::RangeStop>; + auto buffer = MutableBuffer, stop_type>{ + SearchBufferConfig{1, 10}, std::greater<>(), false, stop_type{16.5f}}; + for (uint32_t i = 0; i < 10; ++i) { + buffer.insert({i, static_cast(10 + i), true}); + } + // Similarity in [17, 19] is within the radius. + for (uint32_t i = 9; i >= 7; --i) { + CATCH_REQUIRE(!buffer.done()); + CATCH_REQUIRE(buffer.next().id() == i); + } + CATCH_REQUIRE(buffer.done()); + } +} From 1cb350dfcdd947ed5d4d1c05d24c2066c17a69b5 Mon Sep 17 00:00:00 2001 From: Rafik Saliev Date: Fri, 13 Nov 2026 03:59:43 -0700 Subject: [PATCH 3/5] Add CancelableStop predicate to join search `stop` and `cancel` predicates and update tests --- include/svs/index/vamana/dynamic_index.h | 70 ++++++++++----- .../svs/index/vamana/dynamic_search_buffer.h | 15 ++++ include/svs/index/vamana/greedy_search.h | 10 ++- tests/svs/index/vamana/dynamic_index_2.cpp | 85 +++++++++++++++++++ tests/svs/index/vamana/search_buffer.cpp | 34 ++++++++ 5 files changed, 190 insertions(+), 24 deletions(-) diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 9fe811036..62784c6f5 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -789,7 +789,8 @@ class MutableVamanaIndex { /// buffer unbounded, otherwise the capacity bounds the number of results. /// @param stop The stop predicate invoked by the search buffer with a /// ``MutableBufferState``. The search terminates when it returns ``true`` or when - /// no unvisited candidate is left. + /// no unvisited candidate is left. Wrap it in ``CancelableStop`` to support + /// cancellation, e.g. on timeout. /// template predicated_scratchspace_type @@ -808,22 +809,19 @@ class MutableVamanaIndex { /// @brief Search for a single query, terminating with the scratchspace stop predicate. /// /// After the call, ``scratch.buffer`` holds the valid results with internal IDs, sorted - /// from best to worst. + /// from best to worst. Cancellation is up to the stop predicate (see + /// ``CancelableStop``); a cancelled search keeps the best results found so far. /// template - void predicated_search( - const Query& query, - predicated_scratchspace_type& scratch, - const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) - ) const { + void predicated_search(const Query& query, predicated_scratchspace_type& scratch) + const { extensions::single_search( data_, scratch.buffer, scratch.scratch, query, - greedy_search_closure(scratch.prefetch_parameters, cancel), - *this, - cancel + stop_search_closure(scratch.prefetch_parameters), + *this ); } @@ -833,6 +831,9 @@ class MutableVamanaIndex { /// Up to ``results.n_neighbors()`` results with external IDs are stored per query. /// Missing results are padded with sentinel neighbors. /// + /// Once ``cancel`` returns ``true``, searches in flight stop at their next hop and keep + /// the best results found so far, and queries not yet started get only sentinels. + /// template void predicated_search( QueryResultView results, @@ -848,20 +849,26 @@ class MutableVamanaIndex { threadpool_, threads::StaticPartition{queries.size()}, [&](const auto is, uint64_t SVS_UNUSED(tid)) { - auto scratch = predicated_scratchspace_type{ - predicated_buffer_type{ + using stop_type = CancelableStop; + auto scratch = predicated_scratchspace_type{ + predicated_buffer_type{ predicated_buffer_config(sp.buffer_config_, num_neighbors), distance::comparator(distance_), sp.search_buffer_visited_set_, - stop}, + stop_type{stop, &cancel}}, extensions::single_search_setup(data_, distance_), {sp.prefetch_lookahead_, sp.prefetch_step_}}; + bool cancelled = false; for (auto i : is) { - if (cancel()) { - return; + cancelled = cancelled || cancel(); + if (cancelled) { + for (size_t j = 0; j < num_neighbors; ++j) { + results.set(sentinel, i, j); + } + continue; } - predicated_search(queries.get_datum(i), scratch, cancel); + predicated_search(queries.get_datum(i), scratch); const auto& buffer = scratch.buffer; const size_t found = std::min(buffer.size(), num_neighbors); @@ -895,7 +902,9 @@ class MutableVamanaIndex { /// @param radius The search radius. /// @param sp The search parameters. A zero total buffer capacity makes the number of /// results per query unbounded, otherwise the capacity bounds it. - /// @param cancel Cancellation predicate. + /// @param cancel Cancellation predicate. Once it returns ``true``, searches in flight + /// stop at their next hop and keep the in-radius results found so far, and queries + /// not yet started return no results. /// /// @returns The neighbors with external IDs for each query, sorted from best to worst. /// @@ -906,19 +915,20 @@ class MutableVamanaIndex { const search_parameters_type& sp, const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) ) { - using stop_type = RangeStop>; + using stop_type = CancelableStop>>; const auto compare = distance::comparator(distance_); auto results = std::vector>>(queries.size()); threads::parallel_for( threadpool_, threads::StaticPartition{queries.size()}, [&](const auto is, uint64_t SVS_UNUSED(tid)) { - auto scratch = predicated_scratchspace(sp, stop_type{radius, compare}); + auto scratch = + predicated_scratchspace(sp, stop_type{{radius, compare}, &cancel}); for (auto i : is) { if (cancel()) { return; } - predicated_search(queries.get_datum(i), scratch, cancel); + predicated_search(queries.get_datum(i), scratch); auto& result = results[i]; for (const auto& neighbor : scratch.buffer) { @@ -936,6 +946,26 @@ class MutableVamanaIndex { } private: + // Closure for searches cancelled through the search buffer stop predicate. + auto stop_search_closure(GreedySearchPrefetchParameters prefetch_parameters) const { + return [&, prefetch_parameters]( + const auto& query, auto& accessor, auto& distance, auto& buffer + ) { + greedy_search( + graph_, + data_, + accessor, + query, + distance, + buffer, + vamana::EntryPointInitializer{lib::as_const_span(entry_point_)}, + internal_search_builder(), + prefetch_parameters + ); + buffer.cleanup(); + }; + } + // Map search parameters to a buffer configuration for predicated search. static SearchBufferConfig predicated_buffer_config(const SearchBufferConfig& config, size_t min_capacity) { diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index 2d52544b6..7c444de86 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -19,6 +19,7 @@ #include "svs/index/vamana/search_buffer.h" #include "svs/lib/boundscheck.h" #include "svs/lib/datatype.h" +#include "svs/lib/misc.h" #include "svs/lib/neighbor.h" #include "svs/lib/threads/threadlocal.h" @@ -80,6 +81,20 @@ template > struct RangeStop { } }; +/// @brief Stop predicate that also stops when the search is cancelled (e.g. on timeout). +/// +/// The search then ends gracefully: the buffer keeps the best candidates found so far. +template struct CancelableStop { + static constexpr bool counts_total_candidates = counts_total_candidates_v; + + [[no_unique_address]] Stop stop; + const lib::DefaultPredicate* cancel; + + bool operator()(const MutableBufferState& state) const { + return stop(state) || (*cancel)(); + } +}; + /// /// A search buffer that allows entries to be predicated out. /// The search buffer will still navigate through those entries, but it won't include diff --git a/include/svs/index/vamana/greedy_search.h b/include/svs/index/vamana/greedy_search.h index f12c0129f..60fa81ed2 100644 --- a/include/svs/index/vamana/greedy_search.h +++ b/include/svs/index/vamana/greedy_search.h @@ -120,7 +120,8 @@ template < typename Buffer, typename Initializer, typename Builder, - GreedySearchTracker Tracker> + GreedySearchTracker Tracker, + typename Cancel = lib::Returns>> void greedy_search( const Graph& graph, const Dataset& dataset, @@ -132,7 +133,7 @@ void greedy_search( const Builder& builder, Tracker& search_tracker, GreedySearchPrefetchParameters prefetch_parameters = {}, - const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + const Cancel& cancel = Cancel{} ) { using I = typename Graph::index_type; @@ -212,7 +213,8 @@ template < distance::Distance Dist, typename Buffer, typename Initializer, - typename Builder = NeighborBuilder> + typename Builder = NeighborBuilder, + typename Cancel = lib::Returns>> void greedy_search( const Graph& graph, const Dataset& dataset, @@ -223,7 +225,7 @@ void greedy_search( const Initializer& initializer, const Builder& builder = NeighborBuilder(), GreedySearchPrefetchParameters prefetch_parameters = {}, - const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + const Cancel& cancel = Cancel{} ) { auto null_tracker = NullTracker{}; greedy_search( diff --git a/tests/svs/index/vamana/dynamic_index_2.cpp b/tests/svs/index/vamana/dynamic_index_2.cpp index 0e58e3327..b6f00cd54 100644 --- a/tests/svs/index/vamana/dynamic_index_2.cpp +++ b/tests/svs/index/vamana/dynamic_index_2.cpp @@ -37,6 +37,7 @@ // stl #include +#include #include #include #include @@ -862,6 +863,90 @@ CATCH_TEST_CASE("Dynamic Vamana Index Predicated And Range Search", "[dynamic_in CATCH_REQUIRE(num_full > 0); } + CATCH_SECTION("Cancel predicated search") { + sp.buffer_config({20, 40}); + constexpr auto sentinel = + svs::type_traits::sentinel_v, std::less<>>; + + // Cancellation is checked once per hop and once per query, two fewer per query + // (and one fewer per batch) than in the regular search. + std::atomic regular_calls{0}; + auto expected = svs::QueryResult(num_queries, NUM_NEIGHBORS); + index.search(expected.view(), queries, sp, [&]() { + ++regular_calls; + return false; + }); + std::atomic calls{0}; + auto actual = svs::QueryResult(num_queries, NUM_NEIGHBORS); + index.predicated_search( + actual.view(), + queries, + sp, + svs::index::vamana::DefaultStop{}, + [&]() { + ++calls; + return false; + } + ); + CATCH_REQUIRE(calls + num_queries + 1 == regular_calls); + for (size_t q = 0; q < num_queries; ++q) { + for (size_t j = 0; j < NUM_NEIGHBORS; ++j) { + CATCH_REQUIRE(actual.index(q, j) == expected.index(q, j)); + } + } + + // Cancel midway: started queries keep sorted partial results, the rest are + // sentinels. + calls = 0; + auto partial = svs::QueryResult(num_queries, NUM_NEIGHBORS); + index.predicated_search( + partial.view(), + queries, + sp, + svs::index::vamana::DefaultStop{}, + [&]() { return ++calls > 50; } + ); + size_t num_empty = 0; + for (size_t q = 0; q < num_queries; ++q) { + bool padding = false; + float last = 0; + for (size_t j = 0; j < NUM_NEIGHBORS; ++j) { + auto id = partial.index(q, j); + if (id == sentinel.id()) { + CATCH_REQUIRE(partial.distance(q, j) == sentinel.distance()); + padding = true; + continue; + } + CATCH_REQUIRE(!padding); + CATCH_REQUIRE(index.has_id(id)); + CATCH_REQUIRE(partial.distance(q, j) >= last); + last = partial.distance(q, j); + } + num_empty += partial.index(q, 0) == sentinel.id(); + } + CATCH_REQUIRE(num_empty > 0); + CATCH_REQUIRE(num_empty < num_queries); + } + + CATCH_SECTION("Cancel range search") { + sp.buffer_config({10, buffer_type::unbounded_capacity}); + auto none = index.range_search(queries, radius, sp, []() { return true; }); + for (const auto& result : none) { + CATCH_REQUIRE(result.empty()); + } + + std::atomic calls{0}; + auto partial = + index.range_search(queries, radius, sp, [&]() { return ++calls > 50; }); + check_range_results(partial, num_points); + size_t num_empty = 0; + for (const auto& result : partial) { + num_empty += result.empty(); + } + CATCH_REQUIRE(num_empty > 0); + CATCH_REQUIRE(num_empty < num_queries); + } + CATCH_SECTION("Range search skips deleted points") { std::vector to_delete; for (size_t i = 0; i < num_points; i += 3) { diff --git a/tests/svs/index/vamana/search_buffer.cpp b/tests/svs/index/vamana/search_buffer.cpp index a2726e13f..51ed5c387 100644 --- a/tests/svs/index/vamana/search_buffer.cpp +++ b/tests/svs/index/vamana/search_buffer.cpp @@ -1058,6 +1058,40 @@ CATCH_TEST_CASE("MutableBuffer Stop Predicate", "[core][search_buffer]") { CATCH_REQUIRE(patient.shallow_copy().stop_predicate().min_samples == 30); } + CATCH_SECTION("Cancelable Predicate") { + using svs::index::vamana::CancelableStop; + using svs::index::vamana::counts_total_candidates_v; + static_assert(counts_total_candidates_v>); + static_assert(!counts_total_candidates_v< + CancelableStop>); + + bool cancelled = false; + size_t calls = 0; + auto cancel = svs::lib::DefaultPredicate{[&]() { + ++calls; + return cancelled; + }}; + using stop_type = CancelableStop; + auto buffer = MutableBuffer, stop_type>{ + SearchBufferConfig{2, 2}, std::less<>(), false, stop_type{{}, &cancel}}; + for (uint32_t i = 0; i < 4; ++i) { + buffer.insert({i, static_cast(i), i != 0}); + } + CATCH_REQUIRE(!buffer.done()); + CATCH_REQUIRE(calls == 1); + buffer.next(); + cancelled = true; + CATCH_REQUIRE(buffer.done()); + + // The wrapped predicate takes precedence without consulting `cancel`. + cancelled = false; + buffer.next(); + buffer.next(); + calls = 0; + CATCH_REQUIRE(buffer.done()); + CATCH_REQUIRE(calls == 0); + } + CATCH_SECTION("Range Predicate") { // Window of 2: visit at least the 2 best valid, then while within the radius. using stop_type = svs::index::vamana::RangeStop>; From 0ddea3fb92785fae18a1d20a7a5f558334bb96f4 Mon Sep 17 00:00:00 2001 From: Rafik Saliev Date: Fri, 13 Nov 2026 05:02:25 -0700 Subject: [PATCH 4/5] Add predicated search and range search functionality to static VamanaIndex and related tests --- include/svs/index/vamana/dynamic_index.h | 100 ++--------- include/svs/index/vamana/index.h | 142 ++++++++++++++++ include/svs/index/vamana/predicated_search.h | 165 +++++++++++++++++++ tests/svs/index/vamana/index.cpp | 129 +++++++++++++++ 4 files changed, 454 insertions(+), 82 deletions(-) create mode 100644 include/svs/index/vamana/predicated_search.h diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 62784c6f5..410f25eb4 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -39,6 +39,7 @@ #include "svs/index/vamana/dynamic_search_buffer.h" #include "svs/index/vamana/greedy_search.h" #include "svs/index/vamana/index.h" +#include "svs/index/vamana/predicated_search.h" #include "svs/index/vamana/vamana_build.h" #include "svs/lib/boundscheck.h" #include "svs/lib/preprocessor.h" @@ -797,7 +798,7 @@ class MutableVamanaIndex { predicated_scratchspace(const search_parameters_type& sp, Stop stop) const { return predicated_scratchspace_type{ predicated_buffer_type{ - predicated_buffer_config(sp.buffer_config_, 0), + predicated_buffer_config(sp.buffer_config_), distance::comparator(distance_), sp.search_buffer_visited_set_, std::move(stop)}, @@ -842,51 +843,15 @@ class MutableVamanaIndex { const Stop& stop, const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) ) { - using compare_type = distance::compare_t; - constexpr auto sentinel = type_traits::sentinel_v, compare_type>; - const size_t num_neighbors = results.n_neighbors(); - threads::parallel_for( + detail::predicated_batch_search( + *this, threadpool_, - threads::StaticPartition{queries.size()}, - [&](const auto is, uint64_t SVS_UNUSED(tid)) { - using stop_type = CancelableStop; - auto scratch = predicated_scratchspace_type{ - predicated_buffer_type{ - predicated_buffer_config(sp.buffer_config_, num_neighbors), - distance::comparator(distance_), - sp.search_buffer_visited_set_, - stop_type{stop, &cancel}}, - extensions::single_search_setup(data_, distance_), - {sp.prefetch_lookahead_, sp.prefetch_step_}}; - - bool cancelled = false; - for (auto i : is) { - cancelled = cancelled || cancel(); - if (cancelled) { - for (size_t j = 0; j < num_neighbors; ++j) { - results.set(sentinel, i, j); - } - continue; - } - predicated_search(queries.get_datum(i), scratch); - - const auto& buffer = scratch.buffer; - const size_t found = std::min(buffer.size(), num_neighbors); - for (size_t j = 0; j < found; ++j) { - const auto& neighbor = buffer[j]; - results.set( - Neighbor{ - lib::narrow_cast(translate_internal_id(neighbor.id())), - neighbor.distance()}, - i, - j - ); - } - for (size_t j = found; j < num_neighbors; ++j) { - results.set(sentinel, i, j); - } - } - } + results, + queries, + sp, + stop, + cancel, + [&](Idx id) { return translate_internal_id(id); } ); } @@ -915,34 +880,16 @@ class MutableVamanaIndex { const search_parameters_type& sp, const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) ) { - using stop_type = CancelableStop>>; - const auto compare = distance::comparator(distance_); - auto results = std::vector>>(queries.size()); - threads::parallel_for( + return detail::range_batch_search( + *this, threadpool_, - threads::StaticPartition{queries.size()}, - [&](const auto is, uint64_t SVS_UNUSED(tid)) { - auto scratch = - predicated_scratchspace(sp, stop_type{{radius, compare}, &cancel}); - for (auto i : is) { - if (cancel()) { - return; - } - predicated_search(queries.get_datum(i), scratch); - - auto& result = results[i]; - for (const auto& neighbor : scratch.buffer) { - if (!compare(neighbor.distance(), radius)) { - break; - } - result.emplace_back( - translate_internal_id(neighbor.id()), neighbor.distance() - ); - } - } - } + queries, + radius, + sp, + cancel, + distance::comparator(distance_), + [&](Idx id) { return translate_internal_id(id); } ); - return results; } private: @@ -966,17 +913,6 @@ class MutableVamanaIndex { }; } - // Map search parameters to a buffer configuration for predicated search. - static SearchBufferConfig - predicated_buffer_config(const SearchBufferConfig& config, size_t min_capacity) { - const size_t window = config.get_search_window_size(); - const size_t capacity = config.get_total_capacity(); - if (capacity == 0) { - return SearchBufferConfig{window, search_buffer_type::unbounded_capacity}; - } - return SearchBufferConfig{window, std::max(capacity, min_capacity)}; - } - public: /// /// @brief Return a unique instance of the distance function. diff --git a/include/svs/index/vamana/index.h b/include/svs/index/vamana/index.h index 4d4a47f09..d5dfd1aca 100644 --- a/include/svs/index/vamana/index.h +++ b/include/svs/index/vamana/index.h @@ -26,6 +26,7 @@ #include "svs/index/vamana/calibrate.h" #include "svs/index/vamana/extensions.h" #include "svs/index/vamana/greedy_search.h" +#include "svs/index/vamana/predicated_search.h" #include "svs/index/vamana/search_buffer.h" #include "svs/index/vamana/search_params.h" #include "svs/index/vamana/vamana_build.h" @@ -620,6 +621,147 @@ class VamanaIndex { ); } + ///// Predicated Search + + /// The search buffer type terminating with the stop predicate `Stop`. + template + using predicated_buffer_type = MutableBuffer, Stop>; + + /// The scratchspace type for ``predicated_search`` with the stop predicate `Stop`. + template + using predicated_scratchspace_type = + SearchScratchspace, inner_scratch_type>; + + /// + /// @brief Return a scratchspace for ``predicated_search``. + /// + /// @param sp The search parameters. A zero total buffer capacity makes the search + /// buffer unbounded, otherwise the capacity bounds the number of results. + /// @param stop The stop predicate invoked by the search buffer with a + /// ``MutableBufferState``. The search terminates when it returns ``true`` or when + /// no unvisited candidate is left. Wrap it in ``CancelableStop`` to support + /// cancellation, e.g. on timeout. + /// + template + predicated_scratchspace_type + predicated_scratchspace(const search_parameters_type& sp, Stop stop) const { + return predicated_scratchspace_type{ + predicated_buffer_type{ + predicated_buffer_config(sp.buffer_config_), + distance::comparator(distance_), + sp.search_buffer_visited_set_, + std::move(stop)}, + extensions::single_search_setup(data_, distance_), + {sp.prefetch_lookahead_, sp.prefetch_step_}}; + } + + /// + /// @brief Search for a single query, terminating with the scratchspace stop predicate. + /// + /// After the call, ``scratch.buffer`` holds the results sorted from best to worst. + /// Cancellation is up to the stop predicate (see ``CancelableStop``); a cancelled + /// search keeps the best results found so far. + /// + template + void predicated_search(const Query& query, predicated_scratchspace_type& scratch) + const { + extensions::single_search( + data_, + scratch.buffer, + scratch.scratch, + query, + stop_search_closure(scratch.prefetch_parameters), + *this + ); + } + + /// + /// @brief Batch search terminating with the stop predicate `stop`. + /// + /// Up to ``results.n_neighbors()`` results are stored per query. Missing results are + /// padded with sentinel neighbors. + /// + /// Once ``cancel`` returns ``true``, searches in flight stop at their next hop and keep + /// the best results found so far, and queries not yet started get only sentinels. + /// + template + void predicated_search( + QueryResultView results, + const Queries& queries, + const search_parameters_type& sp, + const Stop& stop, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + ) { + detail::predicated_batch_search( + *this, + threadpool_, + results, + queries, + sp, + stop, + cancel, + [](Idx id) { return id; } + ); + } + + /// + /// @brief Return all neighbors within ``radius`` of each query. + /// + /// A neighbor at distance ``d`` is within the radius if ``compare(d, radius)`` for the + /// index comparator (e.g., ``d < radius`` for L2, ``d > radius`` for inner product). + /// The search visits at least the search window, then continues while the best + /// unvisited candidate is within the radius. + /// + /// @param queries The queries. + /// @param radius The search radius. + /// @param sp The search parameters. A zero total buffer capacity makes the number of + /// results per query unbounded, otherwise the capacity bounds it. + /// @param cancel Cancellation predicate. Once it returns ``true``, searches in flight + /// stop at their next hop and keep the in-radius results found so far, and queries + /// not yet started return no results. + /// + /// @returns The neighbors for each query, sorted from best to worst. + /// + template + std::vector>> range_search( + const Queries& queries, + float radius, + const search_parameters_type& sp, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + ) { + return detail::range_batch_search( + *this, + threadpool_, + queries, + radius, + sp, + cancel, + distance::comparator(distance_), + [](Idx id) { return id; } + ); + } + + private: + // Closure for searches cancelled through the search buffer stop predicate. + auto stop_search_closure(GreedySearchPrefetchParameters prefetch_parameters) const { + return [&, prefetch_parameters]( + const auto& query, auto& accessor, auto& distance, auto& buffer + ) { + greedy_search( + graph_, + data_, + accessor, + query, + distance, + buffer, + vamana::EntryPointInitializer{lib::as_const_span(entry_point_)}, + PredicatedNeighborBuilder(), + prefetch_parameters + ); + }; + } + + public: // TODO (Mark): Make descriptions better. std::string name() const { return "VamanaIndex"; } diff --git a/include/svs/index/vamana/predicated_search.h b/include/svs/index/vamana/predicated_search.h new file mode 100644 index 000000000..1c14f362b --- /dev/null +++ b/include/svs/index/vamana/predicated_search.h @@ -0,0 +1,165 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include "svs/concepts/data.h" +#include "svs/core/distance.h" +#include "svs/core/query_result.h" +#include "svs/index/vamana/dynamic_search_buffer.h" +#include "svs/index/vamana/search_params.h" +#include "svs/lib/misc.h" +#include "svs/lib/neighbor.h" +#include "svs/lib/threads.h" +#include "svs/lib/type_traits.h" + +#include +#include + +namespace svs::index::vamana { + +/// @brief Builder marking every candidate valid, for predicated search without a filter. +struct PredicatedNeighborBuilder { + template + constexpr PredicatedSearchNeighbor operator()(I i, float distance) const { + return PredicatedSearchNeighbor{i, distance}; + } +}; + +/// +/// @brief Return the search buffer configuration for predicated search. +/// +/// A zero total capacity makes the buffer unbounded. Otherwise, the capacity is raised to +/// at least `min_capacity`. +/// +inline SearchBufferConfig +predicated_buffer_config(const SearchBufferConfig& config, size_t min_capacity = 0) { + const size_t window = config.get_search_window_size(); + const size_t capacity = config.get_total_capacity(); + if (capacity == 0) { + return SearchBufferConfig{window, MutableBuffer::unbounded_capacity}; + } + return SearchBufferConfig{window, std::max(capacity, min_capacity)}; +} + +namespace detail { + +// Batch `predicated_search` shared by the Vamana indexes. +// `translate` maps internal IDs of the search buffer to external IDs. +template < + typename Index, + typename Pool, + std::integral I, + data::ImmutableMemoryDataset Queries, + typename Stop, + typename Translate> +void predicated_batch_search( + const Index& index, + Pool& threadpool, + QueryResultView results, + const Queries& queries, + VamanaSearchParameters sp, + const Stop& stop, + const lib::DefaultPredicate& cancel, + const Translate& translate +) { + using compare_type = distance::compare_t; + constexpr auto sentinel = type_traits::sentinel_v, compare_type>; + const size_t num_neighbors = results.n_neighbors(); + sp.buffer_config(predicated_buffer_config(sp.buffer_config_, num_neighbors)); + threads::parallel_for( + threadpool, + threads::StaticPartition{queries.size()}, + [&](const auto is, uint64_t SVS_UNUSED(tid)) { + auto scratch = + index.predicated_scratchspace(sp, CancelableStop{stop, &cancel}); + bool cancelled = false; + for (auto i : is) { + cancelled = cancelled || cancel(); + if (cancelled) { + for (size_t j = 0; j < num_neighbors; ++j) { + results.set(sentinel, i, j); + } + continue; + } + index.predicated_search(queries.get_datum(i), scratch); + + const auto& buffer = scratch.buffer; + const size_t found = std::min(buffer.size(), num_neighbors); + for (size_t j = 0; j < found; ++j) { + const auto& neighbor = buffer[j]; + results.set( + Neighbor{ + lib::narrow_cast(translate(neighbor.id())), + neighbor.distance()}, + i, + j + ); + } + for (size_t j = found; j < num_neighbors; ++j) { + results.set(sentinel, i, j); + } + } + } + ); +} + +// Batch `range_search` shared by the Vamana indexes. +template < + typename Index, + typename Pool, + data::ImmutableMemoryDataset Queries, + typename Cmp, + typename Translate> +std::vector>> range_batch_search( + const Index& index, + Pool& threadpool, + const Queries& queries, + float radius, + const VamanaSearchParameters& sp, + const lib::DefaultPredicate& cancel, + const Cmp& compare, + const Translate& translate +) { + using stop_type = CancelableStop>; + auto results = std::vector>>(queries.size()); + threads::parallel_for( + threadpool, + threads::StaticPartition{queries.size()}, + [&](const auto is, uint64_t SVS_UNUSED(tid)) { + auto scratch = + index.predicated_scratchspace(sp, stop_type{{radius, compare}, &cancel}); + for (auto i : is) { + if (cancel()) { + return; + } + index.predicated_search(queries.get_datum(i), scratch); + + auto& result = results[i]; + for (const auto& neighbor : scratch.buffer) { + if (!compare(neighbor.distance(), radius)) { + break; + } + result.emplace_back(translate(neighbor.id()), neighbor.distance()); + } + } + } + ); + return results; +} + +} // namespace detail +} // namespace svs::index::vamana diff --git a/tests/svs/index/vamana/index.cpp b/tests/svs/index/vamana/index.cpp index 3cd7ea211..a88abd920 100644 --- a/tests/svs/index/vamana/index.cpp +++ b/tests/svs/index/vamana/index.cpp @@ -41,8 +41,15 @@ // svsbenchmark #include "svs-benchmark/benchmark.h" // stl +#include +#include +#include +#include +#include #include #include +#include +#include namespace { @@ -847,3 +854,125 @@ CATCH_TEST_CASE("Vamana Index Bugcheck and Regression", "[vamana][index][regress CATCH_CHECK(build_and_search(64, 3072)); } } + +CATCH_TEST_CASE("Vamana Index Predicated And Range Search", "[vamana][index]") { + const size_t dims = 16; + const size_t num_points = 1000; + const size_t num_queries = 20; + const size_t k = 10; + std::mt19937 rng(4321); + std::normal_distribution normal(0, 1); + // Unit vectors: graphs built for inner product reach unnormalized points poorly. + auto make_data = [&](size_t n) { + auto data = svs::data::SimpleData(n, dims); + std::vector datum(dims); + for (size_t i = 0; i < n; ++i) { + std::generate(datum.begin(), datum.end(), [&]() { return normal(rng); }); + float norm = std::sqrt( + std::inner_product(datum.begin(), datum.end(), datum.begin(), 0.0f) + ); + for (auto& x : datum) { + x /= norm; + } + data.set_datum(i, datum); + } + return data; + }; + const auto points = make_data(num_points); + const auto queries = make_data(num_queries); + auto build_params = svs::index::vamana::VamanaBuildParameters{}; + build_params.graph_max_degree = 32; + build_params.window_size = 64; + + auto run_test = [&](Dist distance) { + auto index = svs::index::vamana::auto_build(build_params, points, distance, 2); + const auto compare = svs::distance::comparator(distance); + using buffer_type = + svs::index::vamana::MutableBuffer>; + + // Brute-force distances and a radius that holds about 30 neighbors per query. + auto distances = std::vector>(num_queries); + float radius = 0; + for (size_t q = 0; q < num_queries; ++q) { + for (size_t i = 0; i < num_points; ++i) { + distances[q].push_back(svs::distance::compute( + distance, queries.get_datum(q), points.get_datum(i) + )); + } + auto sorted = distances[q]; + std::nth_element(sorted.begin(), sorted.begin() + 30, sorted.end(), compare); + radius += sorted[30]; + } + radius /= static_cast(num_queries); + + auto sp = index.get_search_parameters(); + + auto check_range_results = [&](const auto& results, size_t max_results) { + CATCH_REQUIRE(results.size() == num_queries); + size_t num_expected = 0; + size_t num_found = 0; + for (size_t q = 0; q < num_queries; ++q) { + const auto& result = results[q]; + CATCH_REQUIRE(result.size() <= max_results); + auto seen = std::unordered_set{}; + for (size_t j = 0; j < result.size(); ++j) { + CATCH_REQUIRE(compare(result[j].distance(), radius)); + if (j > 0) { + CATCH_REQUIRE( + !compare(result[j].distance(), result[j - 1].distance()) + ); + } + CATCH_REQUIRE(seen.insert(result[j].id()).second); + } + for (size_t i = 0; i < num_points; ++i) { + if (compare(distances[q][i], radius)) { + ++num_expected; + num_found += seen.contains(i); + } + } + } + return static_cast(num_found) / static_cast(num_expected); + }; + + // The default stop reproduces the regular search. + sp.buffer_config({20, 40}); + auto expected = svs::QueryResult(num_queries, k); + index.search(expected.view(), queries, sp); + auto actual = svs::QueryResult(num_queries, k); + index.predicated_search( + actual.view(), queries, sp, svs::index::vamana::DefaultStop{} + ); + for (size_t q = 0; q < num_queries; ++q) { + for (size_t j = 0; j < k; ++j) { + CATCH_REQUIRE(actual.index(q, j) == expected.index(q, j)); + CATCH_REQUIRE(actual.distance(q, j) == expected.distance(q, j)); + } + } + + // Unbounded range search. + sp.buffer_config({10, buffer_type::unbounded_capacity}); + auto results = index.range_search(queries, radius, sp); + CATCH_REQUIRE(check_range_results(results, num_points) >= 0.95); + + // Bounded range search. + sp.buffer_config({10, 20}); + auto bounded = index.range_search(queries, radius, sp); + check_range_results(bounded, 20); + CATCH_REQUIRE(std::any_of(bounded.begin(), bounded.end(), [](const auto& r) { + return r.size() == 20; + })); + + // Cancellation keeps partial results and leaves unstarted queries empty. + sp.buffer_config({10, buffer_type::unbounded_capacity}); + std::atomic calls{0}; + auto partial = + index.range_search(queries, radius, sp, [&]() { return ++calls > 50; }); + check_range_results(partial, num_points); + CATCH_REQUIRE(std::any_of(partial.begin(), partial.end(), [](const auto& r) { + return r.empty(); + })); + }; + + CATCH_SECTION("L2") { run_test(svs::distance::DistanceL2()); } + CATCH_SECTION("Inner Product") { run_test(svs::distance::DistanceIP()); } +} From 0a829c33cd6b87ea514ea03fec80039c23da5fdb Mon Sep 17 00:00:00 2001 From: Rafik Saliev Date: Fri, 13 Nov 2026 06:33:06 -0700 Subject: [PATCH 5/5] Implement predicated search and range search in concurrent::MutableVamanaIndex with corresponding tests --- include/svs/concurrent/dynamic_index.h | 194 +++++++++++++++++++ include/svs/concurrent/greedy_search.h | 10 +- include/svs/index/vamana/predicated_search.h | 22 ++- tests/svs/concurrent/dynamic_index_2.cpp | 182 +++++++++++++++++ 4 files changed, 399 insertions(+), 9 deletions(-) diff --git a/include/svs/concurrent/dynamic_index.h b/include/svs/concurrent/dynamic_index.h index 1fc1a8980..7e96e65dc 100644 --- a/include/svs/concurrent/dynamic_index.h +++ b/include/svs/concurrent/dynamic_index.h @@ -47,6 +47,7 @@ #include "svs/core/translation.h" #include "svs/index/vamana/dynamic_index.h" #include "svs/index/vamana/index.h" +#include "svs/index/vamana/predicated_search.h" #include "svs/lib/boundscheck.h" #include "svs/lib/neighbor.h" #include "svs/lib/preprocessor.h" @@ -1007,6 +1008,199 @@ class MutableVamanaIndex { translate_to_external(results.indices()); } + ///// Predicated Search + + /// The search buffer type terminating with the stop predicate `Stop`. + template + using predicated_buffer_type = + vamana::MutableBuffer, Stop>; + + /// The scratchspace type for ``predicated_search`` with the stop predicate `Stop`. + template + using predicated_scratchspace_type = + SearchScratchspace, inner_scratch_type>; + + /// + /// @brief Return a scratchspace for ``predicated_search``. + /// + /// @param sp The search parameters. A zero total buffer capacity makes the search + /// buffer unbounded, otherwise the capacity bounds the number of results. + /// @param stop The stop predicate invoked by the search buffer with a + /// ``MutableBufferState``. The search terminates when it returns ``true`` or when + /// no unvisited candidate is left. Wrap it in ``CancelableStop`` to support + /// cancellation, e.g. on timeout. + /// + template + predicated_scratchspace_type + predicated_scratchspace(const search_parameters_type& sp, Stop stop) const { + return predicated_scratchspace_type{ + predicated_buffer_type{ + predicated_buffer_config(sp.buffer_config_), + distance::comparator(distance_), + sp.search_buffer_visited_set_, + std::move(stop)}, + extensions::single_search_setup(data_, distance_), + {sp.prefetch_lookahead_, sp.prefetch_step_}}; + } + + /// + /// @brief Search for a single query, terminating with the scratchspace stop predicate. + /// + /// After the call, ``scratch.buffer`` holds the valid results with internal IDs, sorted + /// from best to worst. Consuming those IDs requires external exclusion of compaction, + /// as for the single-query ``search``. Cancellation is up to the stop predicate (see + /// ``CancelableStop``); a cancelled search keeps the best results found so far. + /// + template + void predicated_search(const Query& query, predicated_scratchspace_type& scratch) + const { + auto compact_lock = lock_storage_shared(); + unsafe_predicated_search(query, scratch); + } + + /// + /// @brief Batch search terminating with the stop predicate `stop`. + /// + /// Up to ``results.n_neighbors()`` results with external IDs are stored per query. + /// Missing results are padded with sentinel neighbors. As for ``search``, a result + /// erased by a concurrent consolidation has the ID ``no_external_id``. + /// + /// Once ``cancel`` returns ``true``, searches in flight stop at their next hop and keep + /// the best results found so far, and queries not yet started get only sentinels. + /// + template + void predicated_search( + QueryResultView results, + const Queries& queries, + const search_parameters_type& sp, + const Stop& stop, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + ) { + // Internal IDs must retain their meaning until translation has completed. + auto compact_lock = lock_storage_shared(); + vamana::detail::predicated_batch_search( + LockedSearch{*this}, + threadpool_, + results, + queries, + sp, + stop, + cancel, + [&](Idx id) { return unsafe_translate_internal_id(id); }, + [&]() { return lock_for_translation(); } + ); + } + + /// + /// @brief Return all neighbors within ``radius`` of each query. + /// + /// A neighbor at distance ``d`` is within the radius if ``compare(d, radius)`` for the + /// index comparator (e.g., ``d < radius`` for L2, ``d > radius`` for inner product). + /// The search visits at least the search window, then continues while the best + /// unvisited candidate is within the radius. Results erased by a concurrent + /// consolidation are left out. + /// + /// @param queries The queries. + /// @param radius The search radius. + /// @param sp The search parameters. A zero total buffer capacity makes the number of + /// results per query unbounded, otherwise the capacity bounds it. + /// @param cancel Cancellation predicate. Once it returns ``true``, searches in flight + /// stop at their next hop and keep the in-radius results found so far, and queries + /// not yet started return no results. + /// + /// @returns The neighbors with external IDs for each query, sorted from best to worst. + /// + template + std::vector>> range_search( + const Queries& queries, + float radius, + const search_parameters_type& sp, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + ) { + auto results = [&]() { + // Internal IDs must retain their meaning until translation has completed. + auto compact_lock = lock_storage_shared(); + return vamana::detail::range_batch_search( + LockedSearch{*this}, + threadpool_, + queries, + radius, + sp, + cancel, + distance::comparator(distance_), + [&](Idx id) { return unsafe_translate_internal_id(id); }, + [&]() { return lock_for_translation(); } + ); + }(); + for (auto& result : results) { + std::erase_if(result, [](const auto& neighbor) { + return neighbor.id() == no_external_id; + }); + } + return results; + } + + private: + // Predicated search for callers already holding `lock_storage_shared()`. + template + void unsafe_predicated_search( + const Query& query, predicated_scratchspace_type& scratch + ) const { + extensions::single_search( + data_, + scratch.buffer, + scratch.scratch, + query, + stop_search_closure(scratch.prefetch_parameters), + *this + ); + } + + // Closure for searches cancelled through the search buffer stop predicate. + auto stop_search_closure(GreedySearchPrefetchParameters prefetch_parameters) const { + return [&, prefetch_parameters]( + const auto& query, auto& accessor, auto& distance, auto& buffer + ) { + const Idx seed = entry_point(); + if (seed == NO_ENTRY) { + buffer.clear(); + return; + } + concurrent::greedy_search( + graph_, + data_, + accessor, + query, + distance, + buffer, + vamana::EntryPointInitializer{std::span{&seed, 1}}, + internal_search_builder(), + prefetch_parameters + ); + buffer.cleanup(); + }; + } + + // The index as seen by the shared batch helpers, which run under + // `lock_storage_shared()`: the per-query search must not take it again. + struct LockedSearch { + using distance_type = Dist; + const MutableVamanaIndex& index; + + template + auto predicated_scratchspace(const search_parameters_type& sp, Stop stop) const { + return index.predicated_scratchspace(sp, std::move(stop)); + } + + template + void predicated_search( + const Query& query, predicated_scratchspace_type& scratch + ) const { + index.unsafe_predicated_search(query, scratch); + } + }; + + public: /// /// @brief Return a unique instance of the distance function. /// diff --git a/include/svs/concurrent/greedy_search.h b/include/svs/concurrent/greedy_search.h index fe8d7fd59..a33c6cefd 100644 --- a/include/svs/concurrent/greedy_search.h +++ b/include/svs/concurrent/greedy_search.h @@ -59,7 +59,8 @@ template < typename Buffer, typename Initializer, typename Builder, - GreedySearchTracker Tracker> + GreedySearchTracker Tracker, + typename Cancel = lib::Returns>> void greedy_search( const Graph& graph, const Dataset& dataset, @@ -71,7 +72,7 @@ void greedy_search( const Builder& builder, Tracker& search_tracker, GreedySearchPrefetchParameters prefetch_parameters = {}, - const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + const Cancel& cancel = Cancel{} ) { using I = typename Graph::index_type; @@ -170,7 +171,8 @@ template < distance::Distance Dist, typename Buffer, typename Initializer, - typename Builder = NeighborBuilder> + typename Builder = NeighborBuilder, + typename Cancel = lib::Returns>> void greedy_search( const Graph& graph, const Dataset& dataset, @@ -181,7 +183,7 @@ void greedy_search( const Initializer& initializer, const Builder& builder = NeighborBuilder(), GreedySearchPrefetchParameters prefetch_parameters = {}, - const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + const Cancel& cancel = Cancel{} ) { auto null_tracker = NullTracker{}; concurrent::greedy_search( diff --git a/include/svs/index/vamana/predicated_search.h b/include/svs/index/vamana/predicated_search.h index 1c14f362b..233c9cbeb 100644 --- a/include/svs/index/vamana/predicated_search.h +++ b/include/svs/index/vamana/predicated_search.h @@ -57,15 +57,22 @@ predicated_buffer_config(const SearchBufferConfig& config, size_t min_capacity = namespace detail { +// Default for indexes whose ID translation needs no lock. +struct NoTranslationLock { + constexpr int operator()() const { return 0; } +}; + // Batch `predicated_search` shared by the Vamana indexes. -// `translate` maps internal IDs of the search buffer to external IDs. +// `translate` maps internal IDs of the search buffer to external IDs while the guard +// returned by `lock_translation` is alive. template < typename Index, typename Pool, std::integral I, data::ImmutableMemoryDataset Queries, typename Stop, - typename Translate> + typename Translate, + typename TranslationLock = NoTranslationLock> void predicated_batch_search( const Index& index, Pool& threadpool, @@ -74,7 +81,8 @@ void predicated_batch_search( VamanaSearchParameters sp, const Stop& stop, const lib::DefaultPredicate& cancel, - const Translate& translate + const Translate& translate, + const TranslationLock& lock_translation = {} ) { using compare_type = distance::compare_t; constexpr auto sentinel = type_traits::sentinel_v, compare_type>; @@ -99,6 +107,7 @@ void predicated_batch_search( const auto& buffer = scratch.buffer; const size_t found = std::min(buffer.size(), num_neighbors); + [[maybe_unused]] auto guard = lock_translation(); for (size_t j = 0; j < found; ++j) { const auto& neighbor = buffer[j]; results.set( @@ -123,7 +132,8 @@ template < typename Pool, data::ImmutableMemoryDataset Queries, typename Cmp, - typename Translate> + typename Translate, + typename TranslationLock = NoTranslationLock> std::vector>> range_batch_search( const Index& index, Pool& threadpool, @@ -132,7 +142,8 @@ std::vector>> range_batch_search( const VamanaSearchParameters& sp, const lib::DefaultPredicate& cancel, const Cmp& compare, - const Translate& translate + const Translate& translate, + const TranslationLock& lock_translation = {} ) { using stop_type = CancelableStop>; auto results = std::vector>>(queries.size()); @@ -149,6 +160,7 @@ std::vector>> range_batch_search( index.predicated_search(queries.get_datum(i), scratch); auto& result = results[i]; + [[maybe_unused]] auto guard = lock_translation(); for (const auto& neighbor : scratch.buffer) { if (!compare(neighbor.distance(), radius)) { break; diff --git a/tests/svs/concurrent/dynamic_index_2.cpp b/tests/svs/concurrent/dynamic_index_2.cpp index 044514bba..7160b0bc7 100644 --- a/tests/svs/concurrent/dynamic_index_2.cpp +++ b/tests/svs/concurrent/dynamic_index_2.cpp @@ -36,9 +36,13 @@ // stl #include +#include #include +#include #include #include +#include +#include // The concurrent index is a separate type from `cc::MutableVamanaIndex`; // the pre-existing static and dynamic indexes are untouched by these tests. @@ -725,3 +729,181 @@ CATCH_TEST_CASE( ); } } + +CATCH_TEST_CASE( + "Concurrent Vamana Index Predicated And Range Search", "[concurrent][vamana]" +) { + const size_t dims = 16; + const size_t num_points = 1000; + const size_t num_queries = 20; + const size_t k = 10; + std::mt19937 rng(4321); + std::uniform_real_distribution uniform(0, 1); + auto make_data = [&](size_t n) { + auto data = svs::data::SimpleData(n, dims); + std::vector datum(dims); + for (size_t i = 0; i < n; ++i) { + std::generate(datum.begin(), datum.end(), [&]() { return uniform(rng); }); + data.set_datum(i, datum); + } + return data; + }; + const auto points = make_data(num_points); + const auto queries = make_data(num_queries); + auto data = svs::data::BlockedData(num_points, dims); + for (size_t i = 0; i < num_points; ++i) { + data.set_datum(i, points.get_datum(i)); + } + std::vector ids(num_points); + std::iota(ids.begin(), ids.end(), 0); + + svs::index::vamana::VamanaBuildParameters parameters{1.2, 32, 64, 1000, 28, true}; + auto index = cc::MutableVamanaIndex( + parameters, std::move(data), ids, svs::distance::DistanceL2(), 2 + ); + using buffer_type = svs::index::vamana::MutableBuffer; + + // Brute-force distances and a radius that holds about 30 neighbors per query. + auto distances = std::vector>(num_queries); + float radius = 0; + for (size_t q = 0; q < num_queries; ++q) { + for (size_t i = 0; i < num_points; ++i) { + distances[q].push_back(svs::distance::compute( + svs::distance::DistanceL2(), queries.get_datum(q), points.get_datum(i) + )); + } + auto sorted = distances[q]; + std::nth_element(sorted.begin(), sorted.begin() + 30, sorted.end()); + radius += sorted[30]; + } + radius /= static_cast(num_queries); + + auto check_range_results = [&](const auto& results, size_t max_results) { + CATCH_REQUIRE(results.size() == num_queries); + size_t num_expected = 0; + size_t num_found = 0; + for (size_t q = 0; q < num_queries; ++q) { + const auto& result = results[q]; + CATCH_REQUIRE(result.size() <= max_results); + auto seen = std::unordered_set{}; + float last = 0; + for (const auto& neighbor : result) { + CATCH_REQUIRE(neighbor.distance() < radius); + CATCH_REQUIRE(neighbor.distance() >= last); + CATCH_REQUIRE(neighbor.id() < num_points); + CATCH_REQUIRE(seen.insert(neighbor.id()).second); + last = neighbor.distance(); + } + for (size_t i = 0; i < num_points; ++i) { + if (index.has_id(i) && distances[q][i] < radius) { + ++num_expected; + num_found += seen.contains(i); + } + } + } + return static_cast(num_found) / static_cast(num_expected); + }; + + auto sp = index.get_search_parameters(); + + CATCH_SECTION("Default stop matches search") { + sp.buffer_config({20, 40}); + auto expected = svs::QueryResult(num_queries, k); + index.search(expected.view(), queries, sp); + auto actual = svs::QueryResult(num_queries, k); + index.predicated_search( + actual.view(), queries, sp, svs::index::vamana::DefaultStop{} + ); + for (size_t q = 0; q < num_queries; ++q) { + for (size_t j = 0; j < k; ++j) { + CATCH_REQUIRE(actual.index(q, j) == expected.index(q, j)); + CATCH_REQUIRE(actual.distance(q, j) == expected.distance(q, j)); + } + } + + // The single-query overload leaves internal IDs in the scratch buffer. + auto scratch = index.predicated_scratchspace(sp, svs::index::vamana::DefaultStop{}); + index.predicated_search(queries.get_datum(0), scratch); + for (size_t j = 0; j < k; ++j) { + CATCH_REQUIRE( + index.translate_internal_id(scratch.buffer[j].id()) == expected.index(0, j) + ); + } + } + + CATCH_SECTION("Range search") { + sp.buffer_config({10, buffer_type::unbounded_capacity}); + CATCH_REQUIRE( + check_range_results(index.range_search(queries, radius, sp), num_points) >= 0.95 + ); + + sp.buffer_config({10, 20}); + auto bounded = index.range_search(queries, radius, sp); + check_range_results(bounded, 20); + CATCH_REQUIRE(std::any_of(bounded.begin(), bounded.end(), [](const auto& r) { + return r.size() == 20; + })); + + // Deleted points are never returned. + std::vector to_delete; + for (size_t i = 0; i < num_points; i += 3) { + to_delete.push_back(i); + } + index.delete_entries(to_delete); + sp.buffer_config({10, buffer_type::unbounded_capacity}); + auto results = index.range_search(queries, radius, sp); + CATCH_REQUIRE(check_range_results(results, num_points) >= 0.95); + for (const auto& result : results) { + for (const auto& neighbor : result) { + CATCH_REQUIRE(neighbor.id() % 3 != 0); + } + } + } + + CATCH_SECTION("Cancel") { + sp.buffer_config({10, buffer_type::unbounded_capacity}); + std::atomic calls{0}; + auto partial = + index.range_search(queries, radius, sp, [&]() { return ++calls > 50; }); + check_range_results(partial, num_points); + CATCH_REQUIRE(std::any_of(partial.begin(), partial.end(), [](const auto& r) { + return r.empty(); + })); + } + + CATCH_SECTION("Concurrent with deletion and consolidation") { + sp.buffer_config({10, buffer_type::unbounded_capacity}); + std::atomic done{false}; + std::thread writer{[&] { + for (size_t i = 0; i < num_points / 2; i += 50) { + std::vector batch(50); + std::iota(batch.begin(), batch.end(), i); + index.delete_entries(batch); + index.consolidate(); + if (i % 200 == 0) { + index.compact(); + } + } + done = true; + }}; + size_t rounds = 0; + while (!done || rounds < 2) { + auto results = index.range_search(queries, radius, sp); + for (const auto& result : results) { + for (const auto& neighbor : result) { + CATCH_REQUIRE(neighbor.id() < num_points); + CATCH_REQUIRE(neighbor.distance() < radius); + } + } + auto topk = svs::QueryResult(num_queries, k); + index.predicated_search( + topk.view(), queries, sp, svs::index::vamana::DefaultStop{} + ); + ++rounds; + } + writer.join(); + CATCH_REQUIRE( + check_range_results(index.range_search(queries, radius, sp), num_points) >= 0.95 + ); + } +}