Add configurable IVF centroid routing with persisted initializer types

AGENTS.md

@@ -147,8 +147,13 @@ Recommended:
   pack/unpack, and estimation coverage. Test both `METRIC_L2` and `METRIC_IP` where supported.
 - IVF/HNSW quantized total bits are one sign bit plus `ex_bits`, with totals 1 through 9.
   IVF also accepts `bits == 32`: one-bit filtering plus owned original float32 vectors
-  in place of extra-bit codes. New IVF files have a magic/version header and an explicit
-  padded dimension; legacy quantized files and raw v1 files infer 64-dimension padding.
+  in place of extra-bit codes. New IVF files have a magic/version header, an explicit
+  padded dimension, and a resolved initializer type (v2); v1 files infer the historical
+  Flat/HNSW choice at 20,000 clusters, independently of current auto selection.
+  New IVF auto routing uses Flat below 5,000 clusters, Flat RaBitQ below 60,000,
+  and HNSW otherwise. Flat RaBitQ persists its mean and codes in the main file;
+  HNSW retains its sidecar. Legacy quantized files and raw v1 files infer
+  64-dimension padding.
   IVF search defaults to HACC for 4–9 bits and standard FastScan for 1–3 bits or raw storage.
   SymphonyQG supports
   raw storage (`quantization_bits == 0`) and quantized storage at 4 or 8 bits.

CMakeLists.txt

@@ -55,6 +55,7 @@ function(rabitq_enable_warnings target)
 endfunction()
 
 set(RABITQ_COMMON_SOURCES
+    src/index/ivf_initializer.cpp
     src/index/ivf_search.cpp
     src/index/qg.cpp
     src/utils/cpu_features.cpp

docs/docs/index/ivf.md

@@ -55,6 +55,56 @@ size_t k = centroids.rows();
 index_type ivf(num_points, dim, k, total_bits);
 ```
 
+### Centroid routing
+
+The optional final C++ constructor argument selects `InitializerType::Auto`,
+`Flat`, `FlatRaBitQ`, or `HNSW` from the `rabitqlib::ivf` namespace. Python accepts
+`initializer="auto"`, `"flat"`, `"flat_rabitq"`, or `"hnsw"`:
+
+```python
+index = IvfIndex(dim, num_points, num_clusters, nbits=4,
+                 initializer="flat_rabitq")
+print(index.initializer)  # resolved routing type
+```
+
+```c++
+index_type ivf(num_points, dim, k, total_bits, rabitqlib::METRIC_L2,
+               rabitqlib::RotatorType::FhtKacRotator,
+               rabitqlib::ivf::InitializerType::FlatRaBitQ);
+auto routing = ivf.initializer_type();
+```
+
+- `auto` selects `flat` below 5,000 clusters, `flat_rabitq` from 5,000 through
+  59,999 clusters, and `hnsw` from 60,000 clusters onward. These defaults are
+  practical thresholds; explicitly select a different initializer when it better
+  suits your workload.
+- `flat` computes exact distances to every centroid.
+- `flat_rabitq` scans one-bit centroid codes with FastScan, prunes candidates
+  using twice the standard RaBitQ error margin, and recomputes exact distances
+  for the retained candidates. Returned distances are exact for those centroids, but
+  routing is approximate: pruning can miss a true nearest centroid.
+- `hnsw` searches a graph of float32 centroids approximately.
+
+All choices support L2 and inner product. Flat RaBitQ uses the IVF rotation and
+padded domain, retains float32 centroids for refinement, and stores its mean and
+codes in the main index file. Its code precision is independent of the index's
+`nbits`. It trades additional code storage and an O(cluster count) query work
+buffer for fewer full-vector distance evaluations. Performance depends on
+cluster count, dimension, `nprobe`, and data distribution. Measure recall and
+latency on your workload, and use an explicit initializer to override `auto`.
+
+Existing constructor calls remain source-compatible. C++ consumers must rebuild
+against matching headers and the core library because the IVF object layout has
+changed. Newly built indexes with 5,000–59,999 clusters now use Flat RaBitQ routing. In particular, the 5,000–19,999 range changes
+from exact to approximate centroid selection. To reproduce the previous selection
+policy when rebuilding, explicitly choose `flat` below 20,000 clusters and `hnsw`
+otherwise. Loading an old index preserves its historical routing instead.
+
+Routing is fixed for an index and is also used by `add` when cluster IDs are
+omitted. `initializer_type()` (Python `initializer`) reports the resolved type,
+including after loading; loading never re-runs the current `auto` policy for a
+new-format file.
+
 Finally, call the construct API:
 ```c++
 void IVF::construct(
@@ -96,14 +146,18 @@ After construction, you can directly save the index file to disk:
 ```c++
 ivf.save(output_index_file);
 ```
-New files use the `RABQIDX1` magic, version 1, storage flags, and an explicit
-padded dimension before the index metadata. Loading restores the storage mode,
-rotation state, and padded dimension; raw files also contain the original vectors,
-so querying does not require an external dataset.
-
-Legacy unversioned quantized files and raw v1 files remain readable and retain
-their original 64-dimension padding. Saving a loaded legacy index writes the new
-header without changing its padded dimension or encoded data. Older library
+New files use the `RABQIDX1` magic, version 2, storage flags, an explicit
+padded dimension, and a uint32 initializer type (1 = Flat, 2 = Flat RaBitQ,
+3 = HNSW) before the index metadata. `Auto` is resolved before saving. Loading
+restores the routing type, storage mode, rotation state, and padded dimension.
+Raw files also contain the original vectors, so querying does not require an external dataset.
+
+Version-1 `RABQIDX1` files remain readable with their stored padding. Legacy
+unversioned quantized files and raw v1 files retain their original 64-dimension
+padding. All three historical formats infer routing with the original 20,000
+cluster threshold. HNSW routing still requires the accompanying `.hnsw` file.
+Flat RaBitQ loads its saved codes without random re-quantization. Saving a loaded
+legacy index writes the new header without changing its padded dimension or encoded data. Older library
 versions cannot read the new format.
 
 ### Data Layout
@@ -213,8 +267,8 @@ size_t IVF::remove(const PID* ids_to_remove, size_t n);
   receives the PID `max_elements() + i`, using the count before the call. Existing
   IDs stay unchanged, and removed IDs are never reused.
 - **cluster_ids**: The cluster of each new vector, in `[0, cluster_num)`. When
-  it is `nullptr`, vectors use the same centroid routing as queries. Routing is
-  exhaustive below 20,000 clusters and uses approximate HNSW search otherwise.
+  it is `nullptr`, vectors use the same centroid routing as queries. Routing uses
+  the initializer selected at construction or restored from the index file.
 - **faster**, **num_threads**: Same as in `construct`.
 - **ids_to_remove**: PIDs to remove. Every PID must be below `max_elements()`; nothing is
   removed if one is not. `remove` returns how many were newly removed and can be

include/rabitqlib/index/ivf/initializer.hpp

@@ -6,6 +6,7 @@
 #include <atomic>
 #include <cmath>
 #include <cstddef>
+#include <cstdint>
 #include <cstring>
 #include <exception>
 #include <fstream>
@@ -21,10 +22,28 @@
 
 #include "rabitqlib/defines.hpp"
 #include "rabitqlib/third/hnswlib/hnswlib.h"
+#include "rabitqlib/utils/memory.hpp"
 #include "rabitqlib/utils/space.hpp"
 #include "rabitqlib/utils/tools.hpp"
 
 namespace rabitqlib::ivf {
+enum class InitializerType : uint32_t { Auto = 0, Flat = 1, FlatRaBitQ = 2, HNSW = 3 };
+
+inline InitializerType resolve_initializer_type(InitializerType type, size_t clusters) {
+    switch (type) {
+        case InitializerType::Auto:
+            if (clusters < 5000) {
+                return InitializerType::Flat;
+            }
+            return clusters < 60000 ? InitializerType::FlatRaBitQ : InitializerType::HNSW;
+        case InitializerType::Flat:
+        case InitializerType::FlatRaBitQ:
+        case InitializerType::HNSW:
+            return type;
+    }
+    throw std::invalid_argument("Invalid IVF initializer type");
+}
+
 template <class Function>
 inline void parallel_for(size_t start, size_t end, size_t numThreads, Function fn) {
     if (start >= end) {
@@ -197,6 +216,26 @@ class FlatInitializer : public Initializer {
     }
 };
 
+// Inputs are already in IVF's rotated, padded domain. Persist the codes instead
+// of regenerating them on load, so routing keeps the same quantization state.
+class FlatRaBitQInitializer : public Initializer {
+   private:
+    std::vector<float> centroids_;
+    std::vector<float> global_centroid_;
+    std::vector<std::byte, memory::AlignedAllocator<std::byte, 64, true>> batch_data_;
+    MetricType metric_type_;
+
+   public:
+    explicit FlatRaBitQInitializer(size_t d, size_t k, MetricType metric = METRIC_L2);
+    [[nodiscard]] const float* centroid(PID id) const override;
+    void add_vectors(const float* centroids, size_t num_threads) override;
+    void centroids_distances(
+        const float* query, size_t nprobe, std::vector<AnnCandidate<float>>& candidates
+    ) const override;
+    void save(std::ofstream& output, const char*) const override;
+    void load(std::ifstream& input, const char*) override;
+};
+
 // Keep centroid routing on the same runtime-selected distance kernels as flat IVF.
 class CentroidL2Space : public hnswlib::SpaceInterface<float> {
    private:

include/rabitqlib/index/ivf/ivf.hpp

@@ -47,6 +47,7 @@ class IVF {
     using FloatStorage = std::vector<float, memory::AlignedAllocator<float, 64, true>>;
     using IdStorage = std::vector<PID, memory::AlignedAllocator<PID, 64, true>>;
 
+    InitializerType initializer_type_ = InitializerType::Auto;
     std::unique_ptr<Initializer> initer_;            // initializer for candidate clusters
     ByteStorage batch_storage_;                      // 1-bit code and factors
     ByteStorage ex_storage_;                         // extra-bit codes and packed factors
@@ -68,7 +69,7 @@ class IVF {
     static constexpr uint64_t kRawFormatMagic = 0x3157415251464252ULL;
     static constexpr uint32_t kRawFormatVersion = 1;
     static constexpr uint64_t kFormatMagic = 0x3158444951424152ULL;  // "RABQIDX1"
-    static constexpr uint32_t kFormatVersion = 1;
+    static constexpr uint32_t kFormatVersion = 2;
 
     void
     quantize_cluster(Cluster&, const std::vector<PID>&, const float*, const float*, float*, float*, const quant::RabitqConfig&);
@@ -142,6 +143,7 @@ class IVF {
 
     void swap(IVF& other) noexcept {
         using std::swap;
+        swap(initializer_type_, other.initializer_type_);
         swap(initer_, other.initer_);
         swap(batch_storage_, other.batch_storage_);
         swap(ex_storage_, other.ex_storage_);
@@ -179,7 +181,8 @@ class IVF {
         size_t,
         size_t,
         MetricType metric_type = rabitqlib::METRIC_L2,
-        RotatorType type = RotatorType::FhtKacRotator
+        RotatorType type = RotatorType::FhtKacRotator,
+        InitializerType initializer = InitializerType::Auto
     );
 
     ~IVF();
@@ -190,6 +193,7 @@ class IVF {
     [[nodiscard]] size_t nbits() const { return raw_reranking_ ? 32 : ex_bits_ + 1; }
     [[nodiscard]] MetricType metric_type() const { return metric_type_; }
     [[nodiscard]] RotatorType rotator_type() const { return type_; }
+    [[nodiscard]] InitializerType initializer_type() const { return initializer_type_; }
 
     void construct(const float*, const float*, const PID*, bool, size_t);
 
@@ -276,9 +280,11 @@ inline IVF::IVF(
     size_t cluster_num,
     size_t bits,
     MetricType metric_type,
-    RotatorType type
+    RotatorType type,
+    InitializerType initializer
 )
-    : num_(n)
+    : initializer_type_(resolve_initializer_type(initializer, cluster_num))
+    , num_(n)
     , dim_(dim)
     , padded_dim_(dim)
     , num_cluster_(cluster_num)
@@ -397,9 +403,13 @@ inline void IVF::allocate_memory(const std::vector<size_t>& cluster_sizes) {
     free_memory();
     cluster_lst_.clear();
 
-    if (num_cluster_ < 20000UL) {
+    if (initializer_type_ == InitializerType::Flat) {
         this->initer_ =
             std::make_unique<FlatInitializer>(padded_dim_, num_cluster_, metric_type_);
+    } else if (initializer_type_ == InitializerType::FlatRaBitQ) {
+        this->initer_ = std::make_unique<FlatRaBitQInitializer>(
+            padded_dim_, num_cluster_, metric_type_
+        );
     } else {
         this->initer_ =
             std::make_unique<HNSWInitializer>(padded_dim_, num_cluster_, metric_type_);
@@ -791,6 +801,11 @@ inline void IVF::save(const char* filename) const {
         reinterpret_cast<const char*>(&stored_padded_dim), sizeof(stored_padded_dim)
     );
 
+    const auto stored_initializer = static_cast<uint32_t>(initializer_type_);
+    output.write(
+        reinterpret_cast<const char*>(&stored_initializer), sizeof(stored_initializer)
+    );
+
     /* Save meta data */
     output.write(reinterpret_cast<const char*>(&num_), sizeof(size_t));
     output.write(reinterpret_cast<const char*>(&dim_), sizeof(size_t));
@@ -851,14 +866,23 @@ inline void IVF::load(const char* filename) {
         input.read(reinterpret_cast<char*>(&version), sizeof(version));
         input.read(reinterpret_cast<char*>(&flags), sizeof(flags));
         input.read(reinterpret_cast<char*>(&stored_padded_dim), sizeof(stored_padded_dim));
-        if (version != kFormatVersion) {
+        if (version != 1 && version != kFormatVersion) {
             throw std::runtime_error("Unsupported IVF index version");
         }
         if (flags > 1 || stored_padded_dim > std::numeric_limits<size_t>::max()) {
             throw std::runtime_error("Invalid IVF index metadata");
         }
         loaded.raw_reranking_ = flags == 1;
         loaded.padded_dim_ = static_cast<size_t>(stored_padded_dim);
+        if (version >= 2) {
+            uint32_t initializer = 0;
+            input.read(reinterpret_cast<char*>(&initializer), sizeof(initializer));
+            if (initializer < static_cast<uint32_t>(InitializerType::Flat) ||
+                initializer > static_cast<uint32_t>(InitializerType::HNSW)) {
+                throw std::runtime_error("Invalid IVF initializer type in index file");
+            }
+            loaded.initializer_type_ = static_cast<InitializerType>(initializer);
+        }
     } else if (magic == kRawFormatMagic) {
         loaded.raw_reranking_ = true;
         uint32_t version = 0;
@@ -878,6 +902,12 @@ inline void IVF::load(const char* filename) {
     input.read(reinterpret_cast<char*>(&loaded.type_), sizeof(loaded.type_));
     input.read(reinterpret_cast<char*>(&loaded.metric_type_), sizeof(loaded.metric_type_));
     validate_metric_type(loaded.metric_type_);
+    // Missing routing metadata means a historical file. Its payload layout must
+    // follow the original threshold, independently of the current Auto policy.
+    if (loaded.initializer_type_ == InitializerType::Auto) {
+        loaded.initializer_type_ =
+            loaded.num_cluster_ < 20000 ? InitializerType::Flat : InitializerType::HNSW;
+    }
     if (loaded.num_ == 0 || loaded.num_ > buffer::kSearchBufferMaxPointCount ||
         loaded.dim_ == 0 || loaded.num_cluster_ == 0 ||
         loaded.num_cluster_ > buffer::kSearchBufferMaxPointCount || loaded.ex_bits_ > 8 ||
@@ -911,9 +941,16 @@ inline void IVF::load(const char* filename) {
     } else {
         throw std::runtime_error("Invalid IVF rotator type");
     }
-    if (loaded.num_cluster_ < 20000UL) {
+    if (loaded.initializer_type_ != InitializerType::HNSW) {
         consume(loaded.num_cluster_, sizeof(float) * loaded.padded_dim_);
     }
+    if (loaded.initializer_type_ == InitializerType::FlatRaBitQ) {
+        consume(loaded.padded_dim_, sizeof(float));
+        consume(
+            div_round_up(loaded.num_cluster_, fastscan::kBatchSize),
+            BatchDataMap<float>::data_bytes(loaded.padded_dim_)
+        );
+    }
 
     /* Load number of vectors of each cluster */
     std::vector<size_t> cluster_sizes(loaded.num_cluster_, 0);

include/rabitqlib/third/hnswlib/hnswalg.h

@@ -685,6 +685,9 @@ class HierarchicalNSW : public AlgorithmInterface<dist_t> {
 
     void saveIndex(const std::string &location) {
         std::ofstream output(rabitqlib::io_impl::filesystem_path(location), std::ios::binary);
+        if (!output.is_open())
+            throw std::runtime_error("Cannot open HNSW index file for writing");
+        output.exceptions(std::ios::failbit | std::ios::badbit);
         std::streampos position;
 
         writeBinaryPOD(output, offsetLevel0_);