Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 2 additions & 9 deletions bindings/c/src/dispatcher_dynamic_vamana.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
130 changes: 102 additions & 28 deletions include/svs/core/translation.h
Original file line number Diff line number Diff line change
Expand Up @@ -35,9 +35,14 @@
#include "tsl/robin_map.h"

// stl
#include <algorithm>
#include <functional>
#include <iterator>
#include <limits>
#include <stdexcept>
#include <string_view>
#include <utility>
#include <vector>

namespace svs {

Expand All @@ -50,6 +55,10 @@ class IDTranslator {
tsl::robin_map<external_id_type, internal_id_type>::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<external_id_type>::max();

// Construct the identity transformation of size `n`.
struct Identity {
Identity() = delete;
Expand All @@ -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<size_t>(
static_cast<float>(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<std::pair<external_id_type, internal_id_type>, false>;
// external_to_internal_.bucket_count() * sizeof(bucket_type)
return size() * (sizeof(external_id_type) + sizeof(internal_id_type)) +

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What's wrong with bucket_count() * sizeof(bucket_type) estimation? It should be ok for trivial types.

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);
}

///
Expand Down Expand Up @@ -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);
}
Expand All @@ -140,8 +195,11 @@ class IDTranslator {

template <typename Ext, typename Int>
void insert_translation(Ext external_id, Int internal_id) {
external_to_internal_[external_id] = lib::narrow<internal_id_type>(internal_id);
internal_to_external_[internal_id] = lib::narrow<external_id_type>(external_id);
auto e = lib::narrow<external_id_type>(external_id);
auto i = lib::narrow<internal_id_type>(internal_id);
external_to_internal_[e] = i;
grow_internal(size_t{i} + 1);
internal_to_external_[i] = e;
}

///
Expand All @@ -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;
}

///
Expand All @@ -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];
}

///
Expand All @@ -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});
}

///
Expand Down Expand Up @@ -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);
}
}
Expand All @@ -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);
}
}
Expand All @@ -300,7 +360,7 @@ class IDTranslator {
///
template <class Begin, class End>
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");
}

///
Expand All @@ -314,7 +374,7 @@ class IDTranslator {
check(
begin,
end,
external_to_internal_,
external_contains(),
"Index does not contain external",
std::logical_not()
);
Expand All @@ -328,7 +388,7 @@ class IDTranslator {
///
template <class Begin, class End>
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");
}

///
Expand All @@ -342,7 +402,7 @@ class IDTranslator {
check(
begin,
end,
internal_to_external_,
internal_contains(),
"Index does not contain internal",
std::logical_not()
);
Expand Down Expand Up @@ -403,6 +463,7 @@ class IDTranslator {
auto num_points = lib::load_at<size_t>(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<external_id_type>(is);
auto internal_id = lib::read_binary<internal_id_type>(is);
Expand All @@ -422,23 +483,36 @@ class IDTranslator {
}

private:
template <class Begin, class End, class Map, class Modifier = lib::identity>
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 <class Begin, class End, class Contains, class Modifier = lib::identity>
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_id_type, internal_id_type> external_to_internal_{};
tsl::robin_map<internal_id_type, external_id_type> internal_to_external_{};
std::vector<external_id_type> internal_to_external_{};
};

} // namespace svs
2 changes: 2 additions & 0 deletions include/svs/index/flat/dynamic_flat.h
Original file line number Diff line number Diff line change
Expand Up @@ -259,6 +259,7 @@ template <typename Data, typename Dist> 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.
Expand Down Expand Up @@ -364,6 +365,7 @@ template <typename Data, typename Dist> class DynamicFlatIndex {
}
}
status_.resize(max_index);
translator_.shrink(max_index);
}

///// Consolidation
Expand Down
2 changes: 2 additions & 0 deletions include/svs/index/ivf/dynamic_ivf.h
Original file line number Diff line number Diff line change
Expand Up @@ -494,6 +494,7 @@ class DynamicIVFIndex {
std::vector<size_t> 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
Expand Down Expand Up @@ -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);
Expand Down
33 changes: 11 additions & 22 deletions include/svs/index/vamana/dynamic_index.h
Original file line number Diff line number Diff line change
Expand Up @@ -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<Idx>(0, external_ids.size()));
}

Expand Down Expand Up @@ -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<Idx>(0, external_ids.size()));

// Compute the entry point.
Expand Down Expand Up @@ -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{};
Expand All @@ -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;
}
Expand Down Expand Up @@ -526,18 +521,10 @@ class MutableVamanaIndex {
requires(std::tuple_size_v<Dims> == 2)
void translate_to_external(DenseArray<size_t, Dims, Base>& 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<Idx>(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<Idx>(id);
id = translate_internal_id(internal);
});
}

///
Expand Down Expand Up @@ -806,6 +793,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.
Expand Down Expand Up @@ -1006,6 +994,7 @@ class MutableVamanaIndex {
}
}
status_.resize(max_index);
translator_.shrink(max_index);

// Update entry points.
for (auto& ep : entry_point_) {
Expand Down
Loading
Loading