diff --git a/AGENTS.md b/AGENTS.md index fbe3e5a1..15c7dd9c 100644 --- a/AGENTS.md +++ b/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. diff --git a/CMakeLists.txt b/CMakeLists.txt index 0d93f4fb..9dffcc8f 100644 --- a/CMakeLists.txt +++ b/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 diff --git a/docs/docs/index/ivf.md b/docs/docs/index/ivf.md index 123a1397..bbbef31b 100644 --- a/docs/docs/index/ivf.md +++ b/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 diff --git a/include/rabitqlib/index/ivf/initializer.hpp b/include/rabitqlib/index/ivf/initializer.hpp index 8885369d..733b346f 100644 --- a/include/rabitqlib/index/ivf/initializer.hpp +++ b/include/rabitqlib/index/ivf/initializer.hpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -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 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 centroids_; + std::vector global_centroid_; + std::vector> 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>& 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 { private: diff --git a/include/rabitqlib/index/ivf/ivf.hpp b/include/rabitqlib/index/ivf/ivf.hpp index 97e7672a..93256f44 100644 --- a/include/rabitqlib/index/ivf/ivf.hpp +++ b/include/rabitqlib/index/ivf/ivf.hpp @@ -47,6 +47,7 @@ class IVF { using FloatStorage = std::vector>; using IdStorage = std::vector>; + InitializerType initializer_type_ = InitializerType::Auto; std::unique_ptr 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&, 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& cluster_sizes) { free_memory(); cluster_lst_.clear(); - if (num_cluster_ < 20000UL) { + if (initializer_type_ == InitializerType::Flat) { this->initer_ = std::make_unique(padded_dim_, num_cluster_, metric_type_); + } else if (initializer_type_ == InitializerType::FlatRaBitQ) { + this->initer_ = std::make_unique( + padded_dim_, num_cluster_, metric_type_ + ); } else { this->initer_ = std::make_unique(padded_dim_, num_cluster_, metric_type_); @@ -791,6 +801,11 @@ inline void IVF::save(const char* filename) const { reinterpret_cast(&stored_padded_dim), sizeof(stored_padded_dim) ); + const auto stored_initializer = static_cast(initializer_type_); + output.write( + reinterpret_cast(&stored_initializer), sizeof(stored_initializer) + ); + /* Save meta data */ output.write(reinterpret_cast(&num_), sizeof(size_t)); output.write(reinterpret_cast(&dim_), sizeof(size_t)); @@ -851,7 +866,7 @@ inline void IVF::load(const char* filename) { input.read(reinterpret_cast(&version), sizeof(version)); input.read(reinterpret_cast(&flags), sizeof(flags)); input.read(reinterpret_cast(&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::max()) { @@ -859,6 +874,15 @@ inline void IVF::load(const char* filename) { } loaded.raw_reranking_ = flags == 1; loaded.padded_dim_ = static_cast(stored_padded_dim); + if (version >= 2) { + uint32_t initializer = 0; + input.read(reinterpret_cast(&initializer), sizeof(initializer)); + if (initializer < static_cast(InitializerType::Flat) || + initializer > static_cast(InitializerType::HNSW)) { + throw std::runtime_error("Invalid IVF initializer type in index file"); + } + loaded.initializer_type_ = static_cast(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(&loaded.type_), sizeof(loaded.type_)); input.read(reinterpret_cast(&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::data_bytes(loaded.padded_dim_) + ); + } /* Load number of vectors of each cluster */ std::vector cluster_sizes(loaded.num_cluster_, 0); diff --git a/include/rabitqlib/third/hnswlib/hnswalg.h b/include/rabitqlib/third/hnswlib/hnswalg.h index 701d0de4..7ad27883 100644 --- a/include/rabitqlib/third/hnswlib/hnswalg.h +++ b/include/rabitqlib/third/hnswlib/hnswalg.h @@ -685,6 +685,9 @@ class HierarchicalNSW : public AlgorithmInterface { 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_); diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index 1073c890..55bd8d3a 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/python_bindings/ivf_bindings.cpp @@ -25,6 +25,36 @@ namespace py = pybind11; namespace rabitqlib::python_bindings { +namespace { +ivf::InitializerType initializer_from_string(const std::string& value) { + if (value == "auto") + return ivf::InitializerType::Auto; + if (value == "flat") + return ivf::InitializerType::Flat; + if (value == "flat_rabitq") + return ivf::InitializerType::FlatRaBitQ; + if (value == "hnsw") + return ivf::InitializerType::HNSW; + throw std::invalid_argument( + "initializer must be 'auto', 'flat', 'flat_rabitq', or 'hnsw'" + ); +} + +std::string initializer_to_string(ivf::InitializerType type) { + switch (type) { + case ivf::InitializerType::Auto: + return "auto"; + case ivf::InitializerType::Flat: + return "flat"; + case ivf::InitializerType::FlatRaBitQ: + return "flat_rabitq"; + case ivf::InitializerType::HNSW: + return "hnsw"; + } + throw std::logic_error("Invalid IVF initializer type"); +} +} // namespace + class IvfIndex { public: IvfIndex( @@ -32,7 +62,8 @@ class IvfIndex { size_t max_elements, size_t num_clusters, size_t nbits, - const std::string& metric = "l2" + const std::string& metric = "l2", + const std::string& initializer = "auto" ) : dim_(dim) , max_elements_(max_elements) @@ -45,7 +76,8 @@ class IvfIndex { num_clusters, nbits, metric_, - rabitqlib::RotatorType::FhtKacRotator + rabitqlib::RotatorType::FhtKacRotator, + initializer_from_string(initializer) )) {} void build( @@ -235,6 +267,9 @@ class IvfIndex { [[nodiscard]] size_t num_clusters() const { return num_clusters_; } [[nodiscard]] size_t nbits() const { return nbits_; } [[nodiscard]] bool is_built() const { return built_; } + [[nodiscard]] std::string initializer() const { + return initializer_to_string(index_->initializer_type()); + } [[nodiscard]] std::string metric() const { return metric_to_string(metric_); } private: @@ -257,14 +292,22 @@ void register_ivf(py::module_& m) { py::class_(m, "IvfIndex") .def( - py::init(), + py::init< + size_t, + size_t, + size_t, + size_t, + const std::string&, + const std::string&>(), py::arg("dim"), py::arg("max_elements"), py::arg("num_clusters"), py::arg("nbits"), py::arg("metric") = "l2", + py::arg("initializer") = "auto", "Create IVF with 1-9 quantization bits, or nbits=32 for owned raw-vector " - "reranking." + "reranking. initializer selects auto, flat, flat_rabitq, or hnsw centroid " + "routing." ) .def( "build", @@ -310,5 +353,6 @@ void register_ivf(py::module_& m) { .def_property_readonly("num_clusters", &IvfIndex::num_clusters) .def_property_readonly("nbits", &IvfIndex::nbits) .def_property_readonly("is_built", &IvfIndex::is_built) - .def_property_readonly("metric", &IvfIndex::metric); + .def_property_readonly("metric", &IvfIndex::metric) + .def_property_readonly("initializer", &IvfIndex::initializer); } diff --git a/src/index/ivf_initializer.cpp b/src/index/ivf_initializer.cpp new file mode 100644 index 00000000..3fa0a49c --- /dev/null +++ b/src/index/ivf_initializer.cpp @@ -0,0 +1,179 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/index/estimator.hpp" +#include "rabitqlib/index/ivf/initializer.hpp" +#include "rabitqlib/index/query.hpp" +#include "rabitqlib/quantization/data_layout.hpp" +#include "rabitqlib/quantization/rabitq.hpp" +#include "rabitqlib/utils/buffer.hpp" + +namespace rabitqlib::ivf { +FlatRaBitQInitializer::FlatRaBitQInitializer(size_t d, size_t k, MetricType metric) + : Initializer(d, k), metric_type_(metric) { + validate_metric_type(metric); + const size_t max_size = std::numeric_limits::max(); + if (d == 0 || d % 32 != 0 || d > max_size / 32 || k == 0 || + k > buffer::kSearchBufferMaxPointCount || k > max_size / sizeof(float) / d) { + throw std::invalid_argument("Invalid Flat RaBitQ initializer dimensions"); + } + const size_t batches = div_round_up(k, fastscan::kBatchSize); + const size_t batch_bytes = BatchDataMap::data_bytes(d); + if (batches > max_size / batch_bytes) { + throw std::invalid_argument("Flat RaBitQ initializer storage overflow"); + } + centroids_.resize(k * d); + global_centroid_.resize(d); + batch_data_.resize(batches * batch_bytes); +} + +const float* FlatRaBitQInitializer::centroid(PID id) const { + return centroids_.data() + static_cast(id) * dim_; +} + +void FlatRaBitQInitializer::add_vectors(const float* centroids, size_t num_threads) { + std::memcpy(centroids_.data(), centroids, centroids_.size() * sizeof(float)); + for (size_t d = 0; d < dim_; ++d) { + double sum = 0; + for (size_t i = 0; i < num_cluster_; ++i) { + sum += centroids_[i * dim_ + d]; + } + global_centroid_[d] = static_cast(sum / static_cast(num_cluster_)); + } + const size_t batch_bytes = BatchDataMap::data_bytes(dim_); + parallel_for( + 0, + div_round_up(num_cluster_, fastscan::kBatchSize), + num_threads, + [&](size_t batch, size_t) { + const size_t first = batch * fastscan::kBatchSize; + quant::quantize_split_batch( + centroids_.data() + first * dim_, + global_centroid_.data(), + std::min(fastscan::kBatchSize, num_cluster_ - first), + dim_, + 0, + reinterpret_cast(batch_data_.data()) + batch * batch_bytes, + nullptr, + metric_type_ + ); + } + ); +} + +void FlatRaBitQInitializer::centroids_distances( + const float* query, size_t nprobe, std::vector>& candidates +) const { + const size_t topk = std::min(nprobe, num_cluster_); + candidates.resize(topk); + if (topk == 0) { + return; + } + const auto exact_distance = [&](PID id) { + return metric_type_ == METRIC_L2 ? euclidean_sqr(query, centroid(id), dim_) + : dot_product_dis(query, centroid(id), dim_); + }; + if (topk == num_cluster_) { + for (size_t i = 0; i < topk; ++i) { + candidates[i] = AnnCandidate(static_cast(i), exact_distance(i)); + } + std::sort(candidates.begin(), candidates.end()); + } else { + SplitBatchQuery q_obj(query, dim_, 0, metric_type_, false); + const float norm = std::sqrt(euclidean_sqr(query, global_centroid_.data(), dim_)); + q_obj.set_g_add( + norm, + metric_type_ == METRIC_IP ? dot_product(query, global_centroid_.data(), dim_) + : 0.0F + ); + std::vector lower(num_cluster_); + std::vector> upper(num_cluster_); + std::array estimates{}, lows{}, ips{}; + const size_t batch_bytes = BatchDataMap::data_bytes(dim_); + for (size_t first = 0; first < num_cluster_; first += fastscan::kBatchSize) { + split_batch_estdist( + reinterpret_cast(batch_data_.data()) + + (first / fastscan::kBatchSize) * batch_bytes, + q_obj, + dim_, + estimates.data(), + lows.data(), + ips.data(), + false + ); + const size_t count = std::min(fastscan::kBatchSize, num_cluster_ - first); + for (size_t j = 0; j < count; ++j) { + // Use twice the standard RaBitQ error margin for centroid routing: + // missing a centroid also skips all vectors in its IVF list. + lower[first + j] = estimates[j] - 2.0F * (estimates[j] - lows[j]); + upper[first + j] = AnnCandidate( + static_cast(first + j), 2 * estimates[j] - lower[first + j] + ); + } + } + std::nth_element( + upper.begin(), upper.begin() + static_cast(topk), upper.end() + ); + buffer::SearchBuffer knns(topk); + // Seed with exact distances: an estimated upper bound need not hold for + // every vector. Always fill the result before pruning other candidates. + for (size_t i = 0; i < topk; ++i) { + knns.insert(upper[i].id, exact_distance(upper[i].id)); + } + for (size_t i = topk; i < num_cluster_; ++i) { + const PID id = upper[i].id; + if (lower[id] <= knns.top_dist()) { + knns.insert(id, exact_distance(id)); + } + } + std::copy_n(knns.data().begin(), topk, candidates.begin()); + } + if (metric_type_ == METRIC_L2) { + for (auto& candidate : candidates) { + candidate.distance = std::sqrt(candidate.distance); + } + } +} + +void FlatRaBitQInitializer::save(std::ofstream& output, const char*) const { + output.write( + reinterpret_cast(centroids_.data()), + static_cast(centroids_.size() * sizeof(float)) + ); + output.write( + reinterpret_cast(global_centroid_.data()), + static_cast(global_centroid_.size() * sizeof(float)) + ); + output.write( + reinterpret_cast(batch_data_.data()), + static_cast(batch_data_.size()) + ); + if (!output) { + throw std::runtime_error("Cannot write Flat RaBitQ initializer"); + } +} + +void FlatRaBitQInitializer::load(std::ifstream& input, const char*) { + input.read( + reinterpret_cast(centroids_.data()), + static_cast(centroids_.size() * sizeof(float)) + ); + input.read( + reinterpret_cast(global_centroid_.data()), + static_cast(global_centroid_.size() * sizeof(float)) + ); + input.read( + reinterpret_cast(batch_data_.data()), + static_cast(batch_data_.size()) + ); + if (!input) { + throw std::runtime_error("Cannot read Flat RaBitQ initializer"); + } +} +} // namespace rabitqlib::ivf diff --git a/tests/consumer/main.cpp b/tests/consumer/main.cpp index 6040e951..58a97420 100644 --- a/tests/consumer/main.cpp +++ b/tests/consumer/main.cpp @@ -93,7 +93,19 @@ int main() { return 1; } - rabitqlib::ivf::IVF index(kPoints, kDimension, flat.k, 4); + rabitqlib::ivf::IVF index( + kPoints, + kDimension, + flat.k, + 4, + rabitqlib::METRIC_L2, + rabitqlib::RotatorType::FhtKacRotator, + rabitqlib::ivf::InitializerType::FlatRaBitQ + ); + if (index.initializer_type() != rabitqlib::ivf::InitializerType::FlatRaBitQ) { + std::cerr << "Installed IVF did not preserve explicit centroid routing\n"; + return 1; + } index.construct(data.data(), flat.centroids.data(), flat.assignments.data(), false, 2); std::vector ids(kPoints); std::vector distances(kPoints); diff --git a/tests/python/test_ivf.py b/tests/python/test_ivf.py index da21d258..4323b376 100644 --- a/tests/python/test_ivf.py +++ b/tests/python/test_ivf.py @@ -413,7 +413,7 @@ def test_raw_index_rejects_invalid_files(tmp_path, damage): idx.save(str(path)) payload = bytearray(path.read_bytes()) if damage == "version": - payload[8:12] = (2).to_bytes(4, "little") + payload[8:12] = (99).to_bytes(4, "little") elif damage == "flags": payload[12:16] = (2).to_bytes(4, "little") elif damage.startswith("padding_"): @@ -424,13 +424,13 @@ def test_raw_index_rejects_invalid_files(tmp_path, damage): }[damage] payload[16:24] = padded.to_bytes(8, "little") elif damage == "ex_bits": - payload[48:56] = (1).to_bytes(8, "little") + payload[52:60] = (1).to_bytes(8, "little") elif damage == "dimension": - payload[32:40] = (2**64 - 1).to_bytes(8, "little") + payload[36:44] = (2**64 - 1).to_bytes(8, "little") elif damage == "count": - payload[24:32] = (2**64 - 1).to_bytes(8, "little") + payload[28:36] = (2**64 - 1).to_bytes(8, "little") elif damage == "clusters": - payload[40:48] = (2**64 - 1).to_bytes(8, "little") + payload[44:52] = (2**64 - 1).to_bytes(8, "little") elif damage == "truncated_header": payload = payload[:9] else: @@ -467,14 +467,45 @@ def test_automatic_high_accuracy(nbits, tmp_path): np.testing.assert_array_equal(actual[1], expected[1]) +@pytest.mark.parametrize("failure", ["open", "write"]) +def test_hnsw_sidecar_save_failures(tmp_path, failure): + if failure == "write" and not Path("/dev/full").exists(): + pytest.skip("requires a device that rejects writes") + data = np.zeros((2, 64), dtype=np.float32) + data[1, 0] = 1 + idx = IvfIndex(64, 2, 2, nbits=32, initializer="hnsw") + idx.build(data, data, np.arange(2, dtype=np.uint32)) + expected = idx.search(data, 1, 1) + path = tmp_path / "sidecar_failure.index" + sidecar = Path(str(path) + ".hnsw") + if failure == "open": + sidecar.mkdir() + message = "Cannot open HNSW index file for writing" + else: + sidecar.symlink_to("/dev/full") + message = None # The stream failure message depends on the standard library. + with pytest.raises(RuntimeError, match=message): + idx.save(str(path)) + if failure == "open": + sidecar.rmdir() + else: + sidecar.unlink() + # Failure does not invalidate the in-memory index, and saving can be retried. + idx.save(str(path)) + for obj in (idx, IvfIndex.load(str(path))): + actual = obj.search(data, 1, 1) + np.testing.assert_array_equal(actual[0], expected[0]) + np.testing.assert_array_equal(actual[1], expected[1]) + + def test_hnsw_centroid_routing_roundtrip(tmp_path): - # 20,000 clusters selects HNSW centroid routing rather than the flat router. + # Explicit routing keeps HNSW coverage independent of the auto thresholds. # One vector per centroid makes nearest-centroid IDs and distances exact. rng = np.random.default_rng(314) count, dim = 20000, 64 centroids = rng.standard_normal((count, dim)).astype(np.float32) labels = np.arange(count, dtype=np.uint32) - idx = IvfIndex(dim, count, count, nbits=1) + idx = IvfIndex(dim, count, count, nbits=1, initializer="hnsw") idx.build(centroids, centroids, labels, num_threads=4) selected = np.array([0, 17, 1023, count - 1]) queries = centroids[selected] @@ -498,7 +529,7 @@ def test_hnsw_inner_product_routing_roundtrip(tmp_path): data = centroids[:1].copy() query = np.zeros((4, dim), dtype=np.float32) query[:, 0] = 1.0 - idx = IvfIndex(dim, 1, count, nbits=32, metric="ip") + idx = IvfIndex(dim, 1, count, nbits=32, metric="ip", initializer="hnsw") idx.build(data, centroids, np.array([0], dtype=np.uint32), num_threads=4) path = str(tmp_path / "large_ip_centroid.index") diff --git a/tests/python/test_ivf_initializers.py b/tests/python/test_ivf_initializers.py new file mode 100644 index 00000000..662051a2 --- /dev/null +++ b/tests/python/test_ivf_initializers.py @@ -0,0 +1,158 @@ +"""Explicit centroid routing, persistence, and legacy initializer inference.""" + +from pathlib import Path + +import numpy as np +import pytest +from rabitqlib import IvfIndex + + +@pytest.mark.parametrize("metric", ["l2", "ip"]) +@pytest.mark.parametrize("bits", list(range(1, 10)) + [32]) +@pytest.mark.parametrize("initializer", ["auto", "flat", "flat_rabitq", "hnsw"]) +def test_initializer_roundtrip_and_updates(tmp_path, metric, bits, initializer): + rng = np.random.default_rng(773) + dim, count = 65, 33 + data = rng.standard_normal((count + 1, dim)).astype(np.float32) + data /= np.linalg.norm(data, axis=1, keepdims=True) + index = IvfIndex(dim, count, count, bits, metric, initializer=initializer) + resolved = "flat" if initializer == "auto" else initializer + assert index.initializer == resolved + index.build(data[:count], data[:count], np.arange(count, dtype=np.uint32), 2) + queries = data[[0, 17, 32]] + expected = index.search(queries, 3, 7, num_threads=2) + np.testing.assert_array_equal(expected[0][:, 0], [0, 17, 32]) + path = tmp_path / "索引.index" + index.save(str(path)) + payload = path.read_bytes() + assert int.from_bytes(payload[8:12], "little") == 2 + assert ( + int.from_bytes(payload[24:28], "little") + == { + "flat": 1, + "flat_rabitq": 2, + "hnsw": 3, + }[resolved] + ) + assert Path(str(path) + ".hnsw").exists() == (resolved == "hnsw") + loaded = IvfIndex.load(str(path)) + assert loaded.initializer == resolved + actual = loaded.search(queries, 3, 7, num_threads=2) + for left, right in zip(expected, actual): + np.testing.assert_array_equal(left, right) + loaded.save(str(path)) + assert path.read_bytes() == payload + for obj in (index, loaded): + obj.add(data[count:], num_threads=2) + obj.remove(np.array([0], dtype=np.uint32)) + assert obj.initializer == resolved + for left, right in zip(index.search(queries, 3, 7), loaded.search(queries, 3, 7)): + np.testing.assert_array_equal(left, right) + loaded.save(str(path)) + assert IvfIndex.load(str(path)).initializer == resolved + + +@pytest.mark.parametrize( + "clusters,expected", + [ + (4999, "flat"), + (5000, "flat_rabitq"), + (19999, "flat_rabitq"), + (20000, "flat_rabitq"), + (59999, "flat_rabitq"), + (60000, "hnsw"), + ], +) +def test_auto_initializer_thresholds(clusters, expected): + assert IvfIndex(64, 1, clusters, 1).initializer == expected + for mode in ("flat", "flat_rabitq", "hnsw"): + assert IvfIndex(64, 1, clusters, 1, initializer=mode).initializer == mode + + +@pytest.mark.parametrize("bits", [1, 32]) +@pytest.mark.parametrize("clusters", [3, 5000, 19999, 20000]) +@pytest.mark.parametrize("format_version", ["v1", "legacy"]) +def test_legacy_formats_infer_historical_initializer( + tmp_path, bits, clusters, format_version +): + rng = np.random.default_rng(43) + centroids = rng.standard_normal((clusters, 64)).astype(np.float32) + historical = "flat" if clusters < 20000 else "hnsw" + index = IvfIndex(64, 3, clusters, bits, initializer=historical) + index.build(centroids[:3], centroids, np.arange(3, dtype=np.uint32), 2) + path = tmp_path / "v1.index" + index.save(str(path)) + # v2 must honor the stored type even when today's Auto would choose differently. + assert IvfIndex.load(str(path)).initializer == historical + payload = bytearray(path.read_bytes()) + if format_version == "v1": + # v1 has the same payload but no initializer field after padded_dim. + del payload[24:28] + payload[8:12] = (1).to_bytes(4, "little") + else: + # Dimension 64 also matches the padding inferred by both older formats. + payload = payload[28:] + if bits == 32: + payload = b"RBFQRAW1" + (1).to_bytes(4, "little") + payload + path.write_bytes(payload) + loaded = IvfIndex.load(str(path)) + assert loaded.initializer == index.initializer + for left, right in zip( + index.search(centroids[:3], 1, 2), loaded.search(centroids[:3], 1, 2) + ): + np.testing.assert_array_equal(left, right) + + +def test_initializer_rejects_invalid_argument(): + with pytest.raises(ValueError, match="initializer must be"): + IvfIndex(64, 3, 1, 1, initializer="unknown") + + +@pytest.mark.parametrize("initializer", ["flat", "flat_rabitq"]) +def test_explicit_flat_routing_above_historical_threshold(tmp_path, initializer): + centroids = np.zeros((20000, 64), dtype=np.float32) + centroids[17, 0] = 1 + index = IvfIndex(64, 1, len(centroids), 32, initializer=initializer) + index.build(centroids[17:18], centroids, np.array([17], dtype=np.uint32), 2) + path = tmp_path / "large.index" + index.save(str(path)) + assert not Path(str(path) + ".hnsw").exists() + for obj in (index, IvfIndex.load(str(path))): + assert obj.initializer == initializer + ids, distances = obj.search(centroids[17:18], 1, 1) + np.testing.assert_array_equal(ids, [[0]]) + np.testing.assert_array_equal(distances, [[0]]) + + +@pytest.mark.parametrize("damage", ["auto", "unknown", "truncated_initializer"]) +def test_initializer_rejects_invalid_persistence(tmp_path, damage): + data = np.zeros((3, 65), dtype=np.float32) + index = IvfIndex(65, 3, 1, 1, initializer="flat_rabitq") + index.build(data, data[:1], np.zeros(3, dtype=np.uint32)) + path = tmp_path / "bad.index" + index.save(str(path)) + payload = bytearray(path.read_bytes()) + if damage in ("auto", "unknown"): + payload[24:28] = (0 if damage == "auto" else 99).to_bytes(4, "little") + else: + # The rotator is followed by raw centroids, mean, and the physical code batch. + payload = payload[: 28 + 4 * 8 + 2 + 8 + 96 // 2 + 96 * 4 + 1] + path.write_bytes(payload) + with pytest.raises(RuntimeError): + IvfIndex.load(str(path)) + + +@pytest.mark.parametrize("clusters", [5000, 20000, 59999]) +def test_auto_flat_rabitq_persists_actual_routing(tmp_path, clusters): + centroids = np.zeros((clusters, 64), dtype=np.float32) + centroids[17, 0] = 1 + index = IvfIndex(64, 1, clusters, 32) + index.build(centroids[17:18], centroids, np.array([17], dtype=np.uint32), 2) + path = tmp_path / "auto.index" + index.save(str(path)) + assert not Path(str(path) + ".hnsw").exists() + for obj in (index, IvfIndex.load(str(path))): + assert obj.initializer == "flat_rabitq" + ids, distances = obj.search(centroids[17:18], 1, 1) + np.testing.assert_array_equal(ids, [[0]]) + np.testing.assert_array_equal(distances, [[0]]) diff --git a/tests/unit/rabitqlib/index/flat_rabitq_initializer_test.cpp b/tests/unit/rabitqlib/index/flat_rabitq_initializer_test.cpp new file mode 100644 index 00000000..48f7cc3c --- /dev/null +++ b/tests/unit/rabitqlib/index/flat_rabitq_initializer_test.cpp @@ -0,0 +1,263 @@ +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/index/ivf/initializer.hpp" +#include "rabitqlib/index/ivf/ivf.hpp" + +namespace rabitqlib::ivf { +namespace { +TEST(IvfInitializerTest, ExplicitChoiceOverridesAutomaticThreshold) { + const std::array, 6> cases{{ + {4999, InitializerType::Flat}, + {5000, InitializerType::FlatRaBitQ}, + {19999, InitializerType::FlatRaBitQ}, + {20000, InitializerType::FlatRaBitQ}, + {59999, InitializerType::FlatRaBitQ}, + {60000, InitializerType::HNSW}, + }}; + for (const auto& [clusters, expected] : cases) { + IVF automatic(1, 64, clusters, 1); + EXPECT_EQ(automatic.initializer_type(), expected); + for (auto type : + {InitializerType::Flat, InitializerType::FlatRaBitQ, InitializerType::HNSW}) { + IVF explicit_type( + 1, 64, clusters, 1, METRIC_L2, RotatorType::FhtKacRotator, type + ); + EXPECT_EQ(explicit_type.initializer_type(), type); + } + } + EXPECT_THROW( + ( + IVF(1, + 64, + 1, + 1, + METRIC_L2, + RotatorType::FhtKacRotator, + static_cast(99)) + ), + std::invalid_argument + ); +} + +TEST(IvfInitializerTest, MatrixRotationAndExplicitRoutingRoundTrip) { + constexpr size_t kDim = 17, kCount = 3; + std::vector data(kDim * kCount, 0.0F); + data[0] = 1; + data[kDim + 1] = 1; + data[kDim * 2 + 2] = 1; + const std::array labels{0, 1, 2}; + const std::string path = ::testing::TempDir() + "ivf_matrix_flat_rabitq.index"; + for (auto metric : {METRIC_L2, METRIC_IP}) { + IVF index( + kCount, + kDim, + kCount, + 32, + metric, + RotatorType::MatrixRotator, + InitializerType::FlatRaBitQ + ); + index.construct(data.data(), data.data(), labels.data(), false, 1); + index.save(path.c_str()); + IVF loaded; + loaded.load(path.c_str()); + EXPECT_EQ(loaded.initializer_type(), InitializerType::FlatRaBitQ); + for (size_t i = 0; i < kCount; ++i) { + PID id = kPidMax; + float distance = 0; + loaded.search(data.data() + i * kDim, 1, 1, &id, &distance); + EXPECT_EQ(id, i); + EXPECT_NEAR(distance, 0, 1e-5); + } + } + std::remove(path.c_str()); +} + +TEST(IvfInitializerTest, ReloadChangesRoutingAndFailedLoadPreservesExistingIndex) { + constexpr size_t kDim = 65, kCount = 33, kTopk = 5; + std::mt19937 generator(9923); + std::normal_distribution normal; + std::vector data(kCount * kDim); + std::vector labels(kCount); + for (size_t i = 0; i < kCount; ++i) { + labels[i] = static_cast(i); + for (size_t j = 0; j < kDim; ++j) + data[i * kDim + j] = normal(generator); + } + const std::string path = ::testing::TempDir() + "ivf_initializer_reload.index"; + IVF target; + for (auto type : + {InitializerType::Flat, + InitializerType::FlatRaBitQ, + InitializerType::HNSW, + InitializerType::Flat}) { + IVF source(kCount, kDim, kCount, 5, METRIC_L2, RotatorType::FhtKacRotator, type); + source.construct(data.data(), data.data(), labels.data(), false, 1); + std::array expected_ids{}, actual_ids{}; + std::array expected_distances{}, actual_distances{}; + source.search( + data.data(), kTopk, 7, expected_ids.data(), expected_distances.data() + ); + source.save(path.c_str()); + target.load(path.c_str()); + EXPECT_EQ(target.initializer_type(), type); + target.search(data.data(), kTopk, 7, actual_ids.data(), actual_distances.data()); + EXPECT_EQ(actual_ids, expected_ids); + EXPECT_EQ(actual_distances, expected_distances); + { + std::fstream file(path, std::ios::binary | std::ios::in | std::ios::out); + // A persisted v2 file must contain a concrete initializer, never Auto. + file.seekp(24); + const uint32_t invalid = 0; + file.write(reinterpret_cast(&invalid), sizeof(invalid)); + } + EXPECT_THROW(target.load(path.c_str()), std::runtime_error); + EXPECT_EQ(target.initializer_type(), type); + target.search(data.data(), kTopk, 7, actual_ids.data(), actual_distances.data()); + EXPECT_EQ(actual_ids, expected_ids); + EXPECT_EQ(actual_distances, expected_distances); + } + std::remove(path.c_str()); + std::remove((path + ".hnsw").c_str()); +} + +TEST(FlatRaBitQInitializerTest, TailBatchesMetricsAndExactReturnedDistances) { + std::mt19937 generator(141); + std::normal_distribution normal; + for (size_t dim : {32U, 64U, 96U, 1056U}) { + for (size_t count : {1U, 31U, 32U, 33U, 65U}) { + std::vector centroids(count * dim); + std::generate(centroids.begin(), centroids.end(), [&] { + return normal(generator); + }); + // Unit vectors make self-query the winner under both metrics. + for (size_t i = 0; i < count; ++i) { + const float norm = std::sqrt(l2norm_sqr(centroids.data() + i * dim, dim)); + for (size_t j = 0; j < dim; ++j) + centroids[i * dim + j] /= norm; + } + for (auto metric : {METRIC_L2, METRIC_IP}) { + FlatRaBitQInitializer index(dim, count, metric); + index.add_vectors(centroids.data(), 2); + const float* query = centroids.data() + (count - 1) * dim; + for (size_t nprobe : {0U, 1U, 7U, 99U}) { + std::vector> result; + index.centroids_distances(query, nprobe, result); + ASSERT_EQ(result.size(), std::min(nprobe, count)); + if (result.empty()) + continue; + EXPECT_EQ(result[0].id, count - 1); + EXPECT_TRUE(std::is_sorted(result.begin(), result.end())); + std::vector seen(count); + for (const auto& candidate : result) { + ASSERT_LT(candidate.id, count); + EXPECT_FALSE(seen[candidate.id]); + seen[candidate.id] = true; + const float exact = + metric == METRIC_L2 + ? std::sqrt(euclidean_sqr( + query, index.centroid(candidate.id), dim + )) + : dot_product_dis(query, index.centroid(candidate.id), dim); + EXPECT_FLOAT_EQ(candidate.distance, exact); + } + } + } + } + } +} + +TEST(FlatRaBitQInitializerTest, ZeroResidualsAndInnerProductMagnitude) { + constexpr size_t kDim = 96, kCount = 33; + std::vector centroids(kDim * kCount, 0.0F), query(kDim, 0.0F); + for (auto metric : {METRIC_L2, METRIC_IP}) { + FlatRaBitQInitializer index(kDim, kCount, metric); + index.add_vectors(centroids.data(), 1); + std::vector> result; + index.centroids_distances(query.data(), 7, result); + ASSERT_EQ(result.size(), 7U); + for (const auto& item : result) + EXPECT_FLOAT_EQ(item.distance, metric == METRIC_L2 ? 0 : 1); + } + for (size_t i = 0; i < kCount; ++i) + centroids[i * kDim] = 0.5F; + centroids[17 * kDim] = 100.0F; + query[0] = 1.0F; + FlatRaBitQInitializer index(kDim, kCount, METRIC_IP); + index.add_vectors(centroids.data(), 1); + std::vector> result; + index.centroids_distances(query.data(), 1, result); + ASSERT_EQ(result.size(), 1U); + EXPECT_EQ(result[0].id, 17U); + EXPECT_FLOAT_EQ(result[0].distance, -99.0F); +} + +TEST(FlatRaBitQInitializerTest, PersistsCodesAndSupportsConcurrentSearch) { + constexpr size_t kDim = 96, kCount = 97; + std::mt19937 generator(731); + std::normal_distribution normal; + std::vector centroids(kDim * kCount), queries(8 * kDim); + std::generate(centroids.begin(), centroids.end(), [&] { return normal(generator); }); + std::generate(queries.begin(), queries.end(), [&] { return normal(generator); }); + const std::string path = ::testing::TempDir() + "flat_rabitq_initializer.bin"; + for (auto metric : {METRIC_L2, METRIC_IP}) { + FlatRaBitQInitializer index(kDim, kCount, metric), loaded(kDim, kCount, metric); + index.add_vectors(centroids.data(), 2); + { + std::ofstream out(path, std::ios::binary); + index.save(out, nullptr); + } + { + std::ifstream input(path, std::ios::binary); + loaded.load(input, nullptr); + } + for (size_t i = 0; i < kCount; ++i) { + EXPECT_TRUE( + std::equal(index.centroid(i), index.centroid(i) + kDim, loaded.centroid(i)) + ); + } + std::vector>> expected(8), actual(8); + std::vector threads; + for (size_t i = 0; i < 8; ++i) { + index.centroids_distances(queries.data() + i * kDim, 7, expected[i]); + threads.emplace_back([&, i] { + loaded.centroids_distances(queries.data() + i * kDim, 7, actual[i]); + }); + } + for (auto& thread : threads) + thread.join(); + for (size_t i = 0; i < 8; ++i) { + ASSERT_EQ(actual[i].size(), expected[i].size()); + for (size_t j = 0; j < actual[i].size(); ++j) { + EXPECT_EQ(actual[i][j].id, expected[i][j].id); + EXPECT_FLOAT_EQ(actual[i][j].distance, expected[i][j].distance); + } + } + std::ifstream before(path, std::ios::binary); + const std::string bytes{std::istreambuf_iterator(before), {}}; + before.close(); + { + std::ofstream out(path, std::ios::binary); + loaded.save(out, nullptr); + } + std::ifstream after(path, std::ios::binary); + EXPECT_EQ((std::string{std::istreambuf_iterator(after), {}}), bytes); + } + std::remove(path.c_str()); +} +} // namespace +} // namespace rabitqlib::ivf diff --git a/tests/unit/rabitqlib/index/initializer_test.cpp b/tests/unit/rabitqlib/index/initializer_test.cpp index f0804c5e..7f4238d8 100644 --- a/tests/unit/rabitqlib/index/initializer_test.cpp +++ b/tests/unit/rabitqlib/index/initializer_test.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -39,6 +40,49 @@ TEST(HNSWInitializerTest, PersistsUtf8SidecarPaths) { std::filesystem::remove(directory); } +TEST(HNSWInitializerTest, ReportsSidecarOpenFailureAndCanRetry) { + const std::string path = ::testing::TempDir() + "rabitq_hnsw_save_failure"; + const std::filesystem::path sidecar(path + ".hnsw"); + std::filesystem::create_directory(sidecar); + const std::array centroids{0.0F, 0.0F, 3.0F, 4.0F}; + HNSWInitializer initializer(2, 2); + initializer.add_vectors(centroids.data(), 1); + std::ofstream ignored_output; + try { + initializer.save(ignored_output, path.c_str()); + FAIL() << "Saving the HNSW sidecar to a directory must fail"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ(error.what(), "Cannot open HNSW index file for writing"); + } + std::filesystem::remove(sidecar); + + initializer.save(ignored_output, path.c_str()); + HNSWInitializer loaded(2, 2); + std::ifstream ignored_input; + loaded.load(ignored_input, path.c_str()); + std::vector> candidates(1); + loaded.centroids_distances(centroids.data() + 2, 1, candidates); + EXPECT_EQ(candidates[0].id, 1U); + EXPECT_FLOAT_EQ(candidates[0].distance, 0.0F); + std::filesystem::remove(sidecar); +} + +TEST(HNSWInitializerTest, ReportsBufferedWriteFailure) { + if (!std::filesystem::exists("/dev/full")) { + GTEST_SKIP() << "Requires a device that rejects writes"; + } + CentroidL2Space space(2); + hnswlib::HierarchicalNSW graph(&space, 2); + const std::array point{3.0F, 4.0F}; + graph.addPoint(point.data(), 0); + // This small index fits in the stream buffer; failure can occur at close(). + EXPECT_THROW(graph.saveIndex("/dev/full"), std::ios_base::failure); + const auto result = graph.searchKnn(point.data(), 1); + ASSERT_EQ(result.size(), 1U); + EXPECT_EQ(result.top().second, 0U); + EXPECT_FLOAT_EQ(result.top().first, 0.0F); +} + TEST(CentroidL2SpaceTest, MatchesDispatchedDistanceForUnalignedInputsAndTails) { for (size_t dim : {1U, 7U, 8U, 15U, 16U, 17U, 63U, 64U, 129U}) { CentroidL2Space space(dim); diff --git a/tests/unit/rabitqlib/index/ivf_test.cpp b/tests/unit/rabitqlib/index/ivf_test.cpp index 3bfc8c37..0d91a70d 100644 --- a/tests/unit/rabitqlib/index/ivf_test.cpp +++ b/tests/unit/rabitqlib/index/ivf_test.cpp @@ -342,7 +342,9 @@ TEST(IvfSearchTest, RoutesInnerProductQueriesWithHNSWCentroids) { query[0] = 1.0F; const PID cluster = 0; - IVF index(1, kDim, kCount, 32, METRIC_IP); + IVF index( + 1, kDim, kCount, 32, METRIC_IP, RotatorType::FhtKacRotator, InitializerType::HNSW + ); index.construct(data.data(), centroids.data(), &cluster, false, 4); PID result = kPidMax; float distance = std::numeric_limits::infinity(); @@ -803,7 +805,9 @@ TEST(IvfAddTest, RoutesNewPointsWithHNSWCentroids) { std::array first{}; first[0] = 1.0F; const PID cluster = 0; - IVF index(1, kDim, kCount, 32, METRIC_IP); + IVF index( + 1, kDim, kCount, 32, METRIC_IP, RotatorType::FhtKacRotator, InitializerType::HNSW + ); index.construct(first.data(), centroids.data(), &cluster, false, 4); std::array second{};