diff --git a/include/svs/concepts/graph.h b/include/svs/concepts/graph.h index d3e139eca..03edf1399 100644 --- a/include/svs/concepts/graph.h +++ b/include/svs/concepts/graph.h @@ -160,11 +160,11 @@ template using index_type_t = typename G::index_type; /// }; /// @endcode /// -template +template concept MemoryGraph = requires(T& g, const T& const_g) { // Adding an edge. requires requires(index_type_t src, index_type_t dst) { - { g.add_edge(src, dst) } -> std::convertible_to; + { g.add_edge(src, dst) } -> std::convertible_to; }; // Clear adjacency list. diff --git a/include/svs/concurrent/concurrent.h b/include/svs/concurrent/concurrent.h index 67c61b4c3..ce14c4ace 100644 --- a/include/svs/concurrent/concurrent.h +++ b/include/svs/concurrent/concurrent.h @@ -39,7 +39,6 @@ #include "svs/concurrent/prune.h" #include "svs/concurrent/reverse_edges.h" #include "svs/concurrent/spinlock.h" -#include "svs/concurrent/translation.h" #include "svs/concurrent/vamana_build.h" namespace svs { diff --git a/include/svs/concurrent/consolidate.h b/include/svs/concurrent/consolidate.h index 90af78f66..f8453c554 100644 --- a/include/svs/concurrent/consolidate.h +++ b/include/svs/concurrent/consolidate.h @@ -20,6 +20,7 @@ #include "svs/concurrent/graph.h" #include "svs/concurrent/prune.h" #include "svs/core/distance.h" +#include "svs/index/vamana/consolidate.h" #include "svs/index/vamana/extensions.h" #include "svs/lib/array.h" #include "svs/lib/threads.h" @@ -52,68 +53,7 @@ struct ConsolidationParameters { float alpha; }; -/// -/// A helper struct to store pending updates for the consolidated graph. -/// Represents pending updates using several large allocations rather than many small ones. -/// -template class BulkUpdate { - public: - // Constructor - BulkUpdate(size_t max_batch_size, size_t prune_to) - : neighbors_{make_dims(max_batch_size, prune_to)} - , lengths_{max_batch_size} - , needs_update_{max_batch_size} {} - - /// - /// Pre-conditions: - /// * 0 <= src < max_batch_size - /// * neighbors.size() <= prune_to - /// * May be called concurrently from multiple threads as long as `src` is - /// unique for each thread. - /// - template - void insert(size_t src, const std::vector& neighbors) { - assert(neighbors.size() <= getsize<1>(neighbors_)); - assert(src < getsize<0>(neighbors_)); - - needs_update_.at(src) = true; - // N.B.: There is no way we get even close to crossing what can be - // expressed as a 32-bit unsigned number. - lengths_.at(src) = lib::narrow_cast(neighbors.size()); - std::copy(neighbors.begin(), neighbors.end(), neighbors_.slice(src).begin()); - } - - /// - /// Return `true` if index `src` has a pending update. - /// - bool needs_update(size_t src) const { return needs_update_.at(src); } - - /// - /// Return the pending update for `src`. - /// Pre-conditions: - /// * `needs_update(src) == true` - /// - std::span get_update(size_t src) const { - assert(needs_update(src)); - return neighbors_.slice(src).first(lengths_.at(src)); - } - - /// - /// Prepare the data structure for another batch of processing. - /// **Preconditions:** - /// - /// * Must be called from a single thread. - /// - void prepare() { std::fill(needs_update_.begin(), needs_update_.end(), false); } - - private: - Matrix neighbors_; - Vector lengths_; - // N.B.: Use a `Vector` instead of a `std::vector` because `std::vector` - // specialized on `bool` and will thus cannot be guarantee coherent updates from - // multiple threads. - Vector needs_update_; -}; +using vamana::BulkUpdate; template struct ConsolidateThreadLocal { template using allocator_type = threads::CacheAlignedAllocator; @@ -290,22 +230,6 @@ class GraphConsolidator { } } - /// - /// Write pending updates to the graph. - /// - template - void apply_updates( - BulkUpdate& update_buffer, - const GlobalIds& global_ids, - const threads::UnitRange& local_ids - ) { - for (auto i : local_ids) { - if (update_buffer.needs_update(i)) { - graph_.replace_node(global_ids[i], update_buffer.get_update(i)); - } - } - } - /// /// Gather a superset of the live nodes with an edge to some deleted node in `D`. /// diff --git a/include/svs/concurrent/dynamic_index.h b/include/svs/concurrent/dynamic_index.h index 8c260b1ca..4261b4110 100644 --- a/include/svs/concurrent/dynamic_index.h +++ b/include/svs/concurrent/dynamic_index.h @@ -36,7 +36,6 @@ #include "svs/concurrent/dynamic_search_buffer.h" #include "svs/concurrent/graph.h" #include "svs/concurrent/greedy_search.h" -#include "svs/concurrent/translation.h" #include "svs/concurrent/vamana_build.h" #include "svs/core/data.h" #include "svs/core/distance.h" @@ -45,6 +44,7 @@ #include "svs/core/medioid.h" #include "svs/core/query_result.h" #include "svs/core/recall.h" +#include "svs/core/translation.h" #include "svs/index/vamana/index.h" #include "svs/lib/boundscheck.h" #include "svs/lib/neighbor.h" diff --git a/include/svs/concurrent/dynamic_search_buffer.h b/include/svs/concurrent/dynamic_search_buffer.h index d6811b433..0484f228d 100644 --- a/include/svs/concurrent/dynamic_search_buffer.h +++ b/include/svs/concurrent/dynamic_search_buffer.h @@ -16,20 +16,7 @@ #pragma once -#include "svs/index/vamana/search_buffer.h" -#include "svs/lib/boundscheck.h" -#include "svs/lib/datatype.h" -#include "svs/lib/neighbor.h" -#include "svs/lib/threads/threadlocal.h" - -#include -#include -#include -#include -#include -#include -#include -#include +#include "svs/index/vamana/dynamic_search_buffer.h" namespace svs::index::vamana::concurrent { @@ -38,81 +25,14 @@ namespace svs::index::vamana::concurrent { /// The search buffer will still navigate through those entries, but it won't include /// them in its final result. /// -template > class MutableBuffer { - public: - // Type Aliases - using index_type = Idx; - using value_type = PredicatedSearchNeighbor; - using reference = value_type&; - using const_reference = const value_type&; - using vector_type = std::vector>; - using iterator = typename vector_type::iterator; - using const_iterator = typename vector_type::const_iterator; - using compare_type = Cmp; - using filter_type = VisitedFilter; - - private: - ///// Invariants: - // - // 1. 1 <= target_valid_ <= valid_capacity_ - // - // The second inequality is enforced by the `BufferConfig` class. - // TODO: Enforce the first in `BufferConfig` as well. - // - // 2. best_unvisited_ <= size(). - // 3. roi_end_ <= size(). - // 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_`. - // - // 6. WHEN: `valid_ == valid_capacity_` - // THEN: `back()` must be valid. - // - // 7. The number of valid elements in `[begin(), end())` is equal to `valid_`. - - [[no_unique_address]] Cmp compare_{}; - // Equivalent of the `search_window_size_` in the traditional search buffer. - size_t target_valid_ = 0; - // Number of valid elements can are configured to contain. - // Equivalent to the `search_buffer_capacity_` in the traditional search buffer. - 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. - size_t valid_ = 0; - // A buffer of candidates. - // Unlike the static buffer, this container *does* dynamically change size and does not - // reserve one-past-the-end for copying neighbors. - vector_type candidates_{}; - // An optional visited filter. - std::optional visited_{std::nullopt}; +template > +class MutableBuffer : public vamana::MutableBuffer { + using base_type = vamana::MutableBuffer; public: - MutableBuffer() = default; - - /// Construct a new buffer with the given buffer configuration. - explicit MutableBuffer( - SearchBufferConfig config, Cmp compare = Cmp{}, bool enable_visited = false - ) - : compare_{std::move(compare)} - , target_valid_{config.get_search_window_size()} - , valid_capacity_{config.get_total_capacity()} - , candidates_{valid_capacity_} { - candidates_.clear(); - if (enable_visited) { - enable_visited_set(); - } - } + using base_type::base_type; - /// 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} {} + MutableBuffer() = default; /// Copy the portions of the MutableBuffer that matter for the purposes of scratch /// space. @@ -121,370 +41,28 @@ 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()}; - } - - // TODO: Allow this construction to be noexcept. - SearchBufferConfig config() const { - return SearchBufferConfig{target_valid_, valid_capacity_}; - } - - /// - /// @brief Change the target number of elements to return after search. - /// - /// @param config The new configuration for the buffer. - /// - /// Post conditions: - /// - The target valid capacity for the search buffer will be set to the new capacity. - /// - The actual size (number of contained elements both valid and invalid) will be the - /// minimum of the current size and the new capacity. - /// - void change_maxsize(SearchBufferConfig config) { - // Use temporary variables to ensure the given configuration is valid before - // committing. - size_t target_valid_temp = config.get_search_window_size(); - size_t valid_capacity_temp = config.get_total_capacity(); - - // 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); - } - - // Commit the new sizes. Integer assignment is `noexcept`, so we don't need to worry - // about an exception breaking class invariants. - target_valid_ = target_valid_temp; - valid_capacity_ = valid_capacity_temp; - } - - // Change the maximum number of elements that can be in the search buffer. - void change_maxsize(size_t new_size) { change_maxsize(SearchBufferConfig{new_size}); } - - /// @brief Prepare the buffer for a new search operation. - void clear() { - candidates_.clear(); - best_unvisited_ = 0; - roi_end_ = 0; - valid_ = 0; - if (visited_set_enabled()) { - visited_->reset(); - } - } - - void soft_clear() { - bool use_visited_set = visited_set_enabled(); - if (use_visited_set) { - visited_->reset(); - } - - for (auto& neighbor : candidates_) { - neighbor.clear_visited(); - if (use_visited_set) { - visited_->emplace(neighbor.id()); - } - } - - best_unvisited_ = 0; - } - - size_t capacity() const { return candidates_.capacity(); } - - /// @brief Return the number of valid elements currently contained in the buffer. - size_t valid() const { return valid_; } - - /// @brief Return the target valid capacity as equivalent to buffer capacity - size_t target_capacity() const { return valid_capacity_; } - - /// @brief Return the target valid candidates as equivalent to the search window - size_t target_window() const { return target_valid_; } - - /// @brief Return whether or not the buffer contains its target number of candidates. - bool full() const { return valid() == target_capacity(); } - - /// @brief Return the candidate at index `i`. - /// - /// Element will only be valid if: - /// (A) cleanup() has been invoked. - /// (B) 0 <= i < valid(); - reference operator[](size_t i) { return candidates_[i]; } - - /// @brief Return the candidate at index `i`. - /// - /// Element will only be valid if: - /// (A) cleanup() has been invoked. - /// (B) 0 <= i < valid(); - const_reference operator[](size_t i) const { return candidates_[i]; } - - /// @brief Return the last candidate, whether or not it is valid. - /// - /// It is undefined behavior to call this on an empty buffer. - reference back() { return candidates_.back(); } - - /// @brief Return the last candidate, whether or not it is valid. - /// - /// It is undefined behavior to call this on an empty buffer. - const_reference back() const { return candidates_.back(); } - - /// @brief Return the index of the best unvisited candidate. - size_t best_unvisited() const { return best_unvisited_; } - - /// @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_); - } - - /// @brief Return the best unvisited neighbor in the buffer. - /// The returned result will be convertible to `Neighbor`. - /// - /// Pre-conditions: - /// * `search_buffer.done()` must evaluate to `false`, otherwise an out of bounds - /// access will occur. - /// - /// Post-conditions: - /// * The returned neighbor will be marked as visited. - const_reference next() { - // Get the best unvisited node - reference node = getindex(candidates_, best_unvisited_); - node.set_visited(); - - // 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()) {} - return node; - } - - /// Place the neighbor at the end of the search buffer if `full() != true`. - /// Otherwise, do nothing. - /// - /// NOTE: `push_back` does not necessarily maintain the required invariants by this - /// class. - /// - /// These invariants are restored upon calling `sort()`. - /// - /// Therefore, sequences of `push_back` must always be followed by a call to `sort`. - void push_back(value_type neighbor) { - // Follow the contract of not appending more valid elements than the capacity - // allows. - bool valid = neighbor.valid(); - if (full() && valid) { - return; - } - candidates_.push_back(neighbor); - if (valid) { - ++valid_; - } - roi_end_ = candidates_.size(); - } - - // Define iterators. - constexpr const_iterator begin() const noexcept { return candidates_.begin(); } - constexpr const_iterator end() const noexcept { return begin() + size(); } - constexpr iterator begin() noexcept { return candidates_.begin(); } - constexpr iterator end() noexcept { return begin() + size(); } - - /// - /// Insert `neighbor` into the candidates buffer at position `pos`. - /// Afterwards, try to shrink the candidates buffer - /// - /// May invalidate any iterators for `candidates_`. - /// - void unsafe_insert(value_type neighbor, iterator pos, size_t offset) { - // 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); - 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; - } - } - - bool can_skip(float distance) const { - // If not full - we cannot skip appending this item. - // If we are full, then rely on invariant 6 to compare with the last valid element. - return full() && compare_(back().distance(), distance); - } - - /// Insert the given neighbor into the search buffer. - size_t insert(value_type neighbor) { - if (can_skip(neighbor.distance())) { - return size(); - } - return insert_inner(neighbor); - } - - size_t insert_inner(value_type neighbor) { - const auto start = begin(); - - // Binary search to the first location where `distance` is less than the stored - // neighbor. - auto pos = std::lower_bound( - start, - end(), - neighbor.distance(), - [&](const value_type& other, const float& d) { - return !compare_(d, other.distance()); - } - ); - - // Because repeat ids can exist, we have to search until we're sure no repeat will - // be found. - // - // To do that, we start one before the insertion position. - // Because each instance of repeat ids should have the same distance, we only need - // to look at `ids` until the buffer elements have a distance less than the - // distance of the current node. - if (pos != start) { - auto back = pos; - do { - --back; - const auto& candidate = *back; - if (compare_(candidate.distance(), neighbor.distance())) { - break; - } - if (candidate.id_ == neighbor.id()) { - return size() + 1; - } - } while (back != start); - } - - // 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; - valid_ += static_cast(neighbor.valid()); - unsafe_insert(neighbor, pos, i); - best_unvisited_ = std::min(best_unvisited_, i); - return i; + return MutableBuffer{ + base_type::config(), this->compare_, base_type::visited_set_enabled()}; } // Sort all stored elements in the buffer. void sort() { - assert(valid() <= valid_capacity_); + assert(base_type::valid() <= this->valid_capacity_); // Put the elements in ascending order. - std::sort(begin(), end(), compare_); + std::sort(base_type::begin(), base_type::end(), this->compare_); // Now - fix our invariants. - if (valid_ < target_valid_) { + if (this->valid_ < this->target_valid_) { // Maintain 5A // Invariant 6 has not activated. - roi_end_ = size(); + this->roi_end_ = base_type::size(); // Maintain best_unvisited_ invariant. // Recompute it from scratch, mirroring the 5B branch below. - for (best_unvisited_ = 0; best_unvisited_ < roi_end_; ++best_unvisited_) { - if (!candidates_[best_unvisited_].visited()) { + for (this->best_unvisited_ = 0; this->best_unvisited_ < this->roi_end_; + ++this->best_unvisited_) { + if (!this->candidates_[this->best_unvisited_].visited()) { break; } } @@ -493,12 +71,12 @@ template > class MutableBuffer { // Maintain 5B size_t valid_count = 0; - for (size_t i = 0; i < size(); ++i) { - if (candidates_[i].valid()) { + for (size_t i = 0; i < base_type::size(); ++i) { + if (this->candidates_[i].valid()) { ++valid_count; - if (valid_count == target_valid_) { + if (valid_count == this->target_valid_) { // One past the last valid element. - roi_end_ = i + 1; + this->roi_end_ = i + 1; break; } } @@ -508,118 +86,23 @@ template > class MutableBuffer { // 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()) { + 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 (slack() == 0) { - assert(!candidates_.empty()); - while (!back().valid()) { - candidates_.pop_back(); - assert(!candidates_.empty()); + if (base_type::slack() == 0) { + assert(!this->candidates_.empty()); + while (!base_type::back().valid()) { + this->candidates_.pop_back(); + assert(!this->candidates_.empty()); } } } - - // TODO: Switch over to using iterators for the return values to avoid this. - void cleanup() { - auto new_end = - std::remove_if(begin(), end(), [](const auto& x) { return !x.valid(); }); - candidates_.resize(new_end - begin()); - } - - /// @brief Return the size of the underlying vector. - /// - /// N.B.: THIS FUNCTION IS EASY TO CALL INCORRECTLY. - /// - /// At any given point, the size of the underlying vector and the number of valid - /// candidates in the buffer can be very different. - /// - /// These two are only the same after invoking `clean()`. - /// That is the only context in which non-internal interaction with the search buffer - /// should operate. - size_t size() const { return candidates_.size(); } - - ///// Visited API - bool visited_set_enabled() const { return visited_.has_value(); } - void enable_visited_set() { - if (!visited_set_enabled()) { - visited_.emplace(); - } - } - - /// @brief Enable or disable the visited set based on the argument. - void configure_visited_set(bool enable) { - if (enable) { - enable_visited_set(); - } else { - disable_visited_set(); - } - } - - void disable_visited_set() { - if (visited_set_enabled()) { - visited_.reset(); - } - } - - bool is_visited(Idx i) const { return visited_set_enabled() && unsafe_is_visited(i); } - void prefetch_visited(Idx i) const { - if (visited_set_enabled()) { - unsafe_prefetch_visited(i); - } - } - bool emplace_visited(Idx i) { - return visited_set_enabled() && unsafe_emplace_visited(i); - } - - // Unsafe API - bool unsafe_is_visited(Idx i) const { - assert(visited_); - return visited_->contains(i); - } - void unsafe_prefetch_visited(Idx i) const { - assert(visited_); - return visited_->prefetch(i); - } - bool unsafe_emplace_visited(Idx i) { - assert(visited_); - return visited_->emplace(i); - } - - private: - /// 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()); - } - - /// Return the index of the first preceding valid candidate beginning at the provided - /// index. - /// - /// Requires: - /// * `i` is in the range `[0, size())`. - /// * There exists at least one valid candidate in `[0, i)`. - size_t walk_back(size_t i) const { - while (!candidates_[i].valid()) { - --i; - } - return i; - } }; -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() - << ", valid = " << buffer.valid() << ", size = " << buffer.size() << ")"; -} - } // namespace svs::index::vamana::concurrent diff --git a/include/svs/concurrent/graph.h b/include/svs/concurrent/graph.h index e19d311a1..25582833c 100644 --- a/include/svs/concurrent/graph.h +++ b/include/svs/concurrent/graph.h @@ -742,7 +742,7 @@ class SimpleGraph : public SimpleGraphBase bool operator==(const SimpleGraph& x, const SimpleGraph& y) { - return graphs_equal(x, y); + return svs::graphs::graphs_equal(x, y); } template > diff --git a/include/svs/concurrent/graph_concepts.h b/include/svs/concurrent/graph_concepts.h index bb6eeb6f3..54399886b 100644 --- a/include/svs/concurrent/graph_concepts.h +++ b/include/svs/concurrent/graph_concepts.h @@ -16,16 +16,6 @@ #pragma once -// !! NOTICE TO MAINTAINERS !! -// -// Due to limitations in the Doxygen -> Breathe -> Sphinx interface, support for documenting -// C++ 20 Concepts is limited. -// -// For now, we need to put the member-wise documentation in a code block in the main -// concept docstring. -// -// Hopefully this situation changes in future versions of these systems. - /// /// @file /// @@ -40,13 +30,9 @@ /// @defgroup graph_concept_public Graph Concept Helpers. /// -#include -#include -#include +#include "svs/concepts/graph.h" + #include -#include -#include -#include namespace svs::index::vamana::concurrent::graphs { @@ -60,164 +46,7 @@ enum class AddEdgeResult : uint8_t { Rejected, // The source is no longer eligible for a backlink. }; -// clang-format off - -/// -/// @ingroup graph_concept_entry -/// @brief Main concept modeling immutable in-memory graphs. -/// -/// @code{.cpp} -/// template -/// concept ImmutableMemoryGraph = requires(const T& const_g) { -/// // The encoding of vertices in the graph. -/// // At the time of writing, this is expected to be an integer. -/// // This may be relaxed in the future. -/// typename T::index_type; -/// std::integral; -/// -/// // Has `reference` and `const_reference` type aliases. -/// // These types should be at least forward ranges, but preferably random access -/// // ranges. -/// // -/// // Items yieled by the iterators for these ranges should ``index_type``. -/// typename T::reference; -/// typename T::const_reference; -/// -/// // Return the maximum degree that this particular implementation of the graph is -/// // capable of supporting. -/// // -/// // If the graph supports unbounded adjacency lists, may return `Dynamic`. -/// { const_g.max_degree() } -> std::convertible_to; -/// -/// // Return the number of vertices contained in traph. -/// { const_g.n_nodes() } -> std::convertible_to; -/// -/// // Adjacency list operations. -/// requires requires(typename T::index_type i) { -/// // Return a range over the adjacency list for node ``i``. -/// { const_g.get_node(i) } -> std::same_as; -/// -/// // Return the number of out neighbors for node ``i``. -/// { const_g.get_node_degree(i) } -> std::convertible_to; -/// -/// // Prefetch the adjacency list for node ``i``. -/// // This is a performance optimization only and may be implemented as a no-op -/// // without affecting correctness. -/// const_g.prefetch_node(i); -/// }; -/// }; -/// @endcode -/// -template -concept ImmutableMemoryGraph = requires(const T& const_g) { - // The encoding of vertices in the graph. - typename T::index_type; - requires std::integral; - - // Has `reference` and `const_reference` type aliases. - typename T::reference; - typename T::const_reference; - - { const_g.max_degree() } -> std::convertible_to; - { const_g.n_nodes() } -> std::convertible_to; - - // Get an adjacency list. - requires requires(typename T::index_type i) { - { const_g.get_node(i) } -> std::same_as; - { const_g.get_node_degree(i) } -> std::convertible_to; - const_g.prefetch_node(i); - }; -}; -// clang-format on - -/// -/// @ingroup graph_concept_public -/// @brief Obtain the index type used to encode neighbors in the graph type ``G``. -/// -template using index_type_t = typename G::index_type; - -// clang-format off - -/// -/// @ingroup graph_concept_entry -/// @brief Concept modeling mutable in-memory graphs. -/// -/// @code{.cpp} -/// template -/// concept MemoryGraph = requires(T& g, const T& const_g) { -/// // Add an edge to the graph atomically. Returns an AddEdgeResult indicating: -/// // Added - edge was inserted -/// // AlreadyExists - edge was already present (or self-loop); no insert -/// // Full - node is at max_degree; edge NOT inserted (caller should -/// // route to an overflow buffer if needed) -/// requires requires(index_type_t src, index_type_t dst) { -/// { g.add_edge(src, dst) } -> std::convertible_to; -/// }; -/// -/// // Completely clear the adjacency list for vertex ``i``. -/// requires requires(index_type_t i) { -/// g.clear_node(i); -/// }; -/// -/// // Overwrite the adjacency list for `src`. -/// requires requires( -/// index_type_t src, -/// const std::vector>& neighbors_vector, -/// std::span> neighbors_span -/// ) { -/// g.replace_node(src, neighbors_vector); -/// g.replace_node(src, neighbors_span); -/// }; -/// }; -/// @endcode -/// template -concept MemoryGraph = requires(T& g, const T& const_g) { - // Adding an edge. - requires requires(index_type_t src, index_type_t dst) { - { g.add_edge(src, dst) } -> std::convertible_to; - }; - - // Clear adjacency list. - requires requires(index_type_t i) { - g.clear_node(i); - }; - - // Overwriting an edge. - requires requires( - index_type_t src, - const std::vector>& neighbors_vector, - std::span> neighbors_span - ) { - g.replace_node(src, neighbors_vector); - g.replace_node(src, neighbors_span); - }; -}; -// clang-format on - -/// -/// @ingroup graph_concept_public -/// @brief Compare the equality of two graphs. -/// -/// Two graphs are considered equal if: -/// -/// * The contain the same number of vertices. -/// * The adjacency lists for each vertex compare equal. -/// -template -bool graphs_equal(const Graph1& x, const Graph2& y) { - if (x.num_nodes() != y.num_nodes()) { - return false; - } - - for (size_t i = 0, imax = x.num_nodes(); i < imax; ++i) { - const auto& xa = x.get_node(i); - const auto& ya = y.get_node(i); - if (!std::equal(xa.begin(), xa.end(), ya.begin())) { - return false; - } - } - return true; -} +concept MemoryGraph = svs::graphs::MemoryGraph; } // namespace svs::index::vamana::concurrent::graphs diff --git a/include/svs/concurrent/greedy_search.h b/include/svs/concurrent/greedy_search.h index c8bddb5eb..fe8d7fd59 100644 --- a/include/svs/concurrent/greedy_search.h +++ b/include/svs/concurrent/greedy_search.h @@ -18,7 +18,7 @@ #include "svs/concepts/data.h" #include "svs/concepts/distance.h" -#include "svs/concurrent/graph_concepts.h" +#include "svs/concepts/graph.h" #include "svs/concurrent/spinlock.h" // For `svs::index::vamana::GreedySearchPrefetchParameters`, which this stack reuses // unchanged -- see the using-declaration below. @@ -51,7 +51,7 @@ using svs::index::vamana::EntryPointInitializer; using svs::index::vamana::NeighborBuilder; template < - graphs::ImmutableMemoryGraph Graph, + svs::graphs::ImmutableMemoryGraph Graph, data::ImmutableMemoryDataset Dataset, data::AccessorFor Accessor, typename QueryType, @@ -163,7 +163,7 @@ void greedy_search( // Overload to provide a default search tracker because search trackers are taken by // lvalue reference. template < - graphs::ImmutableMemoryGraph Graph, + svs::graphs::ImmutableMemoryGraph Graph, data::ImmutableMemoryDataset Dataset, data::AccessorFor Accessor, typename QueryType, diff --git a/include/svs/concurrent/iterator.h b/include/svs/concurrent/iterator.h index cbb6a802d..55eed4ca9 100644 --- a/include/svs/concurrent/iterator.h +++ b/include/svs/concurrent/iterator.h @@ -20,6 +20,7 @@ #include "svs/concurrent/dynamic_index.h" #include "svs/index/index.h" #include "svs/index/vamana/index.h" +#include "svs/index/vamana/iterator.h" #include "svs/lib/scopeguard.h" // stl @@ -28,128 +29,26 @@ namespace svs::index::vamana::concurrent { -/// @brief A graph search initializer that uses the existing contents of the search buffer -/// to initialize the next round of graph search. -/// -/// If a previous search exited with an exception, this initializer allows restarting -/// the search from scratch using the traditional method. -template struct RestartInitializer { - template < - typename Buffer, - typename Computer, - graphs::ImmutableMemoryGraph Graph, - typename Builder> - void operator()( - Buffer& buffer, - const Computer& computer, - const Graph& graph, - const Builder& builder, - vamana::NullTracker tracker // Compile error for non-NullTracker. - ) const { - // Restart the search from scratch if requested. - if (hard_restart_) { - vamana::EntryPointInitializer{ - entry_points_}(buffer, computer, graph, builder, tracker); - return; - } - - // Otherwise, sort the buffer to prepare for the next search. - buffer.sort(); - } - - // Entry points for restarting the search from scratch. - std::span entry_points_; - bool hard_restart_; -}; - -namespace detail { -constexpr void checkdims(size_t query_size, size_t index_dims) { - if (query_size != index_dims) { - throw ANNEXCEPTION( - "Incompatible dimensions. Query has {} while the index expects {}.", - query_size, - index_dims - ); - } -} -} // namespace detail - /// @brief A batch iterator for retrieving neighbors from the index in batches. /// /// This iterator abstracts the process of retrieving neighbors in fixed-size batches /// while maintaining internal state for efficient graph traversal. -template class BatchIterator { - public: - static_assert( - std::is_trivially_copyable_v, - "The batch iterator requires a trivial (no-throw copy constructible) type to " - "provide its exception guarantees" - ); - - // Public type aliases - using scratchspace_type = index::scratchspace_t; - using internal_id_type = typename Index::internal_id_type; - using external_id_type = size_t; - - // The value type yielded by the iterator. - using value_type = Neighbor; - - private: - // Private type aliases - using result_buffer_type = std::vector; - - /// @brief Copies results from the scratch buffer to the results buffer. - /// Ensures that only unique neighbors are added to the results buffer. - void copy_from_scratch(size_t batch_size) { - results_.clear(); - const auto& buffer = scratchspace_.buffer; - for (size_t i = 0, imax = buffer.size(); i < imax; ++i) { - auto neighbor = adapt(buffer[i]); - // Erased by a concurrent consolidation. - if (neighbor.id() == no_external_id) { - continue; - } - auto result = yielded_.insert(neighbor.id()); - if (result.second /* inserted */) { - // Rollback insertion into the yielded set if push_back throws. - auto guard = lib::make_dismissable_scope_guard([&]() noexcept { - yielded_.erase(result.first); - }); - results_.push_back(neighbor); - guard.dismiss(); - } - - // Stop if the requested batch size is reached. - if (results_.size() == batch_size) { - break; - } - } - } - - /// @brief Initializes the search buffer with the configured capacity. - void initialize_buffer() { - auto config = SearchBufferConfig{0, extra_search_buffer_capacity_}; - scratchspace_.buffer.change_maxsize(config); - scratchspace_.buffer.clear(); - } - - /// @brief Increments the search window and capacity by `batch_size` for the next - /// iteration. - void increment_buffer(size_t batch_size) { - auto config = scratchspace_.buffer.config(); - config.increment(batch_size); - scratchspace_.buffer.change_maxsize(config); - } +template +class BatchIterator + : public vamana::detail:: + BatchIteratorBase, Index, QueryType, size_t> { + using base_type = vamana::detail:: + BatchIteratorBase, Index, QueryType, size_t>; + friend base_type; + + using base_type::is_exhausted_; + using base_type::parent_; + using base_type::restart_search_; + using base_type::scratchspace_; + using base_type::yielded_; public: - using size_type = typename result_buffer_type::size_type; - using reference = value_type&; - using const_reference = const value_type&; - - /// Random-access iterator to `value_type` over the current batch of results. - using iterator = typename result_buffer_type::iterator; - /// Random-access iterator to `const value_type` over the current batch of results. - using const_iterator = typename result_buffer_type::const_iterator; + using typename base_type::external_id_type; /// @brief Constructs a batch iterator for the given query over the index. /// @param parent The index to search. @@ -162,33 +61,7 @@ template class BatchIterator { std::span query, size_t extra_search_buffer_capacity = svs::UNSIGNED_INTEGER_PLACEHOLDER ) - : parent_{&parent} - , query_{query.begin(), query.end()} - , scratchspace_{parent_->scratchspace()} { - detail::checkdims(query.size(), parent.dimensions()); - - extra_search_buffer_capacity_ = - extra_search_buffer_capacity == svs::UNSIGNED_INTEGER_PLACEHOLDER - ? svs::ITERATOR_EXTRA_BUFFER_CAPACITY_DEFAULT - : extra_search_buffer_capacity; - initialize_buffer(); - } - - /// @brief Updates the iterator with a new query. - /// Resets the internal state and restarts the search when `next(...)` is called. - void update(std::span newquery) { - detail::checkdims(newquery.size(), parent_->dimensions()); - assert(newquery.size() == query_.size()); - - std::copy(newquery.begin(), newquery.end(), query_.begin()); - - initialize_buffer(); - restart_search_ = true; - iteration_ = 0; - yielded_.clear(); - results_.clear(); - is_exhausted_ = false; - } + : base_type{parent, query, extra_search_buffer_capacity} {} /// @brief Adapts an internal neighbor to an external neighbor. template svs::Neighbor adapt(N internal) const { @@ -204,29 +77,6 @@ template class BatchIterator { } } - /// @brief Returns an iterator to the beginning of the results. - iterator begin() { return results_.begin(); } - /// @brief Returns an iterator to the end of the results. - iterator end() { return results_.end(); } - /// @copydoc begin() - const_iterator begin() const { return results_.begin(); } - /// @copydoc end() - const_iterator end() const { return results_.end(); } - /// @copydoc begin() - const_iterator cbegin() const { return results_.cbegin(); } - /// @copydoc begin() - const_iterator cend() const { return results_.cend(); } - - /// @brief Returns a span over the current batch of neighbors. - /// The span is invalidated by calls to `next(...)`. - std::span contents() const { return lib::as_const_span(results_); } - - /// @brief Returns the number of buffered results. - size_t size() const { return results_.size(); } - - /// @brief Return the batch number corresponding to the current buffer. - size_t batch_number() const { return iteration_; } - /// @brief Returns whether iterator can find more neighbors or not for the given query. /// /// The iterator is considered done when all the available nodes have been yielded or @@ -246,125 +96,46 @@ template class BatchIterator { return all_yielded; } - /// @brief Forces the next iteration to restart the search from scratch. - void restart_next_search() { restart_search_ = true; } - - /// @brief Returns the search parameters used for the current batch. - vamana::VamanaSearchParameters parameters_for_current_iteration() const { - auto& buffer = scratchspace_.buffer; - auto& prefetch = scratchspace_.prefetch_parameters; - return VamanaSearchParameters{ - buffer.config(), - buffer.visited_set_enabled(), - prefetch.lookahead, - prefetch.step}; - } - - /// @brief Prepares the next batch of neighbors (up to ``batch_size``) from the index. - /// Handles exceptions gracefully and ensures iterator state consistency. - void next( - size_t batch_size, - const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) - ) { - if (done()) { - results_.clear(); - return; - } - + private: + auto begin_search() { // Keep internal IDs stable through traversal, reranking, and translation. [[maybe_unused]] auto search_guard = parent_->lock_storage_shared(); - if (compaction_epoch_ != parent_->compaction_epoch_) { + const auto epoch = parent_->compaction_epoch_; + if (compaction_epoch_ != epoch) { // Reranking touches cached IDs before the restart initializer runs. // Preserve the grown window so the restart can reach unreturned neighbors. scratchspace_.buffer.clear(); restart_search_ = true; - compaction_epoch_ = parent_->compaction_epoch_; + compaction_epoch_ = epoch; } - increment_buffer(batch_size); - - bool restart_search_copy = std::exchange(restart_search_, true); - - { - parent_->experimental_escape_hatch([&]( - const auto& graph, - const auto& data, - const auto& SVS_UNUSED(distance), - std::span entry_points - ) { - auto search_closure = - [&](const auto& query, const auto& accessor, auto& d, auto& buffer) { - constexpr vamana::extensions::UsesReranking< - std::remove_const_t>> - uses_reranking{}; - if constexpr (uses_reranking()) { - distance::maybe_fix_argument(d, query); - for (size_t j = 0, jmax = buffer.size(); j < jmax; ++j) { - auto& neighbor = buffer[j]; - auto id = neighbor.id(); - auto new_distance = - distance::compute(d, query, data.get_primary(id)); - neighbor.set_distance(new_distance); - } - buffer.sort(); - } + return search_guard; + } - concurrent::greedy_search( - graph, - data, - accessor, - query, - d, - buffer, - RestartInitializer{entry_points, restart_search_copy}, - parent_->internal_search_builder(), - scratchspace_.prefetch_parameters, - cancel - ); + auto lock_for_translation() const { + // Hold the translation lock (translator_mutex_) so a + // concurrent consolidate/compact cannot erase or remap entries mid-copy. + return parent_->lock_for_translation(); + } - if constexpr (Index::needs_id_translation) { - buffer.cleanup(); - buffer.sort(); - } - }; + template auto prepare_neighbor(N internal) const { + return adapt(internal); + } - extensions::single_search( - data, - scratchspace_.buffer, - scratchspace_.scratch, - lib::as_const_span(query_), - search_closure, - *parent_ - ); - }); + template bool accept_neighbor(const N& neighbor) const { + // Erased by a concurrent consolidation. + if (neighbor.id() == no_external_id) { + return false; } + return true; + } - ++iteration_; - restart_search_ = false; - // Hold the translation lock (translator_mutex_) so a - // concurrent consolidate/compact cannot erase or remap entries mid-copy. - { - [[maybe_unused]] auto translation_guard = parent_->lock_for_translation(); - copy_from_scratch(batch_size); - } - // If result is empty after calling next(), mark the iterator as exhausted. - // The iterator will not be able to find any more neighbors. - if (results_.size() == 0 && batch_size > 0) { - is_exhausted_ = true; - } + template N finish_neighbor(N neighbor) const { return neighbor; } + + template static void run_search(Args&&... args) { + concurrent::greedy_search(SVS_FWD(args)...); } - private: - const Index* parent_; // The index being accessed. - std::vector query_; // Local buffer for the query. - scratchspace_type scratchspace_; // Scratch space for search. - std::vector> results_{}; // Filtered results from search. - std::unordered_set yielded_{}; // Set of yielded neighbors. - size_t compaction_epoch_ = 0; // Epoch of scratchspace_ IDs. - size_t iteration_ = 0; // Current iteration number. - bool restart_search_ = true; // Whether the next search should restart from scratch. - size_t extra_search_buffer_capacity_ = - svs::UNSIGNED_INTEGER_PLACEHOLDER; // Extra buffer capacity for the next search. - bool is_exhausted_ = false; // Whether the iterator is exhausted. + size_t compaction_epoch_ = 0; // Epoch of scratchspace_ IDs. }; // Deduction Guides diff --git a/include/svs/concurrent/multi.h b/include/svs/concurrent/multi.h index 5e00fc56f..ec59d79c4 100644 --- a/include/svs/concurrent/multi.h +++ b/include/svs/concurrent/multi.h @@ -17,6 +17,7 @@ #pragma once #include "svs/concurrent/dynamic_index.h" #include "svs/concurrent/iterator.h" +#include "svs/index/vamana/multi.h" #include #include @@ -935,70 +936,18 @@ auto auto_multi_dynamic_assemble( MultiMutableVamanaLoad load_from = MultiMutableVamanaLoad::FROM_MULTI, svs::logging::logger_ptr logger = svs::logging::get() ) { - // Load the dataset - auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); - auto data = svs::detail::dispatch_load(SVS_FWD(data_loader), threadpool); - - // Load the graph. - auto graph = svs::detail::dispatch_load(SVS_FWD(graph_loader), threadpool); - - // Make sure the data and the graph have the same size. - auto datasize = data.size(); - auto graphsize = graph.n_nodes(); - if (datasize != graphsize) { - throw ANNEXCEPTION( - "Reloaded data has {} nodes while the graph has {} nodes!", datasize, graphsize - ); - } - auto [parameters, translator, labels] = - lib::load_from_disk( - config_path, load_from, datasize - ); - - switch (load_from) { - case MultiMutableVamanaLoad::FROM_MULTI: { - if (labels.size() != datasize) { - throw ANNEXCEPTION( - "Labels has {} IDs but should have {}", labels.size(), datasize - ); - } - return MultiMutableVamanaIndex{ - parameters, - std::move(data), - std::move(graph), - std::move(distance), - labels, - std::move(threadpool), - std::move(logger)}; - } - case MultiMutableVamanaLoad::FROM_DYNAMIC: - case MultiMutableVamanaLoad::FROM_STATIC: { - // Make sure that the translator covers all the IDs in the graph and data. - auto translator_size = translator.size(); - if (translator_size != datasize) { - throw ANNEXCEPTION( - "Translator has {} IDs but should have {}", translator_size, datasize - ); - } - - for (size_t i = 0; i < datasize; ++i) { - if (!translator.has_internal(i)) { - throw ANNEXCEPTION("Translator is missing internal id {}", i); - } - } - - return MultiMutableVamanaIndex{ - parameters, - std::move(data), - std::move(graph), - std::move(distance), - std::move(translator), - std::move(threadpool), - std::move(logger)}; + return vamana::detail::assemble_multi_from_disk( + config_path, + SVS_FWD(graph_loader), + SVS_FWD(data_loader), + distance, + threadpool_proto, + load_from, + logger, + [](Args&&... args) { + return MultiMutableVamanaIndex{SVS_FWD(args)...}; } - default: - throw ANNEXCEPTION("Invalid multi vamana load type"); - } + ); } template < @@ -1014,45 +963,17 @@ auto auto_multi_dynamic_assemble( ThreadPoolProto threadpool_proto, svs::logging::logger_ptr logger = svs::logging::get() ) { - using label_type = size_t; - - auto table = lib::detail::read_metadata(is); - - auto parameters = lib::load( - table.template cast().at("parameters").template cast() + return vamana::detail::assemble_multi_from_stream( + is, + graph_loader, + data_loader, + distance, + threadpool_proto, + logger, + [](Args&&... args) { + return MultiMutableVamanaIndex{SVS_FWD(args)...}; + } ); - - auto num_labels = - lib::load(table.template cast().at("num_labels")); - - // Read labels binary data directly from the stream. - std::vector labels(num_labels); - lib::read_binary(is, labels); - - auto data = data_loader(); - auto graph = graph_loader(); - - auto datasize = data.size(); - auto graphsize = graph.n_nodes(); - if (datasize != graphsize) { - throw ANNEXCEPTION( - "Reloaded data has {} nodes while the graph has {} nodes!", datasize, graphsize - ); - } - - if (labels.size() != datasize) { - throw ANNEXCEPTION("Labels has {} IDs but should have {}", labels.size(), datasize); - } - - auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); - return MultiMutableVamanaIndex{ - parameters, - std::move(data), - std::move(graph), - std::move(distance), - labels, - std::move(threadpool), - std::move(logger)}; } } // namespace svs::index::vamana::concurrent diff --git a/include/svs/concurrent/spinlock.h b/include/svs/concurrent/spinlock.h index 6d1085c15..315f44d1d 100644 --- a/include/svs/concurrent/spinlock.h +++ b/include/svs/concurrent/spinlock.h @@ -18,32 +18,9 @@ #include "svs/lib/spinlock.h" -#include - namespace svs::index::vamana::concurrent { -/// -/// @brief ``svs::SpinLock`` plus value-initializing copy/move operations. -/// -/// ``lib::SegmentedVector`` constructs each segment's elements and, on -/// ``resize``, needs the element type to be constructible from an existing element. A -/// ``std::atomic`` member makes ``svs::SpinLock`` neither copyable nor movable, so -/// this subclass supplies those operations. A copied or moved lock is always born -/// *unlocked*: the operations exist to satisfy container requirements, never to transfer -/// ownership of a held lock, and a container is only ever grown while its existing -/// elements are untouched. -/// -class SpinLock : public svs::SpinLock { - public: - SpinLock() = default; - - SpinLock(const SpinLock& /*unused*/) - : svs::SpinLock{} {} - SpinLock& operator=(const SpinLock& /*unused*/) { return *this; } - SpinLock(SpinLock&& /*unused*/) noexcept - : svs::SpinLock{} {} - SpinLock& operator=(SpinLock&& /*unused*/) noexcept { return *this; } - ~SpinLock() = default; -}; +/// @brief Compatibility alias for svs::SpinLock. +using SpinLock = svs::SpinLock; } // namespace svs::index::vamana::concurrent diff --git a/include/svs/concurrent/translation.h b/include/svs/concurrent/translation.h deleted file mode 100644 index 5306689e4..000000000 --- a/include/svs/concurrent/translation.h +++ /dev/null @@ -1,522 +0,0 @@ -/* - * Copyright 2023 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 - -// Index translation. - -// svs -#include "svs/lib/algorithms.h" -#include "svs/lib/boundscheck.h" -#include "svs/lib/exception.h" -#include "svs/lib/file.h" -#include "svs/lib/misc.h" -#include "svs/lib/narrow.h" -#include "svs/lib/readwrite.h" -#include "svs/lib/saveload.h" -#include "svs/lib/threads.h" - -#include "svs/third-party/fmt.h" - -// tsl -#include "tsl/robin_map.h" - -// stl -#include -#include -#include - -namespace svs::index::vamana::concurrent { - -class IDTranslator { - public: - using internal_id_type = uint32_t; - using external_id_type = uint64_t; - - using const_iterator = - tsl::robin_map::const_iterator; - using value_type = typename const_iterator::value_type; - - // Construct the identity transformation of size `n`. - struct Identity { - Identity() = delete; - Identity(size_t n) - : n_{n} {} - size_t n_; - }; - - IDTranslator() = default; - IDTranslator(Identity tag) { - auto ids = threads::UnitRange{0, tag.n_}; - insert(ids, ids); - } - - /// - /// @brief Return the number of translations. - /// - size_t size() const { - if constexpr (checkbounds_v) { - const size_t e2i = external_to_internal_.size(); - const size_t i2e = internal_to_external_.size(); - - if (e2i != i2e) { - throw ANNEXCEPTION("Size mismatch! E2I is {} while I2E is {}!", e2i, i2e); - } - } - return external_to_internal_.size(); - } - - /// - /// @brief Insert the two ranges. - /// - /// @param external Container implementing forward iteration of the external IDs to add. - /// @param internal Container implementing forward iteration of the internal IDs to add. - /// @param check Check that none of the external and internal ids have an assigned - /// mapping yet **and** both only contain unique elements. Only safe to set to - /// ``false`` if this holds true. - /// - /// If any of the checks associated with the ``check`` parameter fail, the container - /// is left unmodified. - /// - template - void insert(const External& external, const Internal& internal, bool check = true) { - insert(external.begin(), external.end(), internal.begin(), internal.end(), check); - } - - template - void insert( - const ExtBegin& ext_begin, - const ExtEnd& ext_end, - const IntBegin& int_begin, - const IntEnd& int_end, - bool check = true - ) { - // Make sure the two iterators are the same length. - auto external_count = std::distance(ext_begin, ext_end); - auto internal_count = std::distance(int_begin, int_end); - if (external_count != internal_count) { - throw ANNEXCEPTION( - "Length of external IDs is {} while the length of internal IDs is {}!", - external_count, - internal_count - ); - } - - // Perform checks before actually modifying data structures to prevent a partially - // completed operation. - if (check) { - if (!lib::all_unique(ext_begin, ext_end)) { - throw ANNEXCEPTION("External IDs contain repeat elements!"); - } - if (!lib::all_unique(int_begin, int_end)) { - throw ANNEXCEPTION("Internal IDs contain repeat elements!"); - } - check_external_free(ext_begin, ext_end); - check_internal_free(int_begin, int_end); - } - - // Now, we actually begin the insertion. - auto i = int_begin; - auto e = ext_begin; - while (i != int_end || e != ext_end) { - insert_translation(*e, *i); - ++e; - ++i; - } - } - - template - void insert_translation(Ext external_id, Int internal_id) { - external_to_internal_[external_id] = lib::narrow(internal_id); - internal_to_external_[internal_id] = lib::narrow(external_id); - } - - /// - /// @brief Insert mappings, replacing any existing mapping whose current internal - /// ID is considered stale by the caller (e.g. the associated slot has been - /// marked for deletion but the translator entry was not yet cleaned up). - /// - /// Throws if an external ID already maps to a non-stale internal ID. - /// The batch is validated in full before any mutation, so on throw the - /// translation maps are left unmodified. - /// - /// @param external Container of external IDs to insert. - /// @param internal Container of internal IDs to insert (same length as external). - /// @param is_stale Callable ``bool(internal_id_type)``. Returns true if the given - /// internal ID is stale and its existing mapping may be overwritten. - /// - template - void replace_stale_and_insert( - const External& external, const Internal& internal, IsStale&& is_stale - ) { - auto ext_begin = external.begin(); - auto ext_end = external.end(); - auto int_begin = internal.begin(); - auto int_end = internal.end(); - - if (std::distance(ext_begin, ext_end) != std::distance(int_begin, int_end)) { - throw ANNEXCEPTION( - "Length of external IDs is {} while the length of internal IDs is {}!", - std::distance(ext_begin, ext_end), - std::distance(int_begin, int_end) - ); - } - if (!lib::all_unique(ext_begin, ext_end)) { - throw ANNEXCEPTION("External IDs contain repeat elements!"); - } - if (!lib::all_unique(int_begin, int_end)) { - throw ANNEXCEPTION("Internal IDs contain repeat elements!"); - } - - // Validation pass - for (auto e = ext_begin; e != ext_end; ++e) { - auto found = external_to_internal_.find(*e); - if (found != external_to_internal_.end() && !is_stale(found->second)) { - throw ANNEXCEPTION("Index already contains external ID {}!", *e); - } - } - - // Commit pass - auto e = ext_begin; - auto i = int_begin; - for (; e != ext_end; ++e, ++i) { - auto found = external_to_internal_.find(*e); - if (found != external_to_internal_.end()) { - internal_to_external_.erase(found->second); - } - insert_translation(*e, *i); - } - } - - /// - /// @brief Return whether the external ID exists. - /// - /// @param e The external ID to check. - /// - bool has_external(external_id_type e) const { - return external_to_internal_.contains(e); - } - - /// - /// @brief Return whether the internal ID exists. - /// - /// @param e The internal ID to check. - /// - bool has_internal(internal_id_type e) const { - return internal_to_external_.contains(e); - } - - /// - /// @brief Return the internal ID mapped to by the external ID. - /// - /// @param e The external ID to translate to an internal ID. - /// - internal_id_type get_internal(external_id_type e) const { - return external_to_internal_.at(e); - } - - /// - /// @brief Return the external ID mapped to by the internal ID. - /// - /// @param i The internal ID to translate to an external ID. - /// - external_id_type get_external(internal_id_type i) const { - return internal_to_external_.at(i); - } - - /// @brief Return the external ID, or a default if not found. - external_id_type - get_external_or(internal_id_type i, external_id_type default_val) const { - auto it = internal_to_external_.find(i); - if (it == internal_to_external_.end()) { - return default_val; - } - return it->second; - } - - /// @brief Return the internal ID, or a default if not found. - internal_id_type - get_internal_or(external_id_type e, internal_id_type default_val) const { - auto it = external_to_internal_.find(e); - if (it == external_to_internal_.end()) { - return default_val; - } - return it->second; - } - - /// - /// @brief Return a start forward iterator over the external->internal IDs. - /// - const_iterator begin() const { return external_to_internal_.begin(); } - - /// - /// @brief Return an end forward iterator over the external->internal IDs. - /// - const_iterator end() const { return external_to_internal_.end(); } - - /// - /// @brief Remap the internal ID. - /// - /// Assumptions: - /// * `from` exists. - /// * `to` does not exist. - /// - void remap_internal_id(internal_id_type from, internal_id_type to) { - assert(has_internal(from)); - assert(!has_internal(to)); - - auto itr = internal_to_external_.find(from); - auto external = itr->second; - - // Updating the internal-to-external ID is easy. - external_to_internal_[external] = to; - internal_to_external_.erase(itr); - internal_to_external_.insert({to, external}); - } - - /// - /// @brief Remap the external ID, keeping the same internal ID. - /// - /// @param from The existing external ID to rename. - /// @param to The external ID to rename it to. - /// @param check Check that ``from`` exists and ``to`` does not yet exist. Only safe - /// to set to ``false`` if this holds true. - /// - void remap_external_id(external_id_type from, external_id_type to, bool check = true) { - if (check) { - if (!has_external(from)) { - throw ANNEXCEPTION("Index does not contain external ID {}!", from); - } - if (has_external(to)) { - throw ANNEXCEPTION("Index already contains external ID {}!", to); - } - } - - auto itr = external_to_internal_.find(from); - auto internal = itr->second; - - // Updating the external-to-internal ID is easy. - internal_to_external_[internal] = to; - external_to_internal_.erase(itr); - external_to_internal_.insert({to, internal}); - } - - /// - /// @brief Delete entries from internal IDs. - /// - /// @param internal_ids A container with the internal ids to delete. Must implement a - /// forward iterator. Furthermore, all entries must be unique (not checked). - /// @param check Check if all ids actually exist in the translation table. This is only - /// safe to set to ``false`` if it can be guaranteed that all entries in - /// ``internal_ids`` exist in the table. - /// - /// Note, if ``check == true`` and an internal id is found not to exist, the underlying - /// translation tables will not be modified. - /// - template - void delete_internal(const Internal& internal_ids, bool check = true) { - // First pass - ensure that all IDs to remove actually exist. - if (check) { - check_internal_exist(internal_ids.begin(), internal_ids.end()); - } - - for (const auto i : internal_ids) { - auto e = get_external(i); - internal_to_external_.erase(i); - external_to_internal_.erase(e); - } - } - - /// - /// @brief Delete entries from external IDs. - /// - /// @param external_ids A container with the external ids to delete. Must implement a - /// forward iterator. Furthermore, all entries must be unique (not checked). - /// @param check Check if all ids actually exist in the translation table. This is only - /// safe to set to ``false`` if it can be guaranteed that all entries in - /// ``external_ids`` exist in the table. - /// - /// Note, if ``check == true`` and an external id is found not to exist, the underlying - /// translation tables will not be modified. - /// - template - void delete_external(const External& external_ids, bool check = true) { - // First pass - ensure that all IDs to remove actually exist. - if (check) { - check_external_exist(external_ids.begin(), external_ids.end()); - } - - for (const auto e : external_ids) { - auto i = get_internal(e); - internal_to_external_.erase(i); - external_to_internal_.erase(e); - } - } - - /// - /// @brief Ensure that **none** of the external ids in the iterator exist yet. - /// - /// @param begin Start forward iterator to the external IDs to check. - /// @param end End forward iterator to the external IDs to check. - /// - template - void check_external_free(const Begin& begin, const End& end) const { - check(begin, end, external_to_internal_, "Index already contains external"); - } - - /// - /// @brief Ensure that **all** of the external ids in the iterator exist. - /// - /// @param begin Start forward iterator to the external IDs to check. - /// @param end End forward iterator to the external IDs to check. - /// - template - void check_external_exist(const Begin& begin, const End& end) const { - check( - begin, - end, - external_to_internal_, - "Index does not contain external", - std::logical_not() - ); - } - - /// - /// @brief Ensure that **none** of the internal ids in the iterator exist yet. - /// - /// @param begin Start forward iterator to the internal IDs to check. - /// @param end End forward iterator to the internal IDs to check. - /// - template - void check_internal_free(const Begin& begin, const End& end) const { - check(begin, end, internal_to_external_, "Index already contains internal"); - } - - /// - /// @brief Ensure that **all** of the internal ids in the iterator exist. - /// - /// @param begin Start forward iterator to the internal IDs to check. - /// @param end End forward iterator to the internal IDs to check. - /// - template - void check_internal_exist(const Begin& begin, const End& end) const { - check( - begin, - end, - internal_to_external_, - "Index does not contain internal", - std::logical_not() - ); - } - - ///// Saving and Loading - static constexpr std::string_view kind = "external to internal id translation"; - static constexpr std::string_view serialization_schema = - "external_to_internal_translation"; - static constexpr lib::Version save_version = lib::Version(0, 0, 0); - - lib::SaveTable metadata() const { - return lib::SaveTable( - serialization_schema, - save_version, - {{"kind", kind}, - {"num_points", lib::save(size())}, - {"external_id_type", lib::save(datatype_v)}, - {"internal_id_type", lib::save(datatype_v)}} - ); - } - - void save(std::ostream& os) const { - for (auto i = begin(), iend = end(); i != iend; ++i) { - // N.B.: Apparently `std::pair` of integers is not trivially copyable ... - lib::write_binary(os, i->first); - lib::write_binary(os, i->second); - } - } - - lib::SaveTable save(const lib::SaveContext& ctx) const { - auto filename = ctx.generate_name("id_translation", "binary"); - // Save the translations to a file. - auto os = lib::open_write(filename); - save(os); - auto table = metadata(); - table.insert("filename", lib::save(filename.filename())); - return table; - } - - static void validate(const lib::ContextFreeLoadTable& table) { - if (kind != lib::load_at(table, "kind")) { - throw ANNEXCEPTION("Mismatched kind!"); - } - - // `svs::name` must be qualified here: the enclosing `svs::index::vamana` declares - // its own `name(SlotMetadata)`, which hides `svs::name()`. - constexpr std::string_view external_id_name = - svs::name>(); - constexpr std::string_view internal_id_name = - svs::name>(); - if (external_id_name != lib::load_at(table, "external_id_type")) { - throw ANNEXCEPTION("Mismatched external id types!"); - } - if (internal_id_name != lib::load_at(table, "internal_id_type")) { - throw ANNEXCEPTION("Mismatched internal id types!"); - } - } - - static IDTranslator load(const lib::ContextFreeLoadTable& table, std::istream& is) { - IDTranslator::validate(table); - auto num_points = lib::load_at(table, "num_points"); - - auto translator = IDTranslator{}; - for (size_t i = 0; i < num_points; ++i) { - auto external_id = lib::read_binary(is); - auto internal_id = lib::read_binary(is); - translator.insert_translation(external_id, internal_id); - } - return translator; - } - - static IDTranslator load(const lib::LoadTable& table) { - IDTranslator::validate(table); - - // Now that we've more-or-less validated the metadata, time to start loading - // the points. - auto resolved = table.resolve_at("filename"); - auto is = lib::open_read(resolved); - return IDTranslator::load(table, is); - } - - private: - template - void check( - const Begin& begin, - const End& end, - const Map& map, - const char* message, - Modifier modify = lib::identity() - ) const { - for (auto i = begin; i != end; ++i) { - if (modify(map.contains(*i))) { - throw ANNEXCEPTION("{} ID {}!", message, *i); - } - } - } - - tsl::robin_map external_to_internal_{}; - tsl::robin_map internal_to_external_{}; -}; - -} // namespace svs::index::vamana::concurrent diff --git a/include/svs/concurrent/vamana_build.h b/include/svs/concurrent/vamana_build.h index e0aa23b8e..1346f5079 100644 --- a/include/svs/concurrent/vamana_build.h +++ b/include/svs/concurrent/vamana_build.h @@ -21,7 +21,6 @@ #include "svs/concepts/distance.h" #include "svs/concurrent/greedy_search.h" #include "svs/concurrent/prune.h" -#include "svs/concurrent/spinlock.h" #include "svs/core/logging.h" #include "svs/index/vamana/build_params.h" #include "svs/index/vamana/extensions.h" diff --git a/include/svs/core/translation.h b/include/svs/core/translation.h index 15d290b5e..88b95311f 100644 --- a/include/svs/core/translation.h +++ b/include/svs/core/translation.h @@ -144,6 +144,65 @@ class IDTranslator { internal_to_external_[internal_id] = lib::narrow(external_id); } + /// + /// @brief Insert mappings, replacing any existing mapping whose current internal + /// ID is considered stale by the caller (e.g. the associated slot has been + /// marked for deletion but the translator entry was not yet cleaned up). + /// + /// Throws if an external ID already maps to a non-stale internal ID. + /// The batch is validated in full before any mutation, so on throw the + /// translation maps are left unmodified. + /// + /// @param external Container of external IDs to insert. + /// @param internal Container of internal IDs to insert (same length as external). + /// @param is_stale Callable ``bool(internal_id_type)``. Returns true if the given + /// internal ID is stale and its existing mapping may be overwritten. + /// + template + void replace_stale_and_insert( + const External& external, const Internal& internal, IsStale&& is_stale + ) { + auto ext_begin = external.begin(); + auto ext_end = external.end(); + auto int_begin = internal.begin(); + auto int_end = internal.end(); + + const auto external_count = std::distance(ext_begin, ext_end); + const auto internal_count = std::distance(int_begin, int_end); + if (external_count != internal_count) { + throw ANNEXCEPTION( + "Length of external IDs is {} while the length of internal IDs is {}!", + external_count, + internal_count + ); + } + if (!lib::all_unique(ext_begin, ext_end)) { + throw ANNEXCEPTION("External IDs contain repeat elements!"); + } + if (!lib::all_unique(int_begin, int_end)) { + throw ANNEXCEPTION("Internal IDs contain repeat elements!"); + } + + // Validation pass + for (auto e = ext_begin; e != ext_end; ++e) { + auto found = external_to_internal_.find(*e); + if (found != external_to_internal_.end() && !is_stale(found->second)) { + throw ANNEXCEPTION("Index already contains external ID {}!", *e); + } + } + + // Commit pass + auto e = ext_begin; + auto i = int_begin; + for (; e != ext_end; ++e, ++i) { + auto found = external_to_internal_.find(*e); + if (found != external_to_internal_.end()) { + internal_to_external_.erase(found->second); + } + insert_translation(*e, *i); + } + } + /// /// @brief Return whether the external ID exists. /// @@ -180,6 +239,26 @@ class IDTranslator { return internal_to_external_.at(i); } + /// @brief Return the external ID, or a default if not found. + external_id_type + get_external_or(internal_id_type i, external_id_type default_val) const { + auto it = internal_to_external_.find(i); + if (it == internal_to_external_.end()) { + return default_val; + } + return it->second; + } + + /// @brief Return the internal ID, or a default if not found. + internal_id_type + get_internal_or(external_id_type e, internal_id_type default_val) const { + auto it = external_to_internal_.find(e); + if (it == external_to_internal_.end()) { + return default_val; + } + return it->second; + } + /// /// @brief Return a start forward iterator over the external->internal IDs. /// diff --git a/include/svs/index/vamana/dynamic_search_buffer.h b/include/svs/index/vamana/dynamic_search_buffer.h index 9e8f304e3..52b071d9a 100644 --- a/include/svs/index/vamana/dynamic_search_buffer.h +++ b/include/svs/index/vamana/dynamic_search_buffer.h @@ -54,7 +54,7 @@ template > class MutableBuffer { /// A visited filter with 65,536 entries with a memory footprint of 128 kiB. using set_type = VisitedFilter; - private: + protected: ///// Invariants: // // 1. 1 <= target_valid_ <= valid_capacity_ @@ -618,7 +618,7 @@ template > class MutableBuffer { return visited_->emplace(i); } - private: + 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. @@ -627,6 +627,7 @@ template > class MutableBuffer { lib::narrow_cast(target_capacity()); } + private: /// Return the index of the first preceding valid candidate beginning at the provided /// index. /// diff --git a/include/svs/index/vamana/iterator.h b/include/svs/index/vamana/iterator.h index 44d76fb81..284b8bfb3 100644 --- a/include/svs/index/vamana/iterator.h +++ b/include/svs/index/vamana/iterator.h @@ -20,6 +20,7 @@ #include "svs/index/index.h" #include "svs/index/vamana/dynamic_index.h" #include "svs/index/vamana/index.h" +#include "svs/lib/misc.h" #include "svs/lib/scopeguard.h" // stl @@ -74,11 +75,14 @@ constexpr void checkdims(size_t query_size, size_t index_dims) { } } // namespace detail -/// @brief A batch iterator for retrieving neighbors from the index in batches. -/// -/// This iterator abstracts the process of retrieving neighbors in fixed-size batches -/// while maintaining internal state for efficient graph traversal. -template class BatchIterator { +namespace detail { + +template < + typename Derived, + typename Index, + typename QueryType, + typename YieldedId = typename Index::internal_id_type> +class BatchIteratorBase { public: static_assert( std::is_trivially_copyable_v, @@ -104,14 +108,17 @@ template class BatchIterator { results_.clear(); const auto& buffer = scratchspace_.buffer; for (size_t i = 0, imax = buffer.size(); i < imax; ++i) { - auto neighbor = buffer[i]; + auto neighbor = derived().prepare_neighbor(buffer[i]); + if (!derived().accept_neighbor(neighbor)) { + continue; + } auto result = yielded_.insert(neighbor.id()); if (result.second /* inserted */) { // Rollback insertion into the yielded set if push_back throws. auto guard = lib::make_dismissable_scope_guard([&]() noexcept { yielded_.erase(result.first); }); - results_.push_back(adapt(neighbor)); + results_.push_back(derived().finish_neighbor(neighbor)); guard.dismiss(); } @@ -147,13 +154,7 @@ template class BatchIterator { /// Random-access iterator to `const value_type` over the current batch of results. using const_iterator = typename result_buffer_type::const_iterator; - /// @brief Constructs a batch iterator for the given query over the index. - /// @param parent The index to search. - /// @param query The query data. - /// @param extra_search_buffer_capacity Additional buffer capacity for the search. - /// When not provided, ``svs::ITERATOR_EXTRA_BUFFER_CAPACITY_DEFAULT = 100`` is - /// used. - BatchIterator( + BatchIteratorBase( const Index& parent, std::span query, size_t extra_search_buffer_capacity = svs::UNSIGNED_INTEGER_PLACEHOLDER @@ -248,11 +249,12 @@ template class BatchIterator { size_t batch_size, const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) ) { - if (done()) { + if (derived().done()) { results_.clear(); return; } + [[maybe_unused]] auto search_guard = derived().begin_search(); increment_buffer(batch_size); bool restart_search_copy = std::exchange(restart_search_, true); @@ -280,7 +282,7 @@ template class BatchIterator { buffer.sort(); } - vamana::greedy_search( + derived().run_search( graph, data, accessor, @@ -311,7 +313,10 @@ template class BatchIterator { ++iteration_; restart_search_ = false; - copy_from_scratch(batch_size); + { + [[maybe_unused]] auto translation_guard = derived().lock_for_translation(); + copy_from_scratch(batch_size); + } // If result is empty after calling next(), mark the iterator as exhausted. // The iterator will not be able to find any more neighbors. if (results_.size() == 0 && batch_size > 0) { @@ -319,12 +324,28 @@ template class BatchIterator { } } - private: + protected: + Derived& derived() { return static_cast(*this); } + const Derived& derived() const { return static_cast(*this); } + + lib::Empty begin_search() { return {}; } + lib::Empty lock_for_translation() const { return {}; } + + template N prepare_neighbor(N neighbor) const { return neighbor; } + template bool accept_neighbor(const N&) const { return true; } + template auto finish_neighbor(N neighbor) const { + return derived().adapt(neighbor); + } + + template static void run_search(Args&&... args) { + vamana::greedy_search(SVS_FWD(args)...); + } + const Index* parent_; // The index being accessed. std::vector query_; // Local buffer for the query. scratchspace_type scratchspace_; // Scratch space for search. std::vector> results_{}; // Filtered results from search. - std::unordered_set yielded_{}; // Set of yielded neighbors. + std::unordered_set yielded_{}; // Set of yielded neighbors. size_t iteration_ = 0; // Current iteration number. bool restart_search_ = true; // Whether the next search should restart from scratch. size_t extra_search_buffer_capacity_ = @@ -332,6 +353,33 @@ template class BatchIterator { bool is_exhausted_ = false; // Whether the iterator is exhausted. }; +} // namespace detail + +/// @brief A batch iterator for retrieving neighbors from the index in batches. +/// +/// This iterator abstracts the process of retrieving neighbors in fixed-size batches +/// while maintaining internal state for efficient graph traversal. +template +class BatchIterator + : public detail::BatchIteratorBase, Index, QueryType> { + using base_type = + detail::BatchIteratorBase, Index, QueryType>; + + public: + /// @brief Constructs a batch iterator for the given query over the index. + /// @param parent The index to search. + /// @param query The query data. + /// @param extra_search_buffer_capacity Additional buffer capacity for the search. + /// When not provided, ``svs::ITERATOR_EXTRA_BUFFER_CAPACITY_DEFAULT = 100`` is + /// used. + BatchIterator( + const Index& parent, + std::span query, + size_t extra_search_buffer_capacity = svs::UNSIGNED_INTEGER_PLACEHOLDER + ) + : base_type{parent, query, extra_search_buffer_capacity} {} +}; + // Deduction Guides template BatchIterator(const Index*, std::span) -> BatchIterator; diff --git a/include/svs/index/vamana/multi.h b/include/svs/index/vamana/multi.h index 277771fb9..0ca02ecb3 100644 --- a/include/svs/index/vamana/multi.h +++ b/include/svs/index/vamana/multi.h @@ -775,23 +775,24 @@ struct MultiVamanaStateLoader { IDTranslator translator_; std::vector labels_; }; -} // namespace detail template < + typename StateLoader, typename GraphLoader, typename DataLoader, typename Distance, - typename ThreadPoolProto> -auto auto_multi_dynamic_assemble( + typename ThreadPoolProto, + typename LoadMode, + typename Factory> +auto assemble_multi_from_disk( const std::filesystem::path& config_path, GraphLoader&& graph_loader, DataLoader&& data_loader, - Distance distance, - ThreadPoolProto threadpool_proto, - /// This flag provides compatibility paths for directly loading dynamic vamana or static - /// vamana datasets. - MultiMutableVamanaLoad load_from = MultiMutableVamanaLoad::FROM_MULTI, - svs::logging::logger_ptr logger = svs::logging::get() + Distance& distance, + ThreadPoolProto& threadpool_proto, + LoadMode load_from, + svs::logging::logger_ptr& logger, + const Factory& make_index ) { // Load the dataset auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); @@ -809,28 +810,27 @@ auto auto_multi_dynamic_assemble( ); } auto [parameters, translator, labels] = - lib::load_from_disk( - config_path, load_from, datasize - ); + lib::load_from_disk(config_path, load_from, datasize); switch (load_from) { - case MultiMutableVamanaLoad::FROM_MULTI: { + case LoadMode::FROM_MULTI: { if (labels.size() != datasize) { throw ANNEXCEPTION( "Labels has {} IDs but should have {}", labels.size(), datasize ); } - return MultiMutableVamanaIndex{ + return make_index( parameters, std::move(data), std::move(graph), std::move(distance), labels, std::move(threadpool), - std::move(logger)}; + std::move(logger) + ); } - case MultiMutableVamanaLoad::FROM_DYNAMIC: - case MultiMutableVamanaLoad::FROM_STATIC: { + case LoadMode::FROM_DYNAMIC: + case LoadMode::FROM_STATIC: { // Make sure that the translator covers all the IDs in the graph and data. auto translator_size = translator.size(); if (translator_size != datasize) { @@ -845,14 +845,15 @@ auto auto_multi_dynamic_assemble( } } - return MultiMutableVamanaIndex{ + return make_index( parameters, std::move(data), std::move(graph), std::move(distance), std::move(translator), std::move(threadpool), - std::move(logger)}; + std::move(logger) + ); } default: throw ANNEXCEPTION("Invalid multi vamana load type"); @@ -863,14 +864,16 @@ template < typename LazyGraphLoader, typename LazyDataLoader, typename Distance, - typename ThreadPoolProto> -auto auto_multi_dynamic_assemble( + typename ThreadPoolProto, + typename Factory> +auto assemble_multi_from_stream( std::istream& is, - LazyGraphLoader graph_loader, - LazyDataLoader data_loader, - Distance distance, - ThreadPoolProto threadpool_proto, - svs::logging::logger_ptr logger = svs::logging::get() + LazyGraphLoader& graph_loader, + LazyDataLoader& data_loader, + Distance& distance, + ThreadPoolProto& threadpool_proto, + svs::logging::logger_ptr& logger, + const Factory& make_index ) { using label_type = size_t; @@ -903,14 +906,73 @@ auto auto_multi_dynamic_assemble( } auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); - return MultiMutableVamanaIndex{ + return make_index( parameters, std::move(data), std::move(graph), std::move(distance), labels, std::move(threadpool), - std::move(logger)}; + std::move(logger) + ); +} + +} // namespace detail + +template < + typename GraphLoader, + typename DataLoader, + typename Distance, + typename ThreadPoolProto> +auto auto_multi_dynamic_assemble( + const std::filesystem::path& config_path, + GraphLoader&& graph_loader, + DataLoader&& data_loader, + Distance distance, + ThreadPoolProto threadpool_proto, + /// This flag provides compatibility paths for directly loading dynamic vamana or static + /// vamana datasets. + MultiMutableVamanaLoad load_from = MultiMutableVamanaLoad::FROM_MULTI, + svs::logging::logger_ptr logger = svs::logging::get() +) { + return detail::assemble_multi_from_disk( + config_path, + SVS_FWD(graph_loader), + SVS_FWD(data_loader), + distance, + threadpool_proto, + load_from, + logger, + [](Args&&... args) { + return MultiMutableVamanaIndex{SVS_FWD(args)...}; + } + ); +} + +template < + typename LazyGraphLoader, + typename LazyDataLoader, + typename Distance, + typename ThreadPoolProto> +auto auto_multi_dynamic_assemble( + std::istream& is, + LazyGraphLoader graph_loader, + LazyDataLoader data_loader, + Distance distance, + ThreadPoolProto threadpool_proto, + svs::logging::logger_ptr logger = svs::logging::get() +) { + return detail::assemble_multi_from_stream( + is, + graph_loader, + data_loader, + distance, + threadpool_proto, + logger, + [](Args&&... args) { + return MultiMutableVamanaIndex{SVS_FWD(args)...}; + } + ); } } // namespace svs::index::vamana diff --git a/tests/svs/concurrent/dynamic_index.cpp b/tests/svs/concurrent/dynamic_index.cpp index 25b92095a..214dec205 100644 --- a/tests/svs/concurrent/dynamic_index.cpp +++ b/tests/svs/concurrent/dynamic_index.cpp @@ -31,6 +31,7 @@ // svs #include "svs/core/recall.h" +#include "svs/lib/spinlock.h" #include "svs/lib/timing.h" // catch2 @@ -608,7 +609,7 @@ CATCH_TEST_CASE( // policy. So pin the directory and require a non-empty payload on top of it. using ReverseList = typename ConcurrentGraph::reverse_edges_type::list_type; const size_t expected_reverse_directory_bytes = - segmented_capacity * (sizeof(ReverseList) + sizeof(cc::SpinLock)); + segmented_capacity * (sizeof(ReverseList) + sizeof(svs::SpinLock)); // Dynamic get_memory_usage() should exactly match the capacity-based graph and data // bytes plus the deterministic metadata implied by the input ids, plus the diff --git a/tests/svs/concurrent/translation.cpp b/tests/svs/concurrent/translation.cpp index 9aef37379..8791bbb61 100644 --- a/tests/svs/concurrent/translation.cpp +++ b/tests/svs/concurrent/translation.cpp @@ -15,7 +15,7 @@ */ // header under test. -#include "svs/concurrent/translation.h" +#include "svs/core/translation.h" // misc utilities #include "svs/lib/narrow.h" @@ -32,7 +32,7 @@ // The concurrent index carries its own ID translator (transactional // `replace_stale_and_insert` plus the non-throwing `_or` lookups). It is a // distinct type from `svs::IDTranslator`, which is left untouched. -using IDTranslator = svs::index::vamana::concurrent::IDTranslator; +using IDTranslator = svs::IDTranslator; namespace { template