IVF per cluster storage

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()

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 <chrono>
+#include <cstddef>
+#include <cstdint>
+#include <exception>
+#include <iostream>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+#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<rabitqlib::PID> base_assigned(base_n);
+    for (size_t i = 0; i < base_n; ++i) {
+        base_assigned[i] = static_cast<rabitqlib::PID>(i % clusters);
+    }
+    // "single" sends every new point to cluster 0; "routed" lets add() pick the nearest.
+    const std::vector<rabitqlib::PID> 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<double>(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<rabitqlib::PID> ids(10);
+        std::vector<float> 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;
+    }
+}

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.
 

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")`.

include/rabitqlib/index/ivf/ivf.hpp

@@ -30,17 +30,52 @@ class Initializer;
 
 class IVF {
    private:
+    friend struct IvfTestAccess;
+
     using ByteStorage =
         std::vector<std::byte, memory::AlignedAllocator<std::byte, 64, true>>;
     using FloatStorage = std::vector<float, memory::AlignedAllocator<float, 64, true>>;
     using IdStorage = std::vector<PID, memory::AlignedAllocator<PID, 64, true>>;
 
+    // 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<char*>(batch_storage.data());
+        }
+
+        [[nodiscard]] const char* batch_data() const {
+            return reinterpret_cast<const char*>(batch_storage.data());
+        }
+
+        [[nodiscard]] char* ex_data() {
+            return raw_reranking_ ? reinterpret_cast<char*>(raw_storage.data())
+                                  : reinterpret_cast<char*>(ex_storage.data());
+        }
+
+        [[nodiscard]] const char* ex_data() const {
+            return raw_reranking_ ? reinterpret_cast<const char*>(raw_storage.data())
+                                  : reinterpret_cast<const char*>(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_;  // 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_;     // initializer for candidate clusters
+    std::vector<ClusterStorage> 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<PID>&, const float*, const float*, float*, float*, const quant::RabitqConfig&);
 
     void
-    extend_cluster(const Cluster&, const std::vector<PID>&, const float*, const float*, PID, char*, char*, PID*, const quant::RabitqConfig&)
+    extend_cluster(const Cluster&, const std::vector<PID>&, 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<size_t>& 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<char*>(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<const char*>(batch_storage_.data());
-    }
-
-    [[nodiscard]] char* ex_data() {
-        return raw_reranking_ ? reinterpret_cast<char*>(raw_storage_.data())
-                              : reinterpret_cast<char*>(ex_storage_.data());
-    }
-
-    [[nodiscard]] const char* ex_data() const {
-        return raw_reranking_ ? reinterpret_cast<const char*>(raw_storage_.data())
-                              : reinterpret_cast<const char*>(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<size_t>&);
 
@@ -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<size_t>::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.
      *

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<uint64_t>(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<rabitqlib::PID>(cluster_id), static_cast<size_t>(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",