diff --git a/benchmarks/hotpaths/CMakeLists.txt b/benchmarks/hotpaths/CMakeLists.txt index 9f356d99..afb26ba6 100644 --- a/benchmarks/hotpaths/CMakeLists.txt +++ b/benchmarks/hotpaths/CMakeLists.txt @@ -21,4 +21,10 @@ if(RABITQ_BUILD_INDEX) "${CMAKE_CURRENT_BINARY_DIR}/sidecar-smoke" 128 64 2) set_tests_properties(SidecarHotpathPrepare PROPERTIES FIXTURES_SETUP sidecar_hotpath) set_tests_properties(SidecarHotpathLoad PROPERTIES FIXTURES_REQUIRED sidecar_hotpath) + + rabitq_hotpath_target(ivf_add index) + add_test(NAME IvfAddHotpathSingle COMMAND rabitq_ivf_add_hotpaths + "${CMAKE_CURRENT_BINARY_DIR}/ivf-add-smoke-single" 64 8 256 16 1 single) + add_test(NAME IvfAddHotpathRouted COMMAND rabitq_ivf_add_hotpaths + "${CMAKE_CURRENT_BINARY_DIR}/ivf-add-smoke-routed" 64 8 256 16 1 routed) endif() diff --git a/benchmarks/hotpaths/ivf_add.cpp b/benchmarks/hotpaths/ivf_add.cpp new file mode 100644 index 00000000..7ecbe11e --- /dev/null +++ b/benchmarks/hotpaths/ivf_add.cpp @@ -0,0 +1,96 @@ +// Times one IVF::add() per repetition, each on a freshly loaded copy of the base index. +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common.hpp" +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/index/ivf/ivf.hpp" +#include "rabitqlib/index/types.hpp" + +namespace { +void run(int argc, char** argv) { + if (argc != 8) { + throw std::invalid_argument( + "Usage: rabitq_ivf_add_hotpaths path dim clusters base_n add_n reps " + "single|routed" + ); + } + const std::string path(argv[1]); + const size_t dim = bench::number(argv[2]); + const size_t clusters = bench::number(argv[3]); + const size_t base_n = bench::number(argv[4]); + const size_t add_n = bench::number(argv[5]); + const size_t reps = bench::number(argv[6]); + const std::string mode(argv[7]); + if (mode != "single" && mode != "routed") { + throw std::invalid_argument("Mode must be single or routed"); + } + if (clusters < 1 || base_n < clusters) { + throw std::invalid_argument("base_n must be at least clusters"); + } + + const auto base_data = bench::values(bench::product(base_n, dim), 42); + const auto add_data = bench::values(bench::product(add_n, dim), 91); + const auto centroids = bench::values(bench::product(clusters, dim), 7); + // Round-robin assignment; clustering quality is irrelevant to storage cost. + std::vector base_assigned(base_n); + for (size_t i = 0; i < base_n; ++i) { + base_assigned[i] = static_cast(i % clusters); + } + // "single" sends every new point to cluster 0; "routed" lets add() pick the nearest. + const std::vector single_ids(add_n, rabitqlib::PID{0}); + const rabitqlib::PID* add_cluster_ids = mode == "single" ? single_ids.data() : nullptr; + + { + rabitqlib::ivf::IVF base(base_n, dim, clusters, 4, rabitqlib::METRIC_L2); + base.construct(base_data.data(), centroids.data(), base_assigned.data(), false, 1); + base.save(path.c_str()); + } + + bench::metadata(dim, base_n, 1); + const char* name = mode == "single" ? "add_single_cluster" : "add_routed"; + uint64_t expected = 0; + for (size_t rep = 0; rep <= reps; ++rep) { + rabitqlib::ivf::IVF index; + index.load(path.c_str()); + + const auto start = bench::Clock::now(); + index.add(add_data.data(), add_n, add_cluster_ids, false, 1); + const double seconds = + std::chrono::duration(bench::Clock::now() - start).count(); + + uint64_t h = 14695981039346656037ULL; + const size_t total = index.max_elements(); + h = bench::hash(&total, sizeof(total), h); + std::vector ids(10); + std::vector distances(10); + index.search(base_data.data(), ids.size(), clusters, ids.data(), distances.data()); + h = bench::hash(ids.data(), ids.size() * sizeof(rabitqlib::PID), h); + h = bench::hash(distances.data(), distances.size() * sizeof(float), h); + + if (rep == 0) { + expected = h; + } else { + if (h != expected) { + throw std::runtime_error("add() result changed between repetitions"); + } + bench::row(name, rep - 1, add_n, seconds, h); + } + } +} +} // namespace + +int main(int argc, char** argv) { + try { + run(argc, argv); + } catch (const std::exception& error) { + std::cerr << error.what() << '\n'; + return 1; + } +} diff --git a/docs/docs/index/ivf.md b/docs/docs/index/ivf.md index 2c1a614b..a545f286 100644 --- a/docs/docs/index/ivf.md +++ b/docs/docs/index/ivf.md @@ -278,6 +278,8 @@ void IVF::add( ); size_t IVF::remove(const PID* ids_to_remove, size_t n); + +void IVF::reserve(PID cluster_id, size_t num_points); ``` - **data**, **n**: `n` new vectors, quantized like `construct` does. Vector `i` @@ -290,6 +292,9 @@ size_t IVF::remove(const PID* ids_to_remove, size_t n); - **ids_to_remove**: PIDs to remove. Every PID must be below `num_stored_points()`; nothing is removed if one is not. `remove` returns how many were newly removed and can be repeated safely. +- **cluster_id**, **num_points**: `reserve` preallocates room for cluster `cluster_id` + to hold `num_points` vectors in total, counting the ones it already holds. It does + nothing if the cluster already has that room. See [Cost of `add`](#cost-of-add). In Python: @@ -298,6 +303,7 @@ new_ids = index.add(vectors, num_threads=1, fast_quantization=False) # To choose clusters explicitly, use this instead of the call above: # new_ids = index.add(vectors, cluster_ids=labels) removed = index.remove(new_ids[:10]) +index.reserve(cluster_id, num_points) ``` Python `vectors` has shape `(n, index.dim)` and is converted to contiguous float32. @@ -314,14 +320,21 @@ access; see the ### Cost of `add` -`add` builds the grown index next to the old one and swaps it in, so the index is -unchanged if it throws. Every call copies the whole storage, so: +Each cluster has its own storage, and `add` only touches the clusters that receive +new vectors. Each of those is rebuilt next to its old storage and swapped in, so the +index is unchanged if `add` throws. As a result: + +- A call needs extra memory for the clusters it touches while it runs. +- Each call costs time proportional to the sizes of the touched clusters, on top of a + cost per new vector. Routed vectors usually touch most clusters, so a call can + still copy most of the index. **Add many vectors per call rather than one at a + time.** -- It needs memory for both copies while it runs. -- Each call has a fixed cost that grows with the size of the whole index, on top of a - cost per new vector. **Add many vectors per call rather than one at a time.** +A rebuilt cluster gets exactly the room it needs. To grow a cluster without +copying it, call `reserve` first: later `add` calls then write into the reserved +room in place until the cluster outgrows it. The reserved room is memory the index +holds but does not use yet. It is not saved, so a loaded index has none. -There is no spare capacity: every `add` moves the data. Measure update latency and peak memory with your index size, bit width and batch size; serialized file size alone does not determine peak process memory. diff --git a/docs/docs/index/usage.md b/docs/docs/index/usage.md index c9062254..c334ef64 100644 --- a/docs/docs/index/usage.md +++ b/docs/docs/index/usage.md @@ -102,7 +102,7 @@ construction, updates, and file I/O. The following contract applies to | Operation on the same index | Concurrent access | | --- | --- | | `search`, `search_batch`, `save`, and property reads | May run together | -| `build`, `add`, `remove`, and `resize`, where available | Require exclusive access | +| `build`, `add`, `remove`, `reserve`, and `resize`, where available | Require exclusive access | A conflicting call immediately raises `RuntimeError("Index is busy: conflicting operation in progress")`. diff --git a/include/rabitqlib/index/ivf/ivf.hpp b/include/rabitqlib/index/ivf/ivf.hpp index b3362066..08b53085 100644 --- a/include/rabitqlib/index/ivf/ivf.hpp +++ b/include/rabitqlib/index/ivf/ivf.hpp @@ -30,17 +30,52 @@ class Initializer; class IVF { private: + friend struct IvfTestAccess; + using ByteStorage = std::vector>; using FloatStorage = std::vector>; using IdStorage = std::vector>; + // One cluster's own storage, independent of every other cluster's. + class ClusterStorage { + public: + explicit ClusterStorage(bool raw_reranking) : raw_reranking_(raw_reranking) {} + + ByteStorage batch_storage; // 1-bit code and factors + ByteStorage ex_storage; // extra-bit codes and packed factors + FloatStorage raw_storage; // original vectors for raw reranking + IdStorage id_storage; // point IDs of this cluster + + [[nodiscard]] char* batch_data() { + return reinterpret_cast(batch_storage.data()); + } + + [[nodiscard]] const char* batch_data() const { + return reinterpret_cast(batch_storage.data()); + } + + [[nodiscard]] char* ex_data() { + return raw_reranking_ ? reinterpret_cast(raw_storage.data()) + : reinterpret_cast(ex_storage.data()); + } + + [[nodiscard]] const char* ex_data() const { + return raw_reranking_ ? reinterpret_cast(raw_storage.data()) + : reinterpret_cast(ex_storage.data()); + } + + [[nodiscard]] PID* ids() { return id_storage.data(); } + + [[nodiscard]] const PID* ids() const { return id_storage.data(); } + + private: + bool raw_reranking_; // raw vectors replace extra-bit codes + }; + InitializerType initializer_type_ = InitializerType::Auto; - std::unique_ptr initializer_; // initializer for candidate clusters - ByteStorage batch_storage_; // 1-bit code and factors - ByteStorage ex_storage_; // extra-bit codes and packed factors - FloatStorage raw_storage_; // original vectors for raw reranking - IdStorage id_storage_; // point IDs organized by cluster + std::unique_ptr initializer_; // initializer for candidate clusters + std::vector cluster_storage_; // one entry per cluster size_t max_elements_ = 0; // Logical capacity; construct fills it, add grows it. size_t num_removed_ = 0; // Recovered from removal markers when loading. size_t dim_ = 0; @@ -64,42 +99,29 @@ class IVF { quantize_cluster(Cluster&, const std::vector&, const float*, const float*, float*, float*, const quant::RabitqConfig&); void - extend_cluster(const Cluster&, const std::vector&, const float*, const float*, PID, char*, char*, PID*, const quant::RabitqConfig&) + extend_cluster(const Cluster&, const std::vector&, const float*, const float*, PID, char*, char*, PID*, char*, const quant::RabitqConfig&) const; // True if `n` points cannot be addressed by the storage size computations [[nodiscard]] bool exceeds_storage(size_t n) const; - [[nodiscard]] size_t ids_bytes() const { return sizeof(PID) * max_elements_; } + [[nodiscard]] static size_t ids_bytes(size_t cluster_size) { + return sizeof(PID) * cluster_size; + } // get num of bytes used for 1-bit code and corresponding factors - [[nodiscard]] size_t batch_data_bytes(const std::vector& cluster_sizes) const; + [[nodiscard]] size_t batch_data_bytes(size_t cluster_size) const; [[nodiscard]] inline size_t rerank_vector_bytes() const; - [[nodiscard]] inline size_t ex_data_bytes() const; + [[nodiscard]] inline size_t ex_data_bytes(size_t cluster_size) const; - [[nodiscard]] char* batch_data() { - return reinterpret_cast(batch_storage_.data()); - } + // Size every buffer of a cluster's storage for `n` points + void resize_storage(ClusterStorage& cs, size_t n) const; - [[nodiscard]] const char* batch_data() const { - return reinterpret_cast(batch_storage_.data()); - } - - [[nodiscard]] char* ex_data() { - return raw_reranking_ ? reinterpret_cast(raw_storage_.data()) - : reinterpret_cast(ex_storage_.data()); - } - - [[nodiscard]] const char* ex_data() const { - return raw_reranking_ ? reinterpret_cast(raw_storage_.data()) - : reinterpret_cast(ex_storage_.data()); - } - - [[nodiscard]] PID* ids() { return id_storage_.data(); } + void reserve_storage(ClusterStorage& cs, size_t n) const; - [[nodiscard]] const PID* ids() const { return id_storage_.data(); } + [[nodiscard]] bool has_capacity(const ClusterStorage& cs, size_t n) const; void allocate_memory(const std::vector&); @@ -166,9 +188,11 @@ class IVF { * data the centroids were trained on. Point `i` of `data` receives PID * `num_stored_points() + i`. The index is unchanged if this throws. * - * The grown index is built next to the old one, so a call needs memory for both and - * takes time proportional to the whole index, not just to `n`: prefer a few large - * calls to many small ones. `construct` expects `max_elements()` rows afterwards. + * Only clusters that receive points are touched. Each one is rebuilt next to its old + * storage, so a call needs extra memory for, and takes time proportional to, those + * clusters plus `n`; prefer a few large calls to many small ones. A cluster with + * enough capacity from `reserve` grows in place instead. `construct` expects + * `max_elements()` rows afterwards. * Not safe to call while another thread searches, adds or removes. * * @param data New points (n * dim) @@ -187,6 +211,20 @@ class IVF { size_t num_threads = std::numeric_limits::max() ); + /** + * @brief Preallocate room for a cluster to hold `num_points` points. + * + * Later `add` calls grow the cluster in place, without copying it, until it holds more + * than `num_points`. Does nothing if the cluster already has that capacity. Reserved + * capacity is not saved: a loaded index has none. The index is unchanged if this + * throws. Not safe to call while another thread searches, adds or removes. + * + * @param cluster_id Cluster to grow + * @param num_points Total number of points the cluster can hold, including its + * current ones + */ + void reserve(PID cluster_id, size_t num_points); + /** * @brief Exclude points from all later search results. * diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index 2ee059d0..a2eef5f2 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/python_bindings/ivf_bindings.cpp @@ -164,6 +164,23 @@ class IvfIndex { return ids; } + void reserve(int64_t cluster_id, int64_t num_points) { + const auto operation = access_.write(); + if (!built_) { + throw std::runtime_error("IvfIndex must be built or loaded before reserve"); + } + if (cluster_id < 0 || static_cast(cluster_id) >= index_->num_clusters()) { + throw std::invalid_argument("cluster_id must be in [0, num_clusters)"); + } + if (num_points < 0) { + throw std::invalid_argument("num_points must be non-negative"); + } + py::gil_scoped_release release; + index_->reserve( + static_cast(cluster_id), static_cast(num_points) + ); + } + size_t remove(py::handle ids) { const auto operation = access_.write(); auto ids_array = ensure_1d_integer_array(ids, "ids"); @@ -347,7 +364,16 @@ void register_ivf(py::module_& m) { "Append vectors without the original data and return their ids. Vectors " "are quantized against the existing centroids, which do not move. Without " "cluster_ids each vector goes to its nearest centroid. Each call copies " - "the whole index, so add many vectors per call rather than one at a time." + "every cluster it adds to, unless reserve left room for the new vectors, " + "so add many vectors per call rather than one at a time." + ) + .def( + "reserve", + &IvfIndex::reserve, + py::arg("cluster_id"), + py::arg("num_points"), + "Preallocate room for a cluster to hold num_points vectors, so later adds " + "to it do not copy it. Reserved room is not saved." ) .def( "remove", diff --git a/src/index/ivf.cpp b/src/index/ivf.cpp index 68267a47..f0564836 100644 --- a/src/index/ivf.cpp +++ b/src/index/ivf.cpp @@ -52,13 +52,9 @@ bool IVF::exceeds_storage(size_t n) const { n > max_size / batch_bytes || (rerank_bytes != 0 && n > max_size / rerank_bytes); } -size_t IVF::batch_data_bytes(const std::vector& cluster_sizes) const { - assert(cluster_sizes.size() == num_clusters_); - size_t total_blocks = 0; - for (auto size : cluster_sizes) { - total_blocks += div_round_up(size, fastscan::kBatchSize); - } - return total_blocks * BatchDataMap::data_bytes(padded_dim_); +size_t IVF::batch_data_bytes(size_t cluster_size) const { + return div_round_up(cluster_size, fastscan::kBatchSize) * + BatchDataMap::data_bytes(padded_dim_); } // Keep the original inlining intent inside this translation unit, especially the @@ -68,24 +64,55 @@ inline size_t IVF::rerank_vector_bytes() const { : ExDataMap::data_bytes(padded_dim_, ex_bits_); } -inline size_t IVF::ex_data_bytes() const { return rerank_vector_bytes() * max_elements_; } +inline size_t IVF::ex_data_bytes(size_t cluster_size) const { + return rerank_vector_bytes() * cluster_size; +} + +void IVF::resize_storage(ClusterStorage& cs, size_t n) const { + cs.batch_storage.resize(batch_data_bytes(n)); + if (rerank_vector_bytes() > 0) { + if (raw_reranking_) { + cs.raw_storage.resize(n * dim_); + } else { + cs.ex_storage.resize(ex_data_bytes(n)); + } + } + cs.id_storage.resize(n); +} + +void IVF::reserve_storage(ClusterStorage& cs, size_t n) const { + cs.batch_storage.reserve(batch_data_bytes(n)); + if (rerank_vector_bytes() > 0) { + if (raw_reranking_) { + cs.raw_storage.reserve(n * dim_); + } else { + cs.ex_storage.reserve(ex_data_bytes(n)); + } + } + cs.id_storage.reserve(n); +} + +bool IVF::has_capacity(const ClusterStorage& cs, size_t n) const { + if (cs.batch_storage.capacity() < batch_data_bytes(n) || cs.id_storage.capacity() < n) { + return false; + } + if (rerank_vector_bytes() == 0) { + return true; + } + return raw_reranking_ ? cs.raw_storage.capacity() >= n * dim_ + : cs.ex_storage.capacity() >= ex_data_bytes(n); +} void IVF::free_memory() { initializer_.reset(); - ByteStorage().swap(batch_storage_); - ByteStorage().swap(ex_storage_); - FloatStorage().swap(raw_storage_); - IdStorage().swap(id_storage_); + std::vector().swap(cluster_storage_); } void IVF::swap(IVF& other) noexcept { using std::swap; swap(initializer_type_, other.initializer_type_); swap(initializer_, other.initializer_); - swap(batch_storage_, other.batch_storage_); - swap(ex_storage_, other.ex_storage_); - swap(raw_storage_, other.raw_storage_); - swap(id_storage_, other.id_storage_); + swap(cluster_storage_, other.cluster_storage_); swap(max_elements_, other.max_elements_); swap(num_removed_, other.num_removed_); swap(dim_, other.dim_); @@ -250,42 +277,25 @@ void IVF::allocate_memory(const std::vector& cluster_sizes) { this->initializer_ = std::make_unique(padded_dim_, num_clusters_, metric_type_); } - batch_storage_ = ByteStorage(batch_data_bytes(cluster_sizes)); - if (rerank_vector_bytes() > 0) { - if (raw_reranking_) { - raw_storage_ = FloatStorage(max_elements_ * dim_); - } else { - ex_storage_ = ByteStorage(ex_data_bytes()); - } + cluster_storage_.assign(num_clusters_, ClusterStorage(raw_reranking_)); + for (size_t i = 0; i < num_clusters_; ++i) { + resize_storage(cluster_storage_[i], cluster_sizes[i]); } - id_storage_ = IdStorage(max_elements_); this->ip_func_ = select_excode_ipfunc(ex_bits_); } /** - * @brief Initialize the cluster list by finding each cluster's storage offsets. + * @brief Initialize the cluster list from each cluster's own independent storage. */ void IVF::init_clusters(const std::vector& cluster_sizes) { this->clusters_.reserve(num_clusters_); - size_t added_vectors = 0; - size_t added_batches = 0; for (size_t i = 0; i < num_clusters_; ++i) { - // find data location for current cluster - size_t num = cluster_sizes[i]; - size_t num_batches = div_round_up(num, fastscan::kBatchSize); - - char* current_batch_data = - batch_data() + (BatchDataMap::data_bytes(padded_dim_) * added_batches); - char* current_ex_data = rerank_vector_bytes() > 0 - ? ex_data() + (added_vectors * rerank_vector_bytes()) - : nullptr; - PID* cluster_ids = ids() + added_vectors; - - this->clusters_.emplace_back(num, current_batch_data, current_ex_data, cluster_ids); - - added_vectors += num; - added_batches += num_batches; + ClusterStorage& cs = cluster_storage_[i]; + char* current_ex_data = rerank_vector_bytes() > 0 ? cs.ex_data() : nullptr; + this->clusters_.emplace_back( + cluster_sizes[i], cs.batch_data(), current_ex_data, cs.ids() + ); } } @@ -344,8 +354,7 @@ void IVF::add( const float* data, size_t n, const PID* cluster_ids, bool faster, size_t num_threads ) { if (!ready_ || initializer_ == nullptr || rotator_ == nullptr || - clusters_.size() != num_clusters_ || batch_storage_.empty() || - id_storage_.size() != max_elements_) { + clusters_.size() != num_clusters_ || cluster_storage_.size() != num_clusters_) { throw std::logic_error("IVF index must be constructed or loaded before add"); } if (n == 0) { @@ -391,33 +400,12 @@ void IVF::add( rows_of[assigned[i]].push_back(static_cast(i)); } std::vector new_sizes(num_clusters_); + std::vector touched; for (size_t i = 0; i < num_clusters_; ++i) { new_sizes[i] = clusters_[i].num() + rows_of[i].size(); - } - - /* Build the grown storage on the side so the index is unchanged on failure */ - ByteStorage new_batch(batch_data_bytes(new_sizes)); - ByteStorage new_ex; - FloatStorage new_raw; - const size_t rerank_bytes = rerank_vector_bytes(); - if (raw_reranking_) { - new_raw = FloatStorage(new_num_points * dim_); - } else if (rerank_bytes > 0) { - new_ex = ByteStorage(new_num_points * rerank_bytes); - } - IdStorage new_ids(new_num_points); - char* new_rerank = raw_reranking_ ? reinterpret_cast(new_raw.data()) - : reinterpret_cast(new_ex.data()); - - std::vector batch_offset(num_clusters_); - std::vector vector_offset(num_clusters_); - size_t added_batches = 0; - size_t added_vectors = 0; - for (size_t i = 0; i < num_clusters_; ++i) { - batch_offset[i] = added_batches; - vector_offset[i] = added_vectors; - added_batches += div_round_up(new_sizes[i], fastscan::kBatchSize); - added_vectors += new_sizes[i]; + if (!rows_of[i].empty()) { + touched.push_back(i); + } } quant::RabitqConfig config; @@ -425,35 +413,80 @@ void IVF::add( config = quant::faster_config(padded_dim_, ex_bits_ + 1); } + const size_t rerank_bytes = rerank_vector_bytes(); const size_t batch_bytes = BatchDataMap::data_bytes(padded_dim_); const PID first_pid = static_cast(max_elements_); + + /* Rebuild clusters without reserved capacity on the side; the rest grow in place */ + std::vector> staged(touched.size()); + std::vector tails(touched.size()); // repacked tails of in-place clusters + for (size_t t = 0; t < touched.size(); ++t) { + const size_t i = touched[t]; + if (has_capacity(cluster_storage_[i], new_sizes[i])) { + const size_t full_batches = clusters_[i].num() / fastscan::kBatchSize; + tails[t].resize(batch_data_bytes(new_sizes[i]) - (full_batches * batch_bytes)); + } else { + staged[t].emplace(raw_reranking_); + resize_storage(*staged[t], new_sizes[i]); + } + } + // Growing within capacity cannot reallocate or throw, so old cluster views stay valid + for (size_t t = 0; t < touched.size(); ++t) { + if (!staged[t]) { + resize_storage(cluster_storage_[touched[t]], new_sizes[touched[t]]); + } + } + // An exception must not leave an OpenMP region, so keep the first one and rethrow it. index_detail::ParallelException failure; #pragma omp parallel for schedule(dynamic) num_threads(resolve_num_threads(num_threads)) - for (std::ptrdiff_t index = 0; index < static_cast(num_clusters_); + for (std::ptrdiff_t index = 0; index < static_cast(touched.size()); ++index) { - const size_t i = static_cast(index); + const auto t = static_cast(index); + const size_t i = touched[t]; try { + ClusterStorage& cs = staged[t] ? *staged[t] : cluster_storage_[i]; + const size_t full_batches = clusters_[i].num() / fastscan::kBatchSize; + char* tail = staged[t] ? cs.batch_data() + (full_batches * batch_bytes) + : reinterpret_cast(tails[t].data()); extend_cluster( clusters_[i], rows_of[i], data, initializer_->centroid(static_cast(i)), first_pid, - reinterpret_cast(new_batch.data()) + (batch_offset[i] * batch_bytes), - rerank_bytes > 0 ? new_rerank + (vector_offset[i] * rerank_bytes) : nullptr, - new_ids.data() + vector_offset[i], + cs.batch_data(), + rerank_bytes > 0 ? cs.ex_data() : nullptr, + cs.ids(), + tail, config ); } catch (...) { failure.capture(); } } - failure.rethrow(); + if (failure.failed()) { + // In-place clusters were only written past their old size; shrinking drops that + for (size_t t = 0; t < touched.size(); ++t) { + if (!staged[t]) { + resize_storage(cluster_storage_[touched[t]], clusters_[touched[t]].num()); + } + } + failure.rethrow(); + } - /* Commit. Rebuilding the cluster list reuses its capacity, so it cannot throw. */ - batch_storage_.swap(new_batch); - ex_storage_.swap(new_ex); - raw_storage_.swap(new_raw); - id_storage_.swap(new_ids); + /* Commit. Nothing below can throw. */ + for (size_t t = 0; t < touched.size(); ++t) { + const size_t i = touched[t]; + if (staged[t]) { + cluster_storage_[i] = std::move(*staged[t]); + } else { + const size_t full_batches = clusters_[i].num() / fastscan::kBatchSize; + std::copy_n( + tails[t].data(), + tails[t].size(), + cluster_storage_[i].batch_storage.data() + (full_batches * batch_bytes) + ); + } + } max_elements_ = new_num_points; clusters_.clear(); init_clusters(new_sizes); @@ -462,9 +495,10 @@ void IVF::add( /** * @brief Copy one cluster into its grown storage and append new points to it. * - * Full batches, rerank data and IDs are copied unchanged. The old partial tail batch is - * unpacked and repacked together with the new points, so the result is byte-identical to - * quantizing the same points in the same order in one construct call. + * Full batches, rerank data and IDs are copied unchanged, unless they are already in place. + * The old partial tail batch is unpacked and repacked together with the new points into + * `tail_dst`, so the result is byte-identical to quantizing the same points in the same + * order in one construct call. */ void IVF::extend_cluster( const Cluster& old_cluster, @@ -475,28 +509,28 @@ void IVF::extend_cluster( char* batch_dst, char* rerank_dst, PID* ids_dst, + char* tail_dst, const quant::RabitqConfig& config ) const { + assert(!rows.empty()); // add() only extends clusters that receive new points const size_t old_num = old_cluster.num(); const size_t batch_bytes = BatchDataMap::data_bytes(padded_dim_); const size_t rerank_bytes = rerank_vector_bytes(); const size_t cols = padded_dim_ / 8; - if (rows.empty()) { - const size_t old_batches = div_round_up(old_num, fastscan::kBatchSize); - std::copy_n(old_cluster.batch_data(), old_batches * batch_bytes, batch_dst); - std::copy_n(old_cluster.ex_data(), old_num * rerank_bytes, rerank_dst); - std::copy_n(old_cluster.ids(), old_num, ids_dst); - return; - } - const size_t full_batches = old_num / fastscan::kBatchSize; const size_t kept = old_num - (full_batches * fastscan::kBatchSize); const size_t tail_num = kept + rows.size(); - std::copy_n(old_cluster.batch_data(), full_batches * batch_bytes, batch_dst); - std::copy_n(old_cluster.ex_data(), old_num * rerank_bytes, rerank_dst); - std::copy_n(old_cluster.ids(), old_num, ids_dst); + if (batch_dst != old_cluster.batch_data()) { + std::copy_n(old_cluster.batch_data(), full_batches * batch_bytes, batch_dst); + } + if (rerank_dst != old_cluster.ex_data()) { + std::copy_n(old_cluster.ex_data(), old_num * rerank_bytes, rerank_dst); + } + if (ids_dst != old_cluster.ids()) { + std::copy_n(old_cluster.ids(), old_num, ids_dst); + } /* Codes and factors of every point in the tail, old points first */ std::vector codes(tail_num * cols); @@ -546,7 +580,7 @@ void IVF::extend_cluster( } /* Repack the tail into batches */ - char* batch = batch_dst + (full_batches * batch_bytes); + char* batch = tail_dst; for (size_t start = 0; start < tail_num; start += fastscan::kBatchSize) { const size_t count = std::min(fastscan::kBatchSize, tail_num - start); BatchDataMap cur_batch(batch, padded_dim_); @@ -562,10 +596,49 @@ void IVF::extend_cluster( } } +void IVF::reserve(PID cluster_id, size_t num_points) { + if (!ready_ || clusters_.size() != num_clusters_ || + cluster_storage_.size() != num_clusters_) { + throw std::logic_error("IVF index must be constructed or loaded before reserve"); + } + if (cluster_id >= num_clusters_) { + throw std::invalid_argument("Cluster ID is out of range"); + } + if (num_points > buffer::kSearchBufferMaxPointCount) { + throw std::invalid_argument("IVF point count exceeds the supported ID range"); + } + if (exceeds_storage(num_points)) { + throw std::invalid_argument("IVF configuration exceeds addressable storage"); + } + ClusterStorage& cs = cluster_storage_[cluster_id]; + if (has_capacity(cs, num_points)) { + return; + } + + const Cluster& old_cluster = clusters_[cluster_id]; + const size_t num = old_cluster.num(); + ClusterStorage grown(raw_reranking_); + reserve_storage(grown, num_points); + resize_storage(grown, num); + std::copy_n(old_cluster.batch_data(), batch_data_bytes(num), grown.batch_data()); + if (rerank_vector_bytes() > 0) { + std::copy_n(old_cluster.ex_data(), ex_data_bytes(num), grown.ex_data()); + } + std::copy_n(old_cluster.ids(), num, grown.ids()); + std::vector cluster_sizes(num_clusters_); + for (size_t i = 0; i < num_clusters_; ++i) { + cluster_sizes[i] = clusters_[i].num(); + } + + /* Commit. Rebuilding the cluster list reuses its capacity, so it cannot throw. */ + cs = std::move(grown); + clusters_.clear(); + init_clusters(cluster_sizes); +} + size_t IVF::remove(const PID* ids_to_remove, size_t n) { if (!ready_ || initializer_ == nullptr || rotator_ == nullptr || - clusters_.size() != num_clusters_ || batch_storage_.empty() || - id_storage_.size() != max_elements_) { + clusters_.size() != num_clusters_ || cluster_storage_.size() != num_clusters_) { throw std::logic_error("IVF index must be constructed or loaded before remove"); } if (n == 0) { @@ -612,8 +685,7 @@ size_t IVF::remove(const PID* ids_to_remove, size_t n) { void IVF::save(const char* filename) const { if (!ready_ || initializer_ == nullptr || rotator_ == nullptr || - clusters_.size() != num_clusters_ || batch_storage_.empty() || - id_storage_.size() != max_elements_) { + clusters_.size() != num_clusters_ || cluster_storage_.size() != num_clusters_) { throw std::logic_error("Cannot save an unconstructed IVF index"); } if (filename == nullptr || filename[0] == '\0') { @@ -658,17 +730,28 @@ void IVF::save(const char* filename) const { /* Save rotator */ this->rotator_->save(output); - /* Save data */ + /* Save data: all batch data, then all ex data, then all IDs, each in cluster order */ this->initializer_->save(output, filename); - output.write( - batch_data(), static_cast(batch_data_bytes(cluster_sizes)) - ); - if (ex_data_bytes() != 0) { - output.write(ex_data(), static_cast(ex_data_bytes())); + for (size_t i = 0; i < num_clusters_; ++i) { + output.write( + cluster_storage_[i].batch_data(), + static_cast(batch_data_bytes(cluster_sizes[i])) + ); + } + if (rerank_vector_bytes() != 0) { + for (size_t i = 0; i < num_clusters_; ++i) { + output.write( + cluster_storage_[i].ex_data(), + static_cast(ex_data_bytes(cluster_sizes[i])) + ); + } + } + for (size_t i = 0; i < num_clusters_; ++i) { + output.write( + reinterpret_cast(cluster_storage_[i].ids()), + static_cast(ids_bytes(cluster_sizes[i])) + ); } - output.write( - reinterpret_cast(ids()), static_cast(ids_bytes()) - ); output.close(); } @@ -813,24 +896,34 @@ void IVF::load(const char* filename) { /* Load rotator */ loaded.rotator_->load(input); - /* Load data */ + /* Load data: all batch data, then all ex data, then all IDs, each in cluster order */ loaded.allocate_memory(cluster_sizes); loaded.initializer_->load(input, filename); - input.read( - loaded.batch_data(), - static_cast(loaded.batch_data_bytes(cluster_sizes)) - ); - if (loaded.ex_data_bytes() != 0) { - input.read(loaded.ex_data(), static_cast(loaded.ex_data_bytes())); + for (size_t i = 0; i < loaded.num_clusters_; ++i) { + input.read( + loaded.cluster_storage_[i].batch_data(), + static_cast(loaded.batch_data_bytes(cluster_sizes[i])) + ); } - input.read( - reinterpret_cast(loaded.ids()), - static_cast(loaded.ids_bytes()) - ); - for (size_t i = 0; i < loaded.max_elements_; ++i) { - if (loaded.ids()[i] >= loaded.max_elements_ || - loaded.ids()[i] >= buffer::kSearchBufferCheckedMask) { - throw std::runtime_error("Invalid point ID in IVF index file"); + if (loaded.rerank_vector_bytes() != 0) { + for (size_t i = 0; i < loaded.num_clusters_; ++i) { + input.read( + loaded.cluster_storage_[i].ex_data(), + static_cast(loaded.ex_data_bytes(cluster_sizes[i])) + ); + } + } + for (size_t i = 0; i < loaded.num_clusters_; ++i) { + input.read( + reinterpret_cast(loaded.cluster_storage_[i].ids()), + static_cast(ids_bytes(cluster_sizes[i])) + ); + } + for (size_t i = 0; i < loaded.num_clusters_; ++i) { + for (PID id : loaded.cluster_storage_[i].id_storage) { + if (id >= loaded.max_elements_ || id >= buffer::kSearchBufferCheckedMask) { + throw std::runtime_error("Invalid point ID in IVF index file"); + } } } @@ -993,8 +1086,7 @@ inline void IVF::validate_search( const float* query, size_t k, size_t nprobe, const PID* results ) const { if (!ready_ || initializer_ == nullptr || rotator_ == nullptr || - clusters_.size() != num_clusters_ || batch_storage_.empty() || - id_storage_.size() != max_elements_) { + clusters_.size() != num_clusters_ || cluster_storage_.size() != num_clusters_) { throw std::logic_error("IVF index must be constructed or loaded before search"); } if (query == nullptr) { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 84ab6e80..fe351aa6 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -42,7 +42,7 @@ if(RABITQ_BUILD_INDEX) endif() set_source_files_properties(unit/rabitqlib/index/historical_persistence_test.cpp PROPERTIES COMPILE_DEFINITIONS - "RABITQ_TEST_FIXTURE_DIR=\"${CMAKE_CURRENT_SOURCE_DIR}/fixtures/persistence/v0.5.2\"") + "RABITQ_TEST_FIXTURE_ROOT=\"${CMAKE_CURRENT_SOURCE_DIR}/fixtures/persistence\"") rabitq_configure_target(rabitq_tests) # Link google test and rabitqlib headers diff --git a/tests/fixtures/persistence/pre-ivf-per-cluster-storage/CMakeLists.txt b/tests/fixtures/persistence/pre-ivf-per-cluster-storage/CMakeLists.txt new file mode 100644 index 00000000..e5838291 --- /dev/null +++ b/tests/fixtures/persistence/pre-ivf-per-cluster-storage/CMakeLists.txt @@ -0,0 +1,13 @@ +cmake_minimum_required(VERSION 3.20) +project(rabitq_pre_ivf_per_cluster_storage_fixtures LANGUAGES CXX) + +if(NOT IS_DIRECTORY "${RABITQ_HISTORICAL_SOURCE}") + message(FATAL_ERROR "Set RABITQ_HISTORICAL_SOURCE to the pinned pre-refactor checkout") +endif() +set(RABITQ_BUILD_TESTS OFF CACHE BOOL "" FORCE) +set(RABITQ_BUILD_SAMPLES OFF CACHE BOOL "" FORCE) +set(RABITQ_BUILD_PYTHON_BINDINGS OFF CACHE BOOL "" FORCE) +set(RABITQ_ENABLE_NATIVE_OPTIMIZATION OFF CACHE BOOL "" FORCE) +add_subdirectory("${RABITQ_HISTORICAL_SOURCE}" historical) +add_executable(generate_fixtures generate.cpp) +target_link_libraries(generate_fixtures PRIVATE rabitqlib::index) diff --git a/tests/fixtures/persistence/pre-ivf-per-cluster-storage/README.md b/tests/fixtures/persistence/pre-ivf-per-cluster-storage/README.md new file mode 100644 index 00000000..b270766b --- /dev/null +++ b/tests/fixtures/persistence/pre-ivf-per-cluster-storage/README.md @@ -0,0 +1,37 @@ +# IVF fixtures from before per-cluster storage + +These files were written by `main` at commit +`64fc943e0e92fa7e18eec1477050b23342a976f8`, the last commit where IVF kept every +cluster in one shared buffer, using [generate.cpp](generate.cpp). +[manifest.json](manifest.json) records the producer and SHA-256 hashes. +Keep these files unchanged; tests never regenerate them with the current writer. + +The two cases cover four-bit and raw-vector (`nbits=32`) IVF for L2. +Each index holds 130 points in 65 original dimensions (96 padded) across +four clusters of 0, 1, 32 and 97 points, so the file has an empty cluster, a +single-point cluster, a full FastScan batch and a partial tail batch. Point 0 is +removed. The `.expected` files hold reference results of a search that probes +every cluster. + +The tests load each file, compare search results, and require that saving it +again reproduces the original bytes. Because the clusters differ in size, a +reader or writer that orders the batch data, extra-bit data and IDs differently +from the shared-buffer layout fails these checks. + +## Deliberate regeneration + +Build the generator only against an isolated copy of the pinned commit: + +```bash +mkdir -p build/fixtures/pre-ivf-per-cluster-storage +git archive 64fc943 | tar -x -C build/fixtures/pre-ivf-per-cluster-storage +cmake -S tests/fixtures/persistence/pre-ivf-per-cluster-storage \ + -B build/fixtures/pre-ivf-generator \ + -DRABITQ_HISTORICAL_SOURCE="$PWD/build/fixtures/pre-ivf-per-cluster-storage" \ + -DCMAKE_BUILD_TYPE=Release +cmake --build build/fixtures/pre-ivf-generator --target generate_fixtures +``` + +Run `generate_fixtures` with a new output directory. Rotation is seeded from +`random_device`, so the bytes differ on every run; each index and its `.expected` +file form a fixed pair. Record new hashes in the manifest before adopting them. diff --git a/tests/fixtures/persistence/pre-ivf-per-cluster-storage/generate.cpp b/tests/fixtures/persistence/pre-ivf-per-cluster-storage/generate.cpp new file mode 100644 index 00000000..e27357c9 --- /dev/null +++ b/tests/fixtures/persistence/pre-ivf-per-cluster-storage/generate.cpp @@ -0,0 +1,67 @@ +// Build only against the pinned pre-refactor source, never the current checkout. +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/metric.hpp" +#include "rabitqlib/index/ivf/ivf.hpp" +#include "rabitqlib/index/types.hpp" + +int main(int argc, char** argv) { + if (argc != 2) { + throw std::invalid_argument("Pass the fixture output directory"); + } + using namespace rabitqlib; + // Cluster sizes 0, 1, 32 and 97 cover empty, single, full-batch and tail clusters. + constexpr size_t kCount = 130, kDim = 65, kClusters = 4, kQueries = 3, kTopK = 5; + const std::filesystem::path directory(argv[1]); + std::filesystem::create_directories(directory); + uint32_t state = 713; + const auto next = [&state] { + state = state * 1664525U + 1013904223U; + return static_cast(static_cast((state >> 8) % 2001) - 1000) / 256.0F; + }; + std::vector data(kCount * kDim), queries(kQueries * kDim); + std::vector centroids(kClusters * kDim); + for (auto& value : data) + value = next(); + for (auto& value : queries) + value = next(); + for (auto& value : centroids) + value = next(); + std::vector clusters(kCount, 3); + clusters[0] = 1; + for (size_t i = 1; i < 33; ++i) + clusters[i] = 2; + for (const std::string kind : {"ivf4", "ivf32"}) { + const std::string name = kind + "_l2"; + const auto path = (directory / (name + ".index")).string(); + std::vector ids(kQueries * kTopK); + std::vector distances(ids.size()); + ivf::IVF index(kCount, kDim, kClusters, kind == "ivf32" ? 32 : 4, METRIC_L2); + index.construct(data.data(), centroids.data(), clusters.data(), false, 1); + const PID removed = 0; + index.remove(&removed, 1); + index.save(path.c_str()); + index.search_batch( + queries.data(), kQueries, kTopK, kClusters, ids.data(), distances.data() + ); + std::ofstream expected(directory / (name + ".expected"), std::ios::binary); + expected.exceptions(std::ios::failbit | std::ios::badbit); + expected << std::setprecision(std::numeric_limits::max_digits10); + expected << kCount << ' ' << kDim << ' ' << kQueries << ' ' << kTopK << '\n'; + for (size_t i = 0; i < queries.size(); ++i) + expected << queries[i] << (i + 1 == queries.size() ? '\n' : ' '); + for (size_t i = 0; i < ids.size(); ++i) + expected << ids[i] << ' ' << distances[i] << '\n'; + expected.close(); + } + return 0; +} diff --git a/tests/fixtures/persistence/pre-ivf-per-cluster-storage/ivf32_l2.expected b/tests/fixtures/persistence/pre-ivf-per-cluster-storage/ivf32_l2.expected new file mode 100644 index 00000000..499fa616 --- /dev/null +++ b/tests/fixtures/persistence/pre-ivf-per-cluster-storage/ivf32_l2.expected @@ -0,0 +1,17 @@ +130 65 3 5 +0.7734375 -1.265625 2.28125 -3.40234375 -1.609375 0.91796875 2.7109375 -0.20703125 0.31640625 0.62890625 1.296875 2.8203125 -1.46484375 -0.171875 -1.53515625 2.46875 3.79296875 -3.61328125 0.52734375 -3.21484375 2.94921875 2.4453125 0.7265625 -1.44921875 2.546875 -2.00390625 -2.8125 -1.609375 -2.1953125 -3.625 2.00390625 -1.5 -0.57421875 3.296875 -2.828125 -3.30078125 -3.03515625 3.25 -2.89453125 0.54296875 1.21875 -2.16015625 -3.34375 2.84765625 -2.24609375 -1.34765625 -0.1484375 -3.7265625 -3.87109375 -0.2890625 1.7421875 -3.08203125 0.66015625 -1.3359375 1.62109375 2.5 -2.8828125 1.203125 -2.82421875 2.0703125 1.05078125 -2.21875 -0.7578125 -1.1796875 1.86328125 -1.00390625 -2.02734375 -2.69140625 3.0703125 -0.796875 -0.78515625 -2.34375 3.60546875 -1.81640625 -1.3671875 -0.80078125 3.63671875 2.00390625 -3.828125 1.8125 1.34765625 0.5234375 0.8359375 3.71484375 1.19140625 0.54296875 -0.06640625 3.25390625 1.68359375 -1.59375 3.03515625 1.22265625 -3.2734375 0.05859375 -0.84765625 -0.96484375 1.046875 -2.58984375 -0.3828125 -1.765625 0.00390625 3.5 0.1484375 0.234375 -1.9609375 -0.375 -3.37109375 -0.40234375 1.34765625 -3.12109375 1 -1.16015625 1.83984375 2.46484375 1.25 -1.3125 -3.7734375 -0.9765625 -1.390625 -0.16015625 -1.6484375 -2.6015625 2.40234375 3.03515625 2.64453125 -1.49609375 -1.9765625 -2.3515625 -1.57421875 -0.18359375 1.37109375 1.4765625 2.46875 -0.75390625 -0.6171875 1.484375 2.1328125 -0.73046875 -1.640625 3.1875 0.6484375 2.80859375 2.79296875 -1.578125 -3.72265625 -3.0546875 -3.57421875 -3.1953125 -3.51171875 -0.9921875 1.171875 0.7734375 -3.08984375 -0.07421875 1.25390625 1.55078125 -2.6171875 -0.5703125 0.21484375 -0.90234375 -0.65234375 0.44140625 1.62109375 0.76171875 2.58984375 -2.8046875 0.90625 -3.09375 -1.3515625 0.5625 1.84765625 1.8515625 2.03125 1.51171875 -3.25 -1.32421875 0.94140625 -1.0390625 2.3984375 -3.0546875 -3.3828125 3.109375 0.01171875 -1.3203125 -0.328125 -0.94140625 0.4296875 2.43359375 -1.5390625 -2.58203125 2.53515625 -2.578125 0.0703125 2.79296875 -1.49609375 +92 420.222473 +46 471.803406 +128 482.1716 +43 495.518188 +3 497.874329 +31 412.941284 +112 416.260925 +129 460.315308 +71 464.002625 +128 464.167297 +32 423.822571 +3 449.275635 +87 462.434082 +12 477.827637 +48 480.80365 diff --git a/tests/fixtures/persistence/pre-ivf-per-cluster-storage/ivf32_l2.index b/tests/fixtures/persistence/pre-ivf-per-cluster-storage/ivf32_l2.index new file mode 100644 index 00000000..aa50d1f3 Binary files /dev/null and b/tests/fixtures/persistence/pre-ivf-per-cluster-storage/ivf32_l2.index differ diff --git a/tests/fixtures/persistence/pre-ivf-per-cluster-storage/ivf4_l2.expected b/tests/fixtures/persistence/pre-ivf-per-cluster-storage/ivf4_l2.expected new file mode 100644 index 00000000..ba55d725 --- /dev/null +++ b/tests/fixtures/persistence/pre-ivf-per-cluster-storage/ivf4_l2.expected @@ -0,0 +1,17 @@ +130 65 3 5 +0.7734375 -1.265625 2.28125 -3.40234375 -1.609375 0.91796875 2.7109375 -0.20703125 0.31640625 0.62890625 1.296875 2.8203125 -1.46484375 -0.171875 -1.53515625 2.46875 3.79296875 -3.61328125 0.52734375 -3.21484375 2.94921875 2.4453125 0.7265625 -1.44921875 2.546875 -2.00390625 -2.8125 -1.609375 -2.1953125 -3.625 2.00390625 -1.5 -0.57421875 3.296875 -2.828125 -3.30078125 -3.03515625 3.25 -2.89453125 0.54296875 1.21875 -2.16015625 -3.34375 2.84765625 -2.24609375 -1.34765625 -0.1484375 -3.7265625 -3.87109375 -0.2890625 1.7421875 -3.08203125 0.66015625 -1.3359375 1.62109375 2.5 -2.8828125 1.203125 -2.82421875 2.0703125 1.05078125 -2.21875 -0.7578125 -1.1796875 1.86328125 -1.00390625 -2.02734375 -2.69140625 3.0703125 -0.796875 -0.78515625 -2.34375 3.60546875 -1.81640625 -1.3671875 -0.80078125 3.63671875 2.00390625 -3.828125 1.8125 1.34765625 0.5234375 0.8359375 3.71484375 1.19140625 0.54296875 -0.06640625 3.25390625 1.68359375 -1.59375 3.03515625 1.22265625 -3.2734375 0.05859375 -0.84765625 -0.96484375 1.046875 -2.58984375 -0.3828125 -1.765625 0.00390625 3.5 0.1484375 0.234375 -1.9609375 -0.375 -3.37109375 -0.40234375 1.34765625 -3.12109375 1 -1.16015625 1.83984375 2.46484375 1.25 -1.3125 -3.7734375 -0.9765625 -1.390625 -0.16015625 -1.6484375 -2.6015625 2.40234375 3.03515625 2.64453125 -1.49609375 -1.9765625 -2.3515625 -1.57421875 -0.18359375 1.37109375 1.4765625 2.46875 -0.75390625 -0.6171875 1.484375 2.1328125 -0.73046875 -1.640625 3.1875 0.6484375 2.80859375 2.79296875 -1.578125 -3.72265625 -3.0546875 -3.57421875 -3.1953125 -3.51171875 -0.9921875 1.171875 0.7734375 -3.08984375 -0.07421875 1.25390625 1.55078125 -2.6171875 -0.5703125 0.21484375 -0.90234375 -0.65234375 0.44140625 1.62109375 0.76171875 2.58984375 -2.8046875 0.90625 -3.09375 -1.3515625 0.5625 1.84765625 1.8515625 2.03125 1.51171875 -3.25 -1.32421875 0.94140625 -1.0390625 2.3984375 -3.0546875 -3.3828125 3.109375 0.01171875 -1.3203125 -0.328125 -0.94140625 0.4296875 2.43359375 -1.5390625 -2.58203125 2.53515625 -2.578125 0.0703125 2.79296875 -1.49609375 +92 416.866028 +46 474.57016 +128 486.291168 +43 489.386963 +3 500.983612 +112 412.081207 +31 417.133514 +14 449.370117 +128 459.453735 +90 460.22171 +32 427.267487 +3 453.46405 +87 462.349762 +127 486.125702 +106 489.006439 diff --git a/tests/fixtures/persistence/pre-ivf-per-cluster-storage/ivf4_l2.index b/tests/fixtures/persistence/pre-ivf-per-cluster-storage/ivf4_l2.index new file mode 100644 index 00000000..b606cf53 Binary files /dev/null and b/tests/fixtures/persistence/pre-ivf-per-cluster-storage/ivf4_l2.index differ diff --git a/tests/fixtures/persistence/pre-ivf-per-cluster-storage/manifest.json b/tests/fixtures/persistence/pre-ivf-per-cluster-storage/manifest.json new file mode 100644 index 00000000..9b91675f --- /dev/null +++ b/tests/fixtures/persistence/pre-ivf-per-cluster-storage/manifest.json @@ -0,0 +1,14 @@ +{ + "source_tag": "main", + "source_commit": "64fc943e0e92fa7e18eec1477050b23342a976f8", + "writer": "tests/fixtures/persistence/pre-ivf-per-cluster-storage/generate.cpp", + "producer": "Linux x86-64, g++ 14.3.0, Release, one thread, native optimization OFF", + "byte_order": "little", + "size_t_bytes": 8, + "sha256": { + "pre-ivf-per-cluster-storage/ivf32_l2.expected": "7843685eadb19c6023df353aa6602e56f5042e4d0c6c2b1ac6019949b7459734", + "pre-ivf-per-cluster-storage/ivf32_l2.index": "4b0f98a898edf9cf157d226c51ba74a313d5fb5c055dc5663e14ce3a631f121b", + "pre-ivf-per-cluster-storage/ivf4_l2.expected": "7de95886d52dfbb5300890f920256fd7927eae2a4528223753e6f44f76a3e20d", + "pre-ivf-per-cluster-storage/ivf4_l2.index": "eb73f9066fef993a5f7e833641adb8bddb1e95cdc04addb2261a8e10bec0c41f" + } +} diff --git a/tests/python/test_ivf.py b/tests/python/test_ivf.py index 163768c2..8261ba12 100644 --- a/tests/python/test_ivf.py +++ b/tests/python/test_ivf.py @@ -602,6 +602,41 @@ def test_add_chunks_match_a_single_add(tmp_path): assert a.read_bytes() == b.read_bytes() +@pytest.mark.parametrize("nbits", [1, 4, 32]) +def test_reserved_adds_match_a_single_add(tmp_path, nbits): + idx, old, new, assigned = _dynamic_index(nbits=nbits) + base = tmp_path / "base.index" + idx.save(str(base)) + whole = IvfIndex.load(str(base)) + whole.add(new, assigned) + parts = IvfIndex.load(str(base)) + for cluster in range(parts.num_clusters): + parts.reserve(cluster, len(old) + len(new)) + for chunk in np.array_split(np.arange(len(new)), [1, 33, 34, 90]): + parts.add(new[chunk], assigned[chunk]) + a, b = tmp_path / "a.index", tmp_path / "b.index" + whole.save(str(a)) + parts.save(str(b)) + assert a.read_bytes() == b.read_bytes() + + +def test_reserve_rejects_invalid_arguments(): + idx, old, _, _ = _dynamic_index() + for bad in (-1, 3, 2**32): + with pytest.raises(ValueError, match="cluster_id"): + idx.reserve(bad, 10) + with pytest.raises(ValueError, match="num_points"): + idx.reserve(0, -1) + with pytest.raises(ValueError, match="supported ID range"): + idx.reserve(0, 2**40) + idx.reserve(0, 0) + assert idx.max_elements == len(old) + + unbuilt = IvfIndex(65, 4, 1, 4) + with pytest.raises(RuntimeError, match="built or loaded"): + unbuilt.reserve(0, 1) + + def test_add_accepts_fast_quantization_and_other_dtypes(): idx, _, new, _ = _dynamic_index() idx.add(new[:10].astype(np.float64), fast_quantization=True) diff --git a/tests/unit/rabitqlib/index/historical_persistence_test.cpp b/tests/unit/rabitqlib/index/historical_persistence_test.cpp index f9b27684..645f0214 100644 --- a/tests/unit/rabitqlib/index/historical_persistence_test.cpp +++ b/tests/unit/rabitqlib/index/historical_persistence_test.cpp @@ -5,6 +5,9 @@ #include #include #include +#include +#include +#include #include #include #include @@ -18,17 +21,38 @@ namespace rabitqlib { namespace { -class HistoricalPersistenceTest : public ::testing::TestWithParam {}; +std::string read_bytes(const std::filesystem::path& path) { + std::ifstream input(path, std::ios::binary); + return {std::istreambuf_iterator(input), std::istreambuf_iterator()}; +} + +// padded_dim follows the padding rule of the writer that produced the fixture set. +struct HistoricalFixtureCase { + std::string directory; // subdirectory under RABITQ_TEST_FIXTURE_ROOT + std::string name; + size_t padded_dim; + size_t count = 33; + size_t ivf_nprobe = 1; + bool resave_identical = false; // the current writer must reproduce the file +}; + +void PrintTo(const HistoricalFixtureCase& value, std::ostream* output) { + *output << value.directory << '/' << value.name; +} + +class HistoricalPersistenceTest : public ::testing::TestWithParam {}; TEST_P(HistoricalPersistenceTest, LoadsSearchesResavesAndPreservesStateAfterFailedLoad) { - const std::filesystem::path directory(RABITQ_TEST_FIXTURE_DIR); - const auto name = GetParam(); + const std::filesystem::path directory = + std::filesystem::path(RABITQ_TEST_FIXTURE_ROOT) / GetParam().directory; + const auto name = GetParam().name; + const auto expected_padded_dim = GetParam().padded_dim; const auto path = (directory / (name + ".index")).string(); std::ifstream expected(directory / (name + ".expected")); ASSERT_TRUE(expected.is_open()); size_t count = 0, dim = 0, nq = 0, k = 0; expected >> count >> dim >> nq >> k; - ASSERT_EQ(count, 33U); + ASSERT_EQ(count, GetParam().count); ASSERT_EQ(dim, 65U); ASSERT_EQ(nq, 3U); ASSERT_EQ(k, 5U); @@ -39,8 +63,8 @@ TEST_P(HistoricalPersistenceTest, LoadsSearchesResavesAndPreservesStateAfterFail for (size_t i = 0; i < ids.size(); ++i) expected >> reference_ids[i] >> reference_distances[i]; ASSERT_TRUE(expected.good()); - const auto output = - std::filesystem::temp_directory_path() / ("rabitq_historical_" + name); + const auto output = std::filesystem::temp_directory_path() / + ("rabitq_historical_" + GetParam().directory + "_" + name); const auto compare = [&] { EXPECT_EQ(ids, reference_ids); for (size_t i = 0; i < ids.size(); ++i) { @@ -66,6 +90,9 @@ TEST_P(HistoricalPersistenceTest, LoadsSearchesResavesAndPreservesStateAfterFail search(); compare(); index.save(output.string().c_str()); + if (GetParam().resave_identical) { + EXPECT_TRUE(read_bytes(output) == read_bytes(path)); + } index.load(output.string().c_str()); check_counts(); search(); @@ -80,9 +107,11 @@ TEST_P(HistoricalPersistenceTest, LoadsSearchesResavesAndPreservesStateAfterFail if (name.find("ivf") == 0) { ivf::IVF index; exercise(index, [&] { - EXPECT_EQ(index.padded_dim(), 128U); + EXPECT_EQ(index.padded_dim(), expected_padded_dim); EXPECT_EQ(index.initializer_type(), ivf::InitializerType::Flat); - index.search_batch(queries.data(), nq, k, 1, ids.data(), distances.data()); + index.search_batch( + queries.data(), nq, k, GetParam().ivf_nprobe, ids.data(), distances.data() + ); }); } else if (name.find("hnsw") == 0) { hnsw::HierarchicalNSW index; @@ -99,7 +128,7 @@ TEST_P(HistoricalPersistenceTest, LoadsSearchesResavesAndPreservesStateAfterFail } else { symqg::QuantizedGraph index; exercise(index, [&] { - EXPECT_EQ(index.padded_dim(), 128U); + EXPECT_EQ(index.padded_dim(), expected_padded_dim); index.set_ef(64); index.search_batch(queries.data(), nq, k, ids.data(), distances.data(), 1); }); @@ -110,20 +139,31 @@ INSTANTIATE_TEST_SUITE_P( V052, HistoricalPersistenceTest, ::testing::Values( - "ivf4_l2", - "ivf4_ip", - "ivf32_l2", - "ivf32_ip", - "hnsw4_l2", - "hnsw4_ip", - "qg0_l2", - "qg0_ip", - "qg4_l2", - "qg4_ip", - "qg8_l2", - "qg8_ip" + HistoricalFixtureCase{"v0.5.2", "ivf4_l2", 128}, + HistoricalFixtureCase{"v0.5.2", "ivf4_ip", 128}, + HistoricalFixtureCase{"v0.5.2", "ivf32_l2", 128}, + HistoricalFixtureCase{"v0.5.2", "ivf32_ip", 128}, + HistoricalFixtureCase{"v0.5.2", "hnsw4_l2", 128}, + HistoricalFixtureCase{"v0.5.2", "hnsw4_ip", 128}, + HistoricalFixtureCase{"v0.5.2", "qg0_l2", 128}, + HistoricalFixtureCase{"v0.5.2", "qg0_ip", 128}, + HistoricalFixtureCase{"v0.5.2", "qg4_l2", 128}, + HistoricalFixtureCase{"v0.5.2", "qg4_ip", 128}, + HistoricalFixtureCase{"v0.5.2", "qg8_l2", 128}, + HistoricalFixtureCase{"v0.5.2", "qg8_ip", 128} + ), + [](const auto& info) { return info.param.name; } +); + +// Four uneven IVF clusters written before per-cluster storage; dim 65 pads to 96. +INSTANTIATE_TEST_SUITE_P( + PreIvfPerClusterStorage, + HistoricalPersistenceTest, + ::testing::Values( + HistoricalFixtureCase{"pre-ivf-per-cluster-storage", "ivf4_l2", 96, 130, 4, true}, + HistoricalFixtureCase{"pre-ivf-per-cluster-storage", "ivf32_l2", 96, 130, 4, true} ), - [](const auto& info) { return info.param; } + [](const auto& info) { return info.param.name; } ); } // namespace } // namespace rabitqlib diff --git a/tests/unit/rabitqlib/index/ivf_test.cpp b/tests/unit/rabitqlib/index/ivf_test.cpp index 715660a0..da258836 100644 --- a/tests/unit/rabitqlib/index/ivf_test.cpp +++ b/tests/unit/rabitqlib/index/ivf_test.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #include #include #include @@ -28,6 +29,22 @@ #include "rabitqlib/index/types.hpp" namespace rabitqlib::ivf { + +// Outside the anonymous namespace so it matches the friend declaration in ivf.hpp. +struct IvfTestAccess { + static Cluster cluster_view(const IVF& index, size_t i) { return index.clusters_[i]; } + + static bool has_spare_capacity(const IVF& index, size_t i) { + const auto& cs = index.cluster_storage_[i]; + return cs.batch_storage.capacity() > cs.batch_storage.size() || + cs.ex_storage.capacity() > cs.ex_storage.size() || + cs.raw_storage.capacity() > cs.raw_storage.size() || + cs.id_storage.capacity() > cs.id_storage.size(); + } + + static std::unique_ptr>& rotator(IVF& index) { return index.rotator_; } +}; + namespace { TEST(IvfSearchTest, BatchSearchMatchesIndividualQueries) { @@ -635,6 +652,30 @@ TEST(IvfAddTest, EveryWayOfAddingTheSamePointsGivesByteIdenticalIndexes) { ASSERT_EQ(done, kNew); EXPECT_TRUE(saved_bytes(chunked, path) == expected) << "chunked"; + // Even clusters are reserved and grow in place; odd ones are rebuilt each time. + IVF reserved; + reserved.load(base_path.c_str()); + std::vector final_sizes(kDynClusters, 0); + for (PID cluster : assigned) { + ++final_sizes[cluster]; + } + for (size_t i = 0; i < kDynClusters; i += 2) { + reserved.reserve(static_cast(i), final_sizes[i]); + } + const Cluster reserved_view = IvfTestAccess::cluster_view(reserved, 0); + done = 0; + for (size_t chunk : std::array{1, 31, 32, 7, 1, kNew - 72}) { + reserved.add( + extra + (done * kDynDim), chunk, assigned.data() + kOld + done, false, 1 + ); + done += chunk; + } + EXPECT_TRUE(saved_bytes(reserved, path) == expected) << "reserved"; + EXPECT_EQ( + IvfTestAccess::cluster_view(reserved, 0).batch_data(), + reserved_view.batch_data() + ); + IVF threaded; threaded.load(base_path.c_str()); threaded.add(extra, kNew, nullptr, false, 4); @@ -645,6 +686,224 @@ TEST(IvfAddTest, EveryWayOfAddingTheSamePointsGivesByteIdenticalIndexes) { std::remove(path.c_str()); } +TEST(IvfAddTest, AddingToOneClusterLeavesOtherClustersUntouched) { + // Cluster sizes 0, 1, 32, 97 cover empty, single, full-batch and tail clusters. + constexpr size_t kOld = 130; + const auto points = random_rows(kOld, 741); + const auto centroids = random_rows(kDynClusters, 742); + std::vector assigned(kOld, 3); + assigned[0] = 1; + std::fill(assigned.begin() + 1, assigned.begin() + 33, PID{2}); + + IVF index(kOld, kDynDim, kDynClusters, 4, METRIC_L2); + index.construct(points.data(), centroids.data(), assigned.data(), false, 1); + + std::vector before; + before.reserve(kDynClusters); + for (size_t i = 0; i < kDynClusters; ++i) { + before.push_back(IvfTestAccess::cluster_view(index, i)); + } + + // Route every new point to cluster 0 (the initially-empty one) explicitly. + constexpr size_t kNew = 5; + const auto extra = random_rows(kNew, 743); + const std::vector new_cluster_ids(kNew, 0); + index.add(extra.data(), kNew, new_cluster_ids.data(), false, 1); + + for (size_t i = 0; i < kDynClusters; ++i) { + SCOPED_TRACE(::testing::Message() << "cluster=" << i); + const Cluster after = IvfTestAccess::cluster_view(index, i); + if (i == 0) { + EXPECT_EQ(after.num(), before[i].num() + kNew); + } else { + // Untouched clusters keep their storage: no reallocation and no copy. + EXPECT_EQ(after.num(), before[i].num()); + EXPECT_EQ(after.batch_data(), before[i].batch_data()); + EXPECT_EQ(after.ex_data(), before[i].ex_data()); + EXPECT_EQ(after.ids(), before[i].ids()); + } + } +} + +TEST(IvfAddTest, UnreservedClustersGrowToExactlyTheirNewSize) { + constexpr size_t kOld = 40; + const auto points = random_rows(kOld + 1, 851); + const auto centroids = random_rows(kDynClusters, 852); + const std::vector assigned(kOld, 0); + + for (size_t bits : {1UL, 4UL, 32UL}) { + SCOPED_TRACE(::testing::Message() << "bits=" << bits); + IVF index(kOld, kDynDim, kDynClusters, bits, METRIC_L2); + index.construct(points.data(), centroids.data(), assigned.data(), false, 1); + const PID route_to_zero = 0; + index.add(points.data() + (kOld * kDynDim), 1, &route_to_zero, false, 1); + EXPECT_EQ(IvfTestAccess::cluster_view(index, 0).num(), kOld + 1); + EXPECT_FALSE(IvfTestAccess::has_spare_capacity(index, 0)); + } +} + +TEST(IvfReserveTest, ReservedClusterGrowsInPlaceUntilItsCapacityIsUsed) { + constexpr size_t kOld = 40; + constexpr size_t kReserved = kOld + 2; + const auto points = random_rows(kReserved + 1, 861); + const auto centroids = random_rows(kDynClusters, 862); + const std::vector assigned(kOld, 0); + const PID route_to_zero = 0; + + for (size_t bits : {1UL, 4UL, 32UL}) { + SCOPED_TRACE(::testing::Message() << "bits=" << bits); + IVF index(kOld, kDynDim, kDynClusters, bits, METRIC_L2); + index.construct(points.data(), centroids.data(), assigned.data(), false, 1); + index.reserve(0, kReserved); + const Cluster reserved = IvfTestAccess::cluster_view(index, 0); + EXPECT_EQ(reserved.num(), kOld); + + for (size_t row = kOld; row < kReserved; ++row) { + index.add(points.data() + (row * kDynDim), 1, &route_to_zero, false, 1); + const Cluster grown = IvfTestAccess::cluster_view(index, 0); + EXPECT_EQ(grown.num(), row + 1); + EXPECT_EQ(grown.batch_data(), reserved.batch_data()); + EXPECT_EQ(grown.ex_data(), reserved.ex_data()); + EXPECT_EQ(grown.ids(), reserved.ids()); + } + + // The old storage is alive while its replacement is allocated, so they differ. + index.add(points.data() + (kReserved * kDynDim), 1, &route_to_zero, false, 1); + EXPECT_NE( + IvfTestAccess::cluster_view(index, 0).batch_data(), reserved.batch_data() + ); + EXPECT_FALSE(IvfTestAccess::has_spare_capacity(index, 0)); + } +} + +TEST(IvfReserveTest, RejectsBadInputAndLeavesSavedBytesUnchanged) { + constexpr size_t kOld = 40; + const auto points = random_rows(kOld, 871); + const auto centroids = random_rows(kDynClusters, 872); + const std::vector assigned(kOld, 0); + const std::string path = ::testing::TempDir() + "rabitq_ivf_reserve.index"; + const auto expect_message = [](const auto& call, const char* message) { + try { + call(); + ADD_FAILURE() << "expected: " << message; + } catch (const std::exception& error) { EXPECT_STREQ(error.what(), message); } + }; + + IVF unbuilt(kOld, kDynDim, kDynClusters, 4); + EXPECT_THROW(unbuilt.reserve(0, 1), std::logic_error); + expect_message( + [&] { unbuilt.reserve(0, 1); }, + "IVF index must be constructed or loaded before reserve" + ); + + IVF index(kOld, kDynDim, kDynClusters, 4); + index.construct(points.data(), centroids.data(), assigned.data(), false, 1); + const std::string expected = saved_bytes(index, path); + const Cluster constructed = IvfTestAccess::cluster_view(index, 0); + + EXPECT_THROW(index.reserve(static_cast(kDynClusters), 1), std::invalid_argument); + expect_message( + [&] { index.reserve(static_cast(kDynClusters), 1); }, + "Cluster ID is out of range" + ); + EXPECT_THROW( + index.reserve(0, buffer::kSearchBufferMaxPointCount + 1), std::invalid_argument + ); + expect_message( + [&] { index.reserve(0, buffer::kSearchBufferMaxPointCount + 1); }, + "IVF point count exceeds the supported ID range" + ); + + // Asking for less than the cluster already holds does nothing. + index.reserve(0, kOld / 2); + EXPECT_EQ(IvfTestAccess::cluster_view(index, 0).batch_data(), constructed.batch_data()); + EXPECT_FALSE(IvfTestAccess::has_spare_capacity(index, 0)); + + // Spare capacity changes neither the file nor what a loaded index holds. + index.reserve(0, kOld * 2); + index.reserve(1, kOld); + EXPECT_TRUE(IvfTestAccess::has_spare_capacity(index, 0)); + EXPECT_TRUE(saved_bytes(index, path) == expected); + IVF loaded; + loaded.load(path.c_str()); + for (size_t i = 0; i < kDynClusters; ++i) { + EXPECT_FALSE(IvfTestAccess::has_spare_capacity(loaded, i)); + } + std::remove(path.c_str()); +} + +// Delegates to a real rotator, but the call numbered `fail_at` throws. +class FailingRotator : public Rotator { + public: + FailingRotator(std::unique_ptr> inner, size_t fail_at) + : Rotator(0, inner->size()), inner_(std::move(inner)), fail_at_(fail_at) {} + + void rotate(const float* src, float* dst) const override { + if (calls_++ == fail_at_) { + throw std::runtime_error("injected rotation failure"); + } + inner_->rotate(src, dst); + } + void load(std::ifstream& input) override { inner_->load(input); } + void save(std::ofstream& output) const override { inner_->save(output); } + void load(const char* data) override { inner_->load(data); } + void save(char* data) const override { inner_->save(data); } + [[nodiscard]] size_t dump_bytes() const override { return inner_->dump_bytes(); } + + std::unique_ptr> release() { return std::move(inner_); } + + private: + std::unique_ptr> inner_; + size_t fail_at_; + mutable size_t calls_ = 0; +}; + +TEST(IvfAddTest, FailedAddLeavesClustersGrowingInPlaceUnchanged) { + // Cluster 0 has a partial tail batch and grows in place; cluster 1 fails afterwards. + constexpr size_t kOld = 70; + constexpr size_t kInPlace = 3; + const auto points = random_rows(kOld + 6, 881); + const auto centroids = random_rows(kDynClusters, 882); + std::vector assigned(kOld, 1); + std::fill(assigned.begin(), assigned.begin() + 40, PID{0}); + const std::array new_clusters{0, 0, 0, 1, 1, 1}; + const float* extra = points.data() + (kOld * kDynDim); + const std::string base_path = ::testing::TempDir() + "rabitq_ivf_add_fail_base.index"; + const std::string path = ::testing::TempDir() + "rabitq_ivf_add_fail.index"; + + for (size_t bits : {1UL, 4UL, 32UL}) { + SCOPED_TRACE(::testing::Message() << "bits=" << bits); + IVF index(kOld, kDynDim, kDynClusters, bits, METRIC_L2); + index.construct(points.data(), centroids.data(), assigned.data(), false, 1); + index.save(base_path.c_str()); + index.reserve(0, 64); + const std::string expected = saved_bytes(index, path); + const Cluster before = IvfTestAccess::cluster_view(index, 0); + + auto& rotator = IvfTestAccess::rotator(index); + rotator = std::make_unique(std::move(rotator), kInPlace); + EXPECT_THROW( + index.add(extra, new_clusters.size(), new_clusters.data(), false, 1), + std::runtime_error + ); + rotator = static_cast(*rotator).release(); + + EXPECT_EQ(index.max_elements(), kOld); + EXPECT_EQ(IvfTestAccess::cluster_view(index, 0).num(), before.num()); + EXPECT_EQ(IvfTestAccess::cluster_view(index, 0).batch_data(), before.batch_data()); + EXPECT_TRUE(saved_bytes(index, path) == expected); + + // Retrying matches an index that never saw the failure. + index.add(extra, new_clusters.size(), new_clusters.data(), false, 1); + IVF reference; + reference.load(base_path.c_str()); + reference.add(extra, new_clusters.size(), new_clusters.data(), false, 1); + EXPECT_TRUE(saved_bytes(index, path) == saved_bytes(reference, path)); + } + std::remove(base_path.c_str()); + std::remove(path.c_str()); +} + TEST(IvfConstructionTest, RebuildMatchesIncrementalEncodingWithSkewedClustersAndTails) { // Cluster sizes 0, 1, 32, 97 cover empty rows and multiple physical tail batches. constexpr size_t kCount = 130;