From 6934b23efbbe71f72ae9cb29e95562fc21b1e963 Mon Sep 17 00:00:00 2001 From: Rafik Saliev Date: Fri, 9 Oct 2026 08:17:37 -0700 Subject: [PATCH 1/2] feat(translator): Refactor IDTranslator to use std::vector for internal-to-external mapping - Added methods to estimate memory usage and manage internal capacity. - Updated memory management in MutableVamanaIndex and DynamicFlatIndex to utilize new translator methods. - Improved error handling for invalid external IDs and ensured proper resizing of internal structures. --- bindings/c/src/dispatcher_dynamic_vamana.cpp | 11 +- include/svs/core/translation.h | 130 +++++++++++++++---- include/svs/index/flat/dynamic_flat.h | 2 + include/svs/index/ivf/dynamic_ivf.h | 2 + include/svs/index/vamana/dynamic_index.h | 17 +-- tests/svs/index/vamana/dynamic_index.cpp | 14 +- 6 files changed, 124 insertions(+), 52 deletions(-) diff --git a/bindings/c/src/dispatcher_dynamic_vamana.cpp b/bindings/c/src/dispatcher_dynamic_vamana.cpp index 96a706f27..3b508fb8e 100644 --- a/bindings/c/src/dispatcher_dynamic_vamana.cpp +++ b/bindings/c/src/dispatcher_dynamic_vamana.cpp @@ -428,17 +428,10 @@ svs::index::vamana::MemoryBreakdown dispatch_dynamic_vamana_memory_estimate( estimate_data_size_blocked(storage, num_vectors, dimension, block_params); // Metadata: single entry point held as Idx, plus the SlotMetadata vector, plus the - // IDTranslator maps. + // IDTranslator tables. size_t metadata_bytes = sizeof(index_type) + sizeof(svs::index::vamana::SlotMetadata) * num_vectors; - // The IDTranslator holds two tsl::robin_map instances (external->internal and - // internal->external), neither of which exposes its allocated byte count. We - // approximate the storage as the id pair held in each of the two directions. This - // ignores the maps' load-factor slack and control bytes, so it is an estimate of - // the hash-map overhead that is accurate to within a few percent. - metadata_bytes += - 2 * num_vectors * - (sizeof(IDTranslator::external_id_type) + sizeof(IDTranslator::internal_id_type)); + metadata_bytes += IDTranslator::estimate_memory_usage(num_vectors); breakdown.metadata_bytes = metadata_bytes; return breakdown; } diff --git a/include/svs/core/translation.h b/include/svs/core/translation.h index 15d290b5e..85452923a 100644 --- a/include/svs/core/translation.h +++ b/include/svs/core/translation.h @@ -35,9 +35,14 @@ #include "tsl/robin_map.h" // stl +#include #include #include +#include +#include #include +#include +#include namespace svs { @@ -50,6 +55,10 @@ class IDTranslator { tsl::robin_map::const_iterator; using value_type = typename const_iterator::value_type; + /// Marks an unused slot in the internal-to-external table; not a valid external ID. + static constexpr external_id_type invalid_external_id = + std::numeric_limits::max(); + // Construct the identity transformation of size `n`. struct Identity { Identity() = delete; @@ -67,16 +76,57 @@ class IDTranslator { /// /// @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(); + size_t size() const { return external_to_internal_.size(); } - if (e2i != i2e) { - throw ANNEXCEPTION("Size mismatch! E2I is {} while I2E is {}!", e2i, i2e); - } + /// + /// @brief Make room for ``n`` translations with internal IDs in ``[0, n)``. + /// + void reserve(size_t n) { + grow_internal(n); + // ``robin_map::reserve`` always rehashes, so only call it when it must grow. + auto e2i_capacity = static_cast( + static_cast(external_to_internal_.bucket_count()) * + external_to_internal_.max_load_factor() + ); + if (n > e2i_capacity) { + external_to_internal_.reserve(n); } - return external_to_internal_.size(); + } + + /// + /// @brief Drop internal slots ``[n, ...)`` and release unused capacity. + /// + /// **Preconditions:** no internal ID ``>= n`` is mapped. + /// + void shrink(size_t n) { + assert(std::all_of( + internal_to_external_.begin() + std::min(n, internal_to_external_.size()), + internal_to_external_.end(), + [](external_id_type e) { return e == invalid_external_id; } + )); + internal_to_external_.resize(std::min(n, internal_to_external_.size())); + internal_to_external_.shrink_to_fit(); + } + + /// + /// @brief Return the number of bytes allocated for both translation tables. + /// + size_t get_memory_usage() const { + // The external->internal map is approximated by its entries, ignoring load-factor + // slack. An exact count depends on tsl::robin_map internals: + // using bucket_type = tsl::detail_robin_hash:: + // bucket_entry, false>; + // external_to_internal_.bucket_count() * sizeof(bucket_type) + return size() * (sizeof(external_id_type) + sizeof(internal_id_type)) + + internal_to_external_.capacity() * sizeof(external_id_type); + } + + /// + /// @brief Estimate ``get_memory_usage()`` for ``n`` translations over slots ``[0, n)``. + /// + static constexpr size_t estimate_memory_usage(size_t n) { + return n * (sizeof(external_id_type) + sizeof(internal_id_type)) + + n * sizeof(external_id_type); } /// @@ -124,6 +174,11 @@ class IDTranslator { if (!lib::all_unique(int_begin, int_end)) { throw ANNEXCEPTION("Internal IDs contain repeat elements!"); } + for (auto e = ext_begin; e != ext_end; ++e) { + if (std::cmp_equal(*e, invalid_external_id)) { + throw ANNEXCEPTION("External ID {} is reserved!", *e); + } + } check_external_free(ext_begin, ext_end); check_internal_free(int_begin, int_end); } @@ -140,8 +195,11 @@ class IDTranslator { 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); + auto e = lib::narrow(external_id); + auto i = lib::narrow(internal_id); + external_to_internal_[e] = i; + grow_internal(size_t{i} + 1); + internal_to_external_[i] = e; } /// @@ -159,7 +217,8 @@ class IDTranslator { /// @param e The internal ID to check. /// bool has_internal(internal_id_type e) const { - return internal_to_external_.contains(e); + return e < internal_to_external_.size() && + internal_to_external_[e] != invalid_external_id; } /// @@ -177,7 +236,10 @@ class IDTranslator { /// @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); + if (!has_internal(i)) { + throw std::out_of_range("IDTranslator: internal ID not found"); + } + return internal_to_external_[i]; } /// @@ -201,13 +263,11 @@ class IDTranslator { 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. + auto external = internal_to_external_[from]; + internal_to_external_[from] = invalid_external_id; + grow_internal(size_t{to} + 1); + internal_to_external_[to] = external; external_to_internal_[external] = to; - internal_to_external_.erase(itr); - internal_to_external_.insert({to, external}); } /// @@ -261,7 +321,7 @@ class IDTranslator { for (const auto i : internal_ids) { auto e = get_external(i); - internal_to_external_.erase(i); + internal_to_external_[i] = invalid_external_id; external_to_internal_.erase(e); } } @@ -287,7 +347,7 @@ class IDTranslator { for (const auto e : external_ids) { auto i = get_internal(e); - internal_to_external_.erase(i); + internal_to_external_[i] = invalid_external_id; external_to_internal_.erase(e); } } @@ -300,7 +360,7 @@ class IDTranslator { /// template void check_external_free(const Begin& begin, const End& end) const { - check(begin, end, external_to_internal_, "Index already contains external"); + check(begin, end, external_contains(), "Index already contains external"); } /// @@ -314,7 +374,7 @@ class IDTranslator { check( begin, end, - external_to_internal_, + external_contains(), "Index does not contain external", std::logical_not() ); @@ -328,7 +388,7 @@ class IDTranslator { /// template void check_internal_free(const Begin& begin, const End& end) const { - check(begin, end, internal_to_external_, "Index already contains internal"); + check(begin, end, internal_contains(), "Index already contains internal"); } /// @@ -342,7 +402,7 @@ class IDTranslator { check( begin, end, - internal_to_external_, + internal_contains(), "Index does not contain internal", std::logical_not() ); @@ -403,6 +463,7 @@ class IDTranslator { auto num_points = lib::load_at(table, "num_points"); auto translator = IDTranslator{}; + translator.reserve(num_points); for (size_t i = 0; i < num_points; ++i) { auto external_id = lib::read_binary(is); auto internal_id = lib::read_binary(is); @@ -422,23 +483,36 @@ class IDTranslator { } private: - template + void grow_internal(size_t n) { + if (n > internal_to_external_.size()) { + internal_to_external_.resize(n, invalid_external_id); + } + } + + auto external_contains() const { + return [this](external_id_type e) { return has_external(e); }; + } + auto internal_contains() const { + return [this](internal_id_type i) { return has_internal(i); }; + } + + template void check( const Begin& begin, const End& end, - const Map& map, + const Contains& contains, const char* message, Modifier modify = lib::identity() ) const { for (auto i = begin; i != end; ++i) { - if (modify(map.contains(*i))) { + if (modify(contains(*i))) { throw ANNEXCEPTION("{} ID {}!", message, *i); } } } tsl::robin_map external_to_internal_{}; - tsl::robin_map internal_to_external_{}; + std::vector internal_to_external_{}; }; } // namespace svs diff --git a/include/svs/index/flat/dynamic_flat.h b/include/svs/index/flat/dynamic_flat.h index e0bd5421c..567eb7940 100644 --- a/include/svs/index/flat/dynamic_flat.h +++ b/include/svs/index/flat/dynamic_flat.h @@ -259,6 +259,7 @@ template class DynamicFlatIndex { // Try to update the id translation now that we have internal ids. // If this fails, we still haven't mutated the index data structure so we're safe // to throw an exception. + translator_.reserve(status_.size()); translator_.insert(external_ids, slots); // Copy the given points into the data. @@ -364,6 +365,7 @@ template class DynamicFlatIndex { } } status_.resize(max_index); + translator_.shrink(max_index); } ///// Consolidation diff --git a/include/svs/index/ivf/dynamic_ivf.h b/include/svs/index/ivf/dynamic_ivf.h index b60ee534c..e6d18a7d5 100644 --- a/include/svs/index/ivf/dynamic_ivf.h +++ b/include/svs/index/ivf/dynamic_ivf.h @@ -494,6 +494,7 @@ class DynamicIVFIndex { std::vector global_ids = allocate_ids(num_points, reuse_empty); // Try to update ID translation + translator_.reserve(status_.size()); translator_.insert(external_ids, global_ids); // Insert points into their assigned clusters @@ -630,6 +631,7 @@ class DynamicIVFIndex { first_empty_ = new_size; // Step 6: Re-add all IDs to translator + translator_.reserve(new_size); translator_.insert(external_ids, new_internal_ids, false); svs::logging::info(logger_, "Compaction complete: {} valid entries", new_size); diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 9336f82e1..32a868847 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -206,6 +206,7 @@ class MutableVamanaIndex { if (graph_.n_nodes() != external_ids.size()) { throw ANNEXCEPTION("Graph node count does not match external IDs size"); } + translator_.reserve(external_ids.size()); translator_.insert(external_ids, threads::UnitRange(0, external_ids.size())); } @@ -302,6 +303,7 @@ class MutableVamanaIndex { use_full_search_history_ = build_parameters_.use_full_search_history; // Setup the initial translation of external to internal ids. + translator_.reserve(external_ids.size()); translator_.insert(external_ids, threads::UnitRange(0, external_ids.size())); // Compute the entry point. @@ -402,8 +404,8 @@ class MutableVamanaIndex { /// /// Reports the allocated memory for graph, data, and metadata components. Uses /// capacity-based accounting for datasets that expose ``capacity()``, so that block - /// over-allocation is reflected. Metadata includes status array, entry points, and an - /// estimated size of the ID translation maps (external/internal ID translation maps). + /// over-allocation is reflected. Metadata includes status array, entry points, and the + /// allocated size of the ID translation tables. MemoryBreakdown get_memory_breakdown() const { using namespace svs::data; MemoryBreakdown usage{}; @@ -413,14 +415,7 @@ class MutableVamanaIndex { size_t metadata_bytes = status_.capacity() * sizeof(SlotMetadata); metadata_bytes += entry_point_.capacity() * sizeof(typename entry_point_type::value_type); - // The IDTranslator holds two tsl::robin_map instances (external->internal and - // internal->external), neither of which exposes its allocated byte count. We - // approximate the storage as the id pair held in each of the two directions. This - // ignores the maps' load-factor slack and control bytes, so it is an estimate of - // the hash-map overhead that is accurate to within a few percent. - metadata_bytes += 2 * translator_.size() * - (sizeof(IDTranslator::external_id_type) + - sizeof(IDTranslator::internal_id_type)); + metadata_bytes += translator_.get_memory_usage(); usage.metadata_bytes = metadata_bytes; return usage; } @@ -806,6 +801,7 @@ class MutableVamanaIndex { // Try to update the id translation now that we have internal ids. // If this fails, we still haven't mutated the index data structure so we're safe // to throw an exception. + translator_.reserve(status_.size()); translator_.insert(external_ids, slots); // Copy the given points into the data and clear the adjacency lists for the graph. @@ -1006,6 +1002,7 @@ class MutableVamanaIndex { } } status_.resize(max_index); + translator_.shrink(max_index); // Update entry points. for (auto& ep : entry_point_) { diff --git a/tests/svs/index/vamana/dynamic_index.cpp b/tests/svs/index/vamana/dynamic_index.cpp index 94daddcf1..a58b98c5b 100644 --- a/tests/svs/index/vamana/dynamic_index.cpp +++ b/tests/svs/index/vamana/dynamic_index.cpp @@ -445,12 +445,16 @@ CATCH_TEST_CASE("MutableVamana Index Memory Usage", "[graph_index][dynamic_index const size_t expected_graph_bytes = index.view_graph().get_data().capacity() * index.view_graph().get_data().element_size(); using Index = decltype(index); + using ExternalId = typename Index::external_id_type; + using InternalId = typename Index::internal_id_type; + const size_t translator_bytes = index.view_translator().get_memory_usage(); + CATCH_REQUIRE( + translator_bytes >= indices.size() * (sizeof(ExternalId) + sizeof(InternalId)) + + indices.size() * sizeof(ExternalId) + ); const size_t expected_metadata_bytes = - data_size * sizeof(svs::index::vamana::SlotMetadata) + - sizeof(typename Index::internal_id_type) + - 2 * indices.size() * - (sizeof(typename Index::external_id_type) + - sizeof(typename Index::internal_id_type)); + data_size * sizeof(svs::index::vamana::SlotMetadata) + sizeof(InternalId) + + translator_bytes; const size_t expected_total_bytes = expected_data_bytes + expected_graph_bytes + expected_metadata_bytes; From b77a7c56b8e781ed426257e85620b8cf025ff10f Mon Sep 17 00:00:00 2001 From: Rafik Saliev Date: Tue, 13 Oct 2026 02:10:10 -0700 Subject: [PATCH 2/2] refactor(vamana): Simplify translate_to_external method using std::for_each --- include/svs/index/vamana/dynamic_index.h | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 32a868847..58088fb40 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -521,18 +521,10 @@ class MutableVamanaIndex { requires(std::tuple_size_v == 2) void translate_to_external(DenseArray& ids) { // N.B.: lib::narrow_cast should be valid because the origin of the IDs is internal. - threads::parallel_for( - threadpool_, - threads::StaticPartition{getsize<0>(ids)}, - [&](const auto is, uint64_t /*tid*/) { - for (auto i : is) { - for (size_t j = 0, jmax = getsize<1>(ids); j < jmax; ++j) { - auto internal = lib::narrow_cast(ids.at(i, j)); - ids.at(i, j) = translate_internal_id(internal); - } - } - } - ); + std::for_each(ids.begin(), ids.end(), [&](auto& id) { + const auto internal = lib::narrow_cast(id); + id = translate_internal_id(internal); + }); } ///