diff --git a/docs/development.md b/docs/development.md index 5f7d14b3..77a3f6a1 100644 --- a/docs/development.md +++ b/docs/development.md @@ -349,7 +349,10 @@ mutating state: non-finite values and magnitudes above the clustering bound rais `std::invalid_argument`. The core `quantize_*` encoders and `Lut` rely on that (the LUT additionally rejects non-finite rotated queries) because quantization converts floats to integers, which is undefined for NaN. Keep the check at the boundary rather -than inside per-vector kernels. +than inside per-vector kernels. IVF and HNSW `train` are the exception: they pass +training data only to k-means, which applies the same rule and bound in parallel +before using it, so they skip the serial pass; `set_centroids` still checks the +trained centroids. **Python conversion** diff --git a/docs/docs/index/hnsw.md b/docs/docs/index/hnsw.md index ac91f5de..16e82a0f 100644 --- a/docs/docs/index/hnsw.md +++ b/docs/docs/index/hnsw.md @@ -106,6 +106,41 @@ void HierarchicalNSW::construct( - **num_threads**: Number of threads to use (default: every available thread; `0` also auto-selects). - **fast_quantization**: Use calibrated scaling instead of per-vector scale search; see [encoding paths](../rabitq/quantizer.md#implementations). +### Training centroids in C++ + +Instead of passing centroids to `construct`, train or set them on an empty +index, then add the points: + +```cpp +#include +#include + +rabitqlib::hnsw::HierarchicalNSW index(max_elements, dim, 4, 16, 200); +index.train(16, num_points, data.data(), + rabitqlib::default_training_parameters(rabitqlib::METRIC_L2, 100)); +// Or supply centroids: index.set_centroids(num_clusters, centroids.data()); +index.add(data.data(), num_points); +``` + +- `train(num_clusters, num_points, data, cp)` trains centroids with + [RaBitQKMeans](../clustering.md#rabitqkmeans); `num_points` must be at least + `num_clusters`. Points are assigned by `add`, not by training. See the + [IVF-RaBitQ defaults](ivf.md#training-centroids-in-c) for + `default_training_parameters`. +- The first `add` into an empty graph links points with their exact vectors, as + `construct` does; with one thread it produces the same index as `construct` + with the same centroids and cluster IDs. If it fails, the graph is left empty + and the centroids are kept. Later adds link through the stored codes; see + [how added points are linked](#how-added-points-are-linked). +- `train` and `set_centroids` replace the centroids only while the index stores + no points, removed ones included; otherwise they throw `std::logic_error`. + A failed call leaves the index unchanged. `construct` needs an untrained index, + so after either call it throws `std::logic_error`; use `add` instead. +- `is_trained()` reports whether centroids are set. Before that, `add` throws + `std::logic_error`; searching an empty index, trained or not, returns no results. +- `save` throws `std::logic_error` for an index without stored points, trained + or not. Earlier releases saved an untrained empty index; such files still load. + ## Querying ```cpp @@ -131,7 +166,8 @@ Python follows the shared [concurrency contract](usage.md#threading-and-file-pat ## Updating an index A constructed or loaded index supports `add()`, `resize()`, -and `remove()` without the original data. Adding points and resizing preserve +and `remove()` without the original data; C++ also adds to a +[trained](#training-centroids-in-c) index. Adding points and resizing preserve the index file format; files with removals have [compatibility restrictions](#how-removal-is-stored). diff --git a/docs/docs/index/ivf.md b/docs/docs/index/ivf.md index 5fb9ca2a..e6ec814a 100644 --- a/docs/docs/index/ivf.md +++ b/docs/docs/index/ivf.md @@ -199,6 +199,47 @@ For example: ivf.construct(data.data(), centroids.data(), cluster_ids.data(), true); ``` +### Training centroids in C++ + +Instead of passing centroids to `construct`, create an index without a point +count, train or set its centroids, then add the points: + +```c++ +#include +#include + +rabitqlib::ivf::IVF index(dim, num_clusters, 4, rabitqlib::METRIC_L2, /*seed=*/42); +index.train(num_points, data.data(), + rabitqlib::default_training_parameters(rabitqlib::METRIC_L2, 42)); +// Or supply centroids: index.set_centroids(centroids.data()); +index.add(data.data(), num_points); +``` + +- The constructor takes `dim`, `num_clusters`, `total_bits`, the metric and the + rotation seed, followed by the optional rotator and initializer types. The + index starts with no point storage. `construct` rebuilds an index from + `max_elements()` rows, so it fails on this index until `add` has stored points. +- An index created with a point count is built by `construct`: `train` and + `set_centroids` throw `std::logic_error` for it. +- `train` uses [RaBitQKMeans](../clustering.md#rabitqkmeans) for up to + `kMaxRaBitQKMeansClusters` (2,000) clusters and + [QGKMeans](../clustering.md#qgkmeans) above that, copying the fields the two + parameter types share and keeping QGKMeans' default graph settings. The + threshold is provisional and has not been measured. Only centroids are + trained; `add` assigns the points. `cp.num_threads` also limits building the + centroid routing. +- `default_training_parameters(metric, seed)` runs 25 iterations on at most + 256 sampled points per centroid and selects spherical k-means for inner + product. The k-means classes keep their own defaults, which train on every point. +- `train` and `set_centroids` replace the centroids only while the index stores + no points, removed ones included; otherwise they throw `std::logic_error`. + A failed call leaves the index unchanged. +- `is_trained()` reports whether centroids are set. Before that, `add` and + `search` throw `std::logic_error`. A trained index without points rejects + every `k` in `search` with `std::invalid_argument`, because `k` may not exceed + the stored point count; `save` throws `std::logic_error` until the index stores + at least one point. + ## Querying To search a saved index, load it first: @@ -271,7 +312,8 @@ Concurrent Python searches may use independent `nprobe`, `high_accuracy`, and ## Updating an index C++ `IVF` and Python `IvfIndex` support `add()` and `remove()` on a constructed -or loaded index without the original dataset. Updates preserve the file format. +or loaded index without the original dataset; C++ also adds to a +[trained](#training-centroids-in-c) index. Updates preserve the file format. ```c++ void IVF::add( diff --git a/include/rabitqlib/clustering/detail/kmeans.hpp b/include/rabitqlib/clustering/detail/kmeans.hpp index 3810ecfc..5447ebc4 100644 --- a/include/rabitqlib/clustering/detail/kmeans.hpp +++ b/include/rabitqlib/clustering/detail/kmeans.hpp @@ -234,9 +234,26 @@ inline void exact_assign( return indices; } +// Copy the fields LloydKMeans reads, which every k-means parameter type shares, +// leaving the method-specific ones of `to` unchanged. Keep this list and the +// using-declarations below the same. +template +void copy_shared_parameters(const From& from, To& to) { + to.early_stop_threshold = from.early_stop_threshold; + to.final_assignment = from.final_assignment; + to.max_points_per_centroid = from.max_points_per_centroid; + to.min_points_per_centroid = from.min_points_per_centroid; + to.niter = from.niter; + to.num_threads = from.num_threads; + to.seed = from.seed; + to.spherical = from.spherical; + to.verbose = from.verbose; +} + template class LloydKMeans : public Parameters { public: + // Keep in sync with copy_shared_parameters. using Parameters::early_stop_threshold; using Parameters::final_assignment; using Parameters::max_points_per_centroid; @@ -251,7 +268,19 @@ class LloydKMeans : public Parameters { validate_parameters(); } - void train(size_t n, const float* x) { + void train(size_t n, const float* x) { fit(n, x, true); } + + /** + * Trains the same centroids as train() but skips its final assignment of all n + * vectors, which can cost more than every sampled iteration together. Leaves + * assignments and distances empty and final_obj zero; centroids, + * iteration_stats and assign() match a train() call. + */ + void train_centroids(size_t n, const float* x) { fit(n, x, false); } + + private: + // assign_all selects the final assignment that train() returns. + void fit(size_t n, const float* x, bool assign_all) { validate_parameters(); if (x == nullptr) { throw std::invalid_argument( @@ -372,6 +401,12 @@ class LloydKMeans : public Parameters { } assigner->store_model_state(assignment_state_); + if (!assign_all) { + std::vector().swap(assignments); + std::vector().swap(distances); + final_obj = 0.0; + return; + } // Return assignments for the final centroids, not the centroids from // the beginning of the last update step. if (sampled) { @@ -400,11 +435,12 @@ class LloydKMeans : public Parameters { final_obj = std::accumulate(distances.begin(), distances.end(), 0.0); } + public: /** * Assigns n row-major vectors to the current centroids with final_assignment. * Each vector's result is independent of the others in the call and of - * previous labels. Requires a successful train(); the outputs hold n elements - * each and must not overlap x or centroids. + * previous labels. Requires a successful train() or train_centroids(); the + * outputs hold n elements each and must not overlap x or centroids. */ void assign( size_t n, const float* x, PID* point_assignments, float* point_distances diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index f3906664..001ffc65 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -30,6 +30,9 @@ class VisitedSetPool; namespace detail { struct DistanceEstimate; } +namespace rabitqkmeans { +struct RaBitQKMeansParameters; +} } // namespace rabitqlib namespace rabitqlib::hnsw { @@ -73,6 +76,9 @@ class HierarchicalNSW { return num_stored_points_.load() - num_removed_; } [[nodiscard]] size_t num_stored_points() const { return num_stored_points_; } + // Centroids are set, by train, set_centroids or construct, or the index was loaded + // with clusters. + [[nodiscard]] bool is_trained() const { return num_clusters_ != 0; } // Cluster of a stored record by internal ID (0 through num_stored_points() - 1; // these coincide with labels only for points added in order). Removed points // keep their cluster. @@ -88,9 +94,39 @@ class HierarchicalNSW { return stored_cluster_id(internal_id); } + // Requires at least one stored point; a trained index must add points first. void save(const char* filename) const; void load(const char* filename); + /** + * @brief Replace the centroids of an initialized index without stored points. + * + * The index is unchanged if this throws. Later adds route and quantize against + * these centroids; the first add into the empty graph links like construct. + * + * @param num_clusters Number of centroids, from 1 through 2^31 + * @param centroids Centroid vectors (num_clusters * dim) in the original coordinates + */ + void set_centroids(size_t num_clusters, const float* centroids); + + /** + * @brief Train centroids with RaBitQKMeans and set them as set_centroids does. + * + * Only centroids are trained: points are not assigned until add. The index is + * unchanged if this throws. + * + * @param num_clusters Number of centroids, at most num_points + * @param num_points Number of training points + * @param data Training points (num_points * dim) + * @param cp Clustering parameters, for example default_training_parameters() + */ + void train( + size_t num_clusters, + size_t num_points, + const float* data, + const rabitqkmeans::RaBitQKMeansParameters& cp + ); + // Centroids and data use the original dimension; cluster_ids has num_points entries. // Input buffers must remain valid until this synchronous call returns or throws; // the index retains no pointers to them afterward. @@ -107,10 +143,12 @@ class HierarchicalNSW { ); /** - * @brief Insert points into a constructed or loaded index. + * @brief Insert points into a trained, constructed or loaded index. * - * Linked with the stored quantized codes, not the exact vectors construct uses, - * which are not kept. Centroids and rotation are never retrained. Point `i` gets + * Into an empty graph, points are linked with their exact vectors as construct + * links them, and a failure leaves the graph empty. Otherwise they are linked with + * the stored quantized codes, since the exact vectors are not kept. Centroids and + * rotation are never retrained. Point `i` gets * the label `num_stored_points() + i` whatever `num_threads` is. Throws rather than * growing past `max_elements()`. Not safe to call while another thread searches. * Preparation failures leave the graph and point count unchanged. A failure @@ -168,7 +206,8 @@ class HierarchicalNSW { // Queries use the original dimension. Each result is (distance, external label), // where distance estimates squared L2 or 1 - dot according to the index metric. - // ef_search is raised to k and clamped to the stored point count. + // ef_search is raised to k and clamped to the stored point count. An empty index, + // trained or not, returns no results. std::vector>> search( const float* queries, size_t num_queries, @@ -437,6 +476,16 @@ class HierarchicalNSW { void add_point(PID label, PID cluster_id, const quant::RabitqConfig& config); + // Link validated points into an empty graph that has centroids, scoring with their + // exact vectors. Labels are 0 through num_points - 1. A failure empties the graph. + void link_exact( + size_t num_points, + const float* data, + const PID* cluster_ids, + size_t num_threads, + bool fast_quantization + ); + // Insertion steps shared by raw and quantized scoring. distance(id) scores an // existing node against the node being inserted; prefetch(id) may warm its data. template diff --git a/include/rabitqlib/index/ivf/ivf.hpp b/include/rabitqlib/index/ivf/ivf.hpp index be504632..ef62fa90 100644 --- a/include/rabitqlib/index/ivf/ivf.hpp +++ b/include/rabitqlib/index/ivf/ivf.hpp @@ -23,6 +23,9 @@ namespace buffer { template class SearchBuffer; } +namespace rabitqkmeans { +struct RaBitQKMeansParameters; +} } // namespace rabitqlib namespace rabitqlib::ivf { @@ -83,7 +86,7 @@ class IVF { size_t padded_dim_ = 0; size_t num_clusters_ = 0; bool raw_reranking_ = false; // raw vectors replace extra-bit codes - bool ready_ = false; // construction or loading completed + bool ready_ = false; // centroids set, construction or loading completed size_t ex_bits_ = 0; // total bits = ex_bits_ + 1; 0 with raw reranking RotatorType rotator_type_ = RotatorType::FhtKacRotator; std::unique_ptr> rotator_; @@ -146,6 +149,11 @@ class IVF { [[nodiscard]] bool has_capacity(const ClusterStorage& cs, size_t num_points) const; + // Validate the settings both constructors share and derive padded_dim_ + void configure(size_t total_bits); + + [[nodiscard]] std::unique_ptr make_initializer() const; + // Reset the index and allocate a new initializer and per-cluster storage void allocate_storage(const std::vector& cluster_sizes); @@ -196,6 +204,31 @@ class IVF { InitializerType initializer_type = InitializerType::Auto ); + /** + * @brief Create an empty index to fill with train or set_centroids, then add. + * + * The index starts with no point storage. construct rebuilds an index from + * max_elements() rows, so it fails here until add has stored points. + * + * @param dim Original input dimension + * @param num_clusters Number of centroids that train or set_centroids provides + * @param total_bits 1-9 for quantized storage (including the sign bit), or 32 for + * raw-vector reranking + * @param metric_type Distance metric + * @param seed Rotation seed + * @param rotator_type Rotation type + * @param initializer_type Centroid routing strategy + */ + explicit IVF( + size_t dim, + size_t num_clusters, + size_t total_bits, + MetricType metric_type, + uint32_t seed, + RotatorType rotator_type = RotatorType::FhtKacRotator, + InitializerType initializer_type = InitializerType::Auto + ); + ~IVF(); // Logical capacity, including slots occupied by removed points. add grows it. @@ -210,6 +243,40 @@ class IVF { [[nodiscard]] MetricType metric_type() const { return metric_type_; } [[nodiscard]] RotatorType rotator_type() const { return rotator_type_; } [[nodiscard]] InitializerType initializer_type() const { return initializer_type_; } + // Centroids are set, by train, set_centroids, construct or load. + [[nodiscard]] bool is_trained() const { return ready_; } + + /** + * @brief Train num_clusters() centroids on `data` and set them as set_centroids does. + * + * Uses RaBitQKMeans up to kMaxRaBitQKMeansClusters clusters, and QGKMeans with the + * fields the two parameter types share and default graph parameters above that. + * Only centroids are trained: points are not assigned until add. Requires a + * configured index without stored points that was not created with a point + * count; the index is unchanged if this throws. + * + * @param num_points Number of training points, at least num_clusters() + * @param data Training points (num_points * dim) + * @param cp Clustering parameters, for example default_training_parameters(); + * cp.num_threads also limits building the centroid initializer + */ + void train( + size_t num_points, const float* data, const rabitqkmeans::RaBitQKMeansParameters& cp + ); + + /** + * @brief Replace the centroids of an index without stored points. + * + * Builds the centroid initializer and empty clusters, so the index can add points. + * An index created with a point count is built by construct instead, so this + * throws std::logic_error for it. The index is unchanged if this throws. + * + * @param centroids Centroid vectors (num_clusters() * dim) + * @param num_threads Threads used to build the centroid initializer + */ + void set_centroids( + const float* centroids, size_t num_threads = std::numeric_limits::max() + ); /** * @brief Quantize every point into its cluster and build the centroid initializer. @@ -233,7 +300,7 @@ class IVF { ); /** - * @brief Append points to a constructed or loaded index without the original data. + * @brief Append points to a trained, constructed or loaded index. * * Points are quantized against the centroids and rotator the index already holds, so * the centroids never move: recall can degrade if the added data drifts away from the @@ -292,6 +359,7 @@ class IVF { */ size_t remove(const PID* ids_to_remove, size_t num_points); + // Requires at least one stored point; a trained index must add points first. void save(const char* filename) const; void load(const char* filename); diff --git a/include/rabitqlib/index/training.hpp b/include/rabitqlib/index/training.hpp new file mode 100644 index 00000000..ce35dded --- /dev/null +++ b/include/rabitqlib/index/training.hpp @@ -0,0 +1,35 @@ +#pragma once + +#include +#include + +#include "rabitqlib/clustering/rabitqkmeans.hpp" // IWYU pragma: export +#include "rabitqlib/core/metric.hpp" + +namespace rabitqlib { + +/** + * Centroid training parameters for IVF::train and HierarchicalNSW::train: 25 + * iterations on a sample of at most 256 points per centroid, spherical k-means for + * inner product, and the given seed. The k-means classes keep their own defaults, + * which train on all points. + */ +inline rabitqkmeans::RaBitQKMeansParameters default_training_parameters( + MetricType metric_type, uint32_t seed +) { + validate_metric_type(metric_type); + rabitqkmeans::RaBitQKMeansParameters cp; + cp.niter = 25; + cp.max_points_per_centroid = 256; + cp.spherical = metric_type == METRIC_IP; + cp.seed = seed; + return cp; +} + +namespace ivf { +// IVF::train uses RaBitQKMeans up to this many clusters and QGKMeans above it. +// Tentative: this threshold has not been measured. +inline constexpr size_t kMaxRaBitQKMeansClusters = 2000; +} // namespace ivf + +} // namespace rabitqlib diff --git a/src/index/hnsw.cpp b/src/index/hnsw.cpp index 54cd04f5..72dc4fa8 100644 --- a/src/index/hnsw.cpp +++ b/src/index/hnsw.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #include #include #include @@ -22,6 +23,7 @@ #include #include "file_replacement.hpp" +#include "rabitqlib/clustering/rabitqkmeans.hpp" #include "rabitqlib/core/data_layout.hpp" #include "rabitqlib/core/encoder.hpp" #include "rabitqlib/core/estimator.hpp" @@ -141,6 +143,9 @@ void HierarchicalNSW::save(const char* filename) const { if (level0_record_bytes_ == 0) { throw std::logic_error("HNSW index must be initialized before save"); } + if (num_stored_points_ == 0) { + throw std::logic_error("HNSW index has no points to save; add points first"); + } if (filename == nullptr || filename[0] == '\0') { throw std::invalid_argument("HNSW index filename must not be empty"); } @@ -499,7 +504,7 @@ void HierarchicalNSW::construct( bool fast_quantization ) { if (num_clusters_ != 0 || num_stored_points_ != 0) { - throw std::logic_error("HNSW index is already constructed or loaded"); + throw std::logic_error("HNSW index is already trained, constructed or loaded"); } if (num_clusters == 0 || num_clusters > buffer::kSearchBufferMaxPointCount) { throw std::invalid_argument("HNSW cluster count is out of range"); @@ -534,6 +539,23 @@ void HierarchicalNSW::construct( ); } + link_exact(num_points, data, cluster_ids, num_threads, fast_quantization); + } catch (...) { + // Also discard the centroids, retaining capacity and rotation for a retry. + centroids_memory_.reset(); + num_clusters_ = 0; + throw; + } +} + +void HierarchicalNSW::link_exact( + size_t num_points, + const float* data, + const PID* cluster_ids, + size_t num_threads, + bool fast_quantization +) { + try { quant::RabitqConfig config; if (fast_quantization) { config = quant::faster_config(padded_dim_, ex_bits_ + 1); @@ -547,7 +569,7 @@ void HierarchicalNSW::construct( ); } catch (...) { // parallel_for joins every worker before propagating an exception. - // Discard the partial graph, retaining capacity and rotation for a retry. + // Discard the partial graph, retaining capacity, rotation and centroids. std::fill_n(element_levels_.begin(), num_stored_points_.load(), 0); link_lists_.clear(); num_stored_points_ = 0; @@ -555,14 +577,75 @@ void HierarchicalNSW::construct( label_lookup_.clear(); entry_point_id_ = kPidMax; max_level_ = -1; - centroids_memory_.reset(); - num_clusters_ = 0; construction_data_ = nullptr; throw; } construction_data_ = nullptr; } +void HierarchicalNSW::set_centroids(size_t num_clusters, const float* centroids) { + if (level0_record_bytes_ == 0) { + throw std::logic_error("HNSW index must be initialized before setting centroids"); + } + if (num_stored_points_ != 0) { + throw std::logic_error("HNSW cannot replace the centroids of an index with points"); + } + if (num_clusters == 0 || num_clusters > buffer::kSearchBufferMaxPointCount) { + throw std::invalid_argument("HNSW cluster count is out of range"); + } + if (centroids == nullptr) { + throw std::invalid_argument("HNSW centroids must not be null"); + } + if (num_clusters > std::numeric_limits::max() / sizeof(float) / padded_dim_) { + throw std::invalid_argument("HNSW centroids exceed addressable storage"); + } + check_coordinates(centroids, num_clusters * dim_, dim_, "HNSW centroids"); + + // construct rotates centroids the same way, so both paths quantize identically. + AlignedStorage rotated( + memory::huge_page_allocate(num_clusters * padded_dim_ * sizeof(float)) + ); + for (size_t i = 0; i < num_clusters; ++i) { + rotator_->rotate( + centroids + (i * dim_), + reinterpret_cast(rotated.get()) + (i * padded_dim_) + ); + } + centroids_memory_ = std::move(rotated); + num_clusters_ = num_clusters; +} + +void HierarchicalNSW::train( + size_t num_clusters, + size_t num_points, + const float* data, + const rabitqkmeans::RaBitQKMeansParameters& cp +) { + if (level0_record_bytes_ == 0) { + throw std::logic_error("HNSW index must be initialized before training"); + } + if (num_stored_points_ != 0) { + throw std::logic_error("HNSW cannot retrain an index with points"); + } + if (num_clusters == 0 || num_clusters > buffer::kSearchBufferMaxPointCount) { + throw std::invalid_argument("HNSW cluster count is out of range"); + } + if (data == nullptr) { + throw std::invalid_argument("HNSW training data must not be null"); + } + if (num_points < num_clusters) { + throw std::invalid_argument("HNSW training needs at least num_clusters points"); + } + if (num_points > std::numeric_limits::max() / sizeof(float) / dim_) { + throw std::invalid_argument("HNSW training data exceeds addressable storage"); + } + // k-means checks the data for finite, in-range coordinates in parallel with + // the bound check_coordinates uses; set_centroids checks its centroids. + rabitqkmeans::RaBitQKMeans kmeans(dim_, num_clusters, cp); + kmeans.train_centroids(num_points, data); + set_centroids(num_clusters, kmeans.centroids.data()); +} + template PID HierarchicalNSW::greedy_descend( PID current_id, int from_level, int to_level, Distance distance @@ -1292,7 +1375,9 @@ std::vector HierarchicalNSW::add( ) { if (data_level0_memory_ == nullptr || centroids_memory_ == nullptr || rotator_ == nullptr || num_clusters_ == 0) { - throw std::logic_error("HNSW index must be constructed or loaded before add"); + throw std::logic_error( + "HNSW index must be trained, constructed or loaded before add" + ); } if (num_points == 0) { return {}; @@ -1336,6 +1421,13 @@ std::vector HierarchicalNSW::add( }); } + if (num_stored_points_ == 0) { + // The exact vectors are at hand, so link them as construct would. + link_exact(num_points, data, assigned.data(), num_threads, fast_quantization); + std::iota(labels.begin(), labels.end(), PID{0}); + return labels; + } + quant::RabitqConfig config; if (fast_quantization) { config = quant::faster_config(padded_dim_, ex_bits_ + 1); diff --git a/src/index/ivf.cpp b/src/index/ivf.cpp index 3de7bb81..0984d5bd 100644 --- a/src/index/ivf.cpp +++ b/src/index/ivf.cpp @@ -22,6 +22,8 @@ #include "file_replacement.hpp" #include "ivf_search.hpp" +#include "rabitqlib/clustering/qgkmeans.hpp" +#include "rabitqlib/clustering/rabitqkmeans.hpp" #include "rabitqlib/core/data_layout.hpp" #include "rabitqlib/core/encoder.hpp" #include "rabitqlib/core/estimator.hpp" @@ -37,6 +39,7 @@ #include "rabitqlib/index/detail/parallel_for.hpp" #include "rabitqlib/index/ivf/cluster.hpp" #include "rabitqlib/index/ivf/initializer.hpp" +#include "rabitqlib/index/training.hpp" #include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/input_checks.hpp" #include "rabitqlib/utils/path.hpp" @@ -50,6 +53,9 @@ namespace { constexpr float kRemovedFAdd = std::numeric_limits::infinity(); // Files without routing metadata used HNSW routing from this many clusters. constexpr size_t kHistoricalHnswMinClusters = 20000; +// An index from the point-count constructor is built by construct, never trained. +constexpr const char* kPointCountConstructOnly = + "IVF index created with a point count must be built by construct"; } // namespace IVF::IVF() noexcept = default; @@ -162,28 +168,149 @@ IVF::IVF( , ex_bits_(total_bits == 32 ? 0 : total_bits - 1) , rotator_type_(rotator_type) , metric_type_(metric_type) { - validate_metric_type(metric_type); - if ((total_bits < 1 || total_bits > 9) && total_bits != 32) { - throw std::invalid_argument("IVF bits must be in [1, 9] or 32 for raw reranking"); - } + configure(total_bits); if (num_points == 0 || num_points > buffer::kSearchBufferMaxPointCount) { throw std::invalid_argument("IVF point count exceeds the supported ID range"); } - if (num_clusters == 0 || num_clusters > buffer::kSearchBufferMaxPointCount) { + if (exceeds_storage(num_points)) { + throw std::invalid_argument("IVF configuration exceeds addressable storage"); + } + rotator_.reset(choose_rotator(dim, rotator_type, padded_dim_)); +} + +IVF::IVF( + size_t dim, + size_t num_clusters, + size_t total_bits, + MetricType metric_type, + uint32_t seed, + RotatorType rotator_type, + InitializerType initializer_type +) + : initializer_type_(resolve_initializer_type(initializer_type, num_clusters)) + , dim_(dim) + , padded_dim_(dim) + , num_clusters_(num_clusters) + , raw_reranking_(total_bits == 32) + , ex_bits_(total_bits == 32 ? 0 : total_bits - 1) + , rotator_type_(rotator_type) + , metric_type_(metric_type) { + configure(total_bits); + rotator_.reset(choose_rotator(dim, rotator_type, padded_dim_, seed)); +} + +void IVF::configure(size_t total_bits) { + validate_metric_type(metric_type_); + if ((total_bits < 1 || total_bits > 9) && total_bits != 32) { + throw std::invalid_argument("IVF bits must be in [1, 9] or 32 for raw reranking"); + } + if (num_clusters_ == 0 || num_clusters_ > buffer::kSearchBufferMaxPointCount) { throw std::invalid_argument("IVF cluster count exceeds the supported ID range"); } - if (dim == 0 || dim > (std::numeric_limits::max() / 32) - 64) { + if (dim_ == 0 || dim_ > (std::numeric_limits::max() / 32) - 64) { throw std::invalid_argument("IVF dimension is invalid or too large"); } padded_dim_ = round_up_to_multiple(dim_, 32); const size_t max_size = std::numeric_limits::max(); - if (exceeds_storage(num_points) || - num_clusters > max_size / sizeof(float) / padded_dim_ || - (rotator_type == RotatorType::MatrixRotator && - dim > max_size / sizeof(float) / padded_dim_)) { + if (num_clusters_ > max_size / sizeof(float) / padded_dim_ || + (rotator_type_ == RotatorType::MatrixRotator && + dim_ > max_size / sizeof(float) / padded_dim_)) { throw std::invalid_argument("IVF configuration exceeds addressable storage"); } - rotator_.reset(choose_rotator(dim, rotator_type, padded_dim_)); +} + +std::unique_ptr IVF::make_initializer() const { + if (initializer_type_ == InitializerType::Flat) { + return std::make_unique(padded_dim_, num_clusters_, metric_type_); + } + if (initializer_type_ == InitializerType::FlatRaBitQ) { + return std::make_unique( + padded_dim_, num_clusters_, metric_type_ + ); + } + return std::make_unique(padded_dim_, num_clusters_, metric_type_); +} + +void IVF::set_centroids(const float* centroids, size_t num_threads) { + if (rotator_ == nullptr) { + throw std::logic_error("IVF must be configured before setting centroids"); + } + if (num_stored_points() != 0) { + throw std::logic_error("IVF cannot replace the centroids of an index with points"); + } + if (!ready_ && max_elements_ != 0) { + throw std::logic_error(kPointCountConstructOnly); + } + if (centroids == nullptr) { + throw std::invalid_argument("IVF centroids must not be null"); + } + check_coordinates(centroids, num_clusters_ * dim_, dim_, "IVF centroids"); + + // construct rotates centroids the same way, so both paths route identically. + std::vector rotated_centroids(num_clusters_ * padded_dim_); + for (size_t i = 0; i < num_clusters_; ++i) { + rotator_->rotate(centroids + (i * dim_), &rotated_centroids[i * padded_dim_]); + } + auto initializer = make_initializer(); + initializer->add_vectors( + rotated_centroids.data(), std::min(resolve_num_threads(num_threads), num_clusters_) + ); + std::vector cluster_storage( + num_clusters_, ClusterStorage(raw_reranking_) + ); + std::vector clusters; + clusters.reserve(num_clusters_); + for (auto& cs : cluster_storage) { + clusters.emplace_back( + 0, cs.batch_data(), has_rerank_data() ? cs.rerank_data() : nullptr, cs.ids() + ); + } + const auto ip_func = select_excode_ipfunc(ex_bits_); + + // Commit. Moving the vectors keeps the empty views' storage, and nothing throws. + initializer_ = std::move(initializer); + cluster_storage_ = std::move(cluster_storage); + clusters_ = std::move(clusters); + ip_func_ = ip_func; + ready_ = true; +} + +void IVF::train( + size_t num_points, const float* data, const rabitqkmeans::RaBitQKMeansParameters& cp +) { + if (rotator_ == nullptr) { + throw std::logic_error("IVF must be configured before training"); + } + if (num_stored_points() != 0) { + throw std::logic_error("IVF cannot retrain an index with points"); + } + if (!ready_ && max_elements_ != 0) { + throw std::logic_error(kPointCountConstructOnly); + } + if (data == nullptr) { + throw std::invalid_argument("IVF training data must not be null"); + } + if (num_points < num_clusters_) { + throw std::invalid_argument("IVF training needs at least num_clusters points"); + } + if (num_points > std::numeric_limits::max() / sizeof(float) / dim_) { + throw std::invalid_argument("IVF training data exceeds addressable storage"); + } + // k-means checks the data for finite, in-range coordinates in parallel with + // the bound check_coordinates uses; set_centroids checks its centroids. + std::vector centroids; + if (num_clusters_ <= kMaxRaBitQKMeansClusters) { + rabitqkmeans::RaBitQKMeans kmeans(dim_, num_clusters_, cp); + kmeans.train_centroids(num_points, data); + centroids = std::move(kmeans.centroids); + } else { + qgkmeans::QGKMeansParameters qg_cp; // Default graph settings + clustering::detail::copy_shared_parameters(cp, qg_cp); + qgkmeans::QGKMeans kmeans(dim_, num_clusters_, qg_cp); + kmeans.train_centroids(num_points, data); + centroids = std::move(kmeans.centroids); + } + set_centroids(centroids.data(), cp.num_threads); } void IVF::construct( @@ -258,17 +385,7 @@ void IVF::allocate_storage(const std::vector& cluster_sizes) { free_storage(); clusters_.clear(); - if (initializer_type_ == InitializerType::Flat) { - initializer_ = - std::make_unique(padded_dim_, num_clusters_, metric_type_); - } else if (initializer_type_ == InitializerType::FlatRaBitQ) { - initializer_ = std::make_unique( - padded_dim_, num_clusters_, metric_type_ - ); - } else { - initializer_ = - std::make_unique(padded_dim_, num_clusters_, metric_type_); - } + initializer_ = make_initializer(); 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]); @@ -344,7 +461,9 @@ void IVF::add( size_t num_threads ) { if (!ready_) { - throw std::logic_error("IVF index must be constructed or loaded before add"); + throw std::logic_error( + "IVF index must be trained, constructed or loaded before add" + ); } if (num_points == 0) { return; @@ -686,6 +805,9 @@ void IVF::save(const char* filename) const { if (!ready_) { throw std::logic_error("IVF index must be constructed or loaded before save"); } + if (max_elements_ == 0) { + throw std::logic_error("IVF index has no points to save; add points first"); + } if (filename == nullptr || filename[0] == '\0') { throw std::invalid_argument("IVF save filename must not be empty"); } @@ -1088,7 +1210,9 @@ inline void IVF::validate_search( const float* query, size_t k, size_t nprobe, const PID* result_ids ) const { if (!ready_) { - throw std::logic_error("IVF index must be constructed or loaded before search"); + throw std::logic_error( + "IVF index must be trained, constructed or loaded before search" + ); } if (query == nullptr) { throw std::invalid_argument("IVF search query must not be null"); diff --git a/tests/failures/hnsw_add_allocation_test.cpp b/tests/failures/hnsw_add_allocation_test.cpp index f8003666..6f28f99e 100644 --- a/tests/failures/hnsw_add_allocation_test.cpp +++ b/tests/failures/hnsw_add_allocation_test.cpp @@ -151,6 +151,48 @@ TEST(HnswConstructionAllocationTest, EveryAllocationFailureLeavesEmptyReusableIn } } +TEST(HnswFirstAddAllocationTest, EveryAllocationFailureKeepsCentroidsAndAnEmptyGraph) { + Fixture fixture; + for (const size_t threads : {1U, 4U}) { + SCOPED_TRACE(threads); + size_t failures = 0; + bool succeeded = false; + for (size_t allocation = 0; allocation < 512; ++allocation) { + SCOPED_TRACE(allocation); + HierarchicalNSW index(kCapacity, kDim, 4, 2, 10, 1); + index.set_centroids(1, fixture.centroid.data()); + bool failed = false; + { + AllocationFailure failure(allocation); + try { + index.add(fixture.data.data(), kAdded, nullptr, false, threads); + } catch (const std::bad_alloc&) { failed = true; } + } + if (!failed) { + succeeded = true; + break; + } + ++failures; + EXPECT_EQ(index.num_stored_points(), 0U); + EXPECT_TRUE(index.is_trained()); + EXPECT_EQ(index.num_clusters(), 1U); + const auto empty = index.search(fixture.data.data(), 1, 1, kCapacity, 1); + ASSERT_EQ(empty.size(), 1U); + EXPECT_TRUE(empty[0].empty()); + ASSERT_NO_THROW( + index.add(fixture.data.data(), kAdded, nullptr, false, threads) + ); + const auto results = index.search(fixture.data.data(), kAdded, 1, kCapacity, 1); + for (size_t i = 0; i < kAdded; ++i) { + ASSERT_EQ(results[i].size(), 1U); + EXPECT_EQ(results[i][0].second, i); + } + } + EXPECT_GT(failures, 0U); + EXPECT_TRUE(succeeded); + } +} + TEST(HnswResizeAllocationTest, FailurePreservesEveryLinkOwner) { Fixture fixture; size_t failures = 0; diff --git a/tests/python/test_hnsw.py b/tests/python/test_hnsw.py index eaf3fad8..a0ec3e00 100644 --- a/tests/python/test_hnsw.py +++ b/tests/python/test_hnsw.py @@ -43,7 +43,7 @@ def test_rejected_rebuild_preserves_index(base_data, clusters, tmp_path, reload) idx = HnswIndex.load(path) before = idx.search(base_data[:10], k=3, ef=_EF) with pytest.raises( - RuntimeError, match="^HNSW index is already constructed or loaded$" + RuntimeError, match="^HNSW index is already trained, constructed or loaded$" ): idx.build(base_data, centroids[:1], np.zeros(N_VECTORS, dtype=np.uint32)) assert idx.is_built diff --git a/tests/unit/rabitqlib/clustering/kmeans_sampling_test.cpp b/tests/unit/rabitqlib/clustering/kmeans_sampling_test.cpp index 3ed8a515..93b73fb5 100644 --- a/tests/unit/rabitqlib/clustering/kmeans_sampling_test.cpp +++ b/tests/unit/rabitqlib/clustering/kmeans_sampling_test.cpp @@ -630,5 +630,99 @@ TEST(KMeansSamplingTest, QGSampledRetrainingOnOwnCentroids) { check_sampled_retraining_on_own_centroids(); } +// train_centroids trains what train trains, with and without sampling, but assigns +// nothing; assign still works afterwards. +template +void check_train_centroids_skips_final_assignment() { + constexpr size_t kNumPoints = 600; + const auto points = clustered_points(kNumPoints, 23); + const size_t num_clusters = kClusters; + for (const size_t max_points_per_centroid : {size_t{0}, size_t{8}}) { + for (const auto mode : + {FinalAssignmentMode::Approximate, FinalAssignmentMode::Exact}) { + SCOPED_TRACE( + ::testing::Message() + << "max_points_per_centroid=" << max_points_per_centroid + << " exact=" << (mode == FinalAssignmentMode::Exact) + ); + auto parameters = test_parameters(); + parameters.max_points_per_centroid = max_points_per_centroid; + parameters.final_assignment = mode; + Clustering trained(kDim, num_clusters, parameters); + trained.train(kNumPoints, points.data()); + Clustering centroids_only(kDim, num_clusters, parameters); + centroids_only.train_centroids(kNumPoints, points.data()); + + EXPECT_EQ(centroids_only.centroids, trained.centroids); + ASSERT_EQ( + centroids_only.iteration_stats.size(), trained.iteration_stats.size() + ); + for (size_t i = 0; i < trained.iteration_stats.size(); ++i) { + EXPECT_EQ( + centroids_only.iteration_stats[i].obj, trained.iteration_stats[i].obj + ); + } + EXPECT_TRUE(centroids_only.assignments.empty()); + EXPECT_TRUE(centroids_only.distances.empty()); + EXPECT_EQ(centroids_only.final_obj, 0.0); + + std::vector expected_ids(kNumPoints), actual_ids(kNumPoints); + std::vector expected_distances(kNumPoints), actual_distances(kNumPoints); + trained.assign( + kNumPoints, points.data(), expected_ids.data(), expected_distances.data() + ); + centroids_only.assign( + kNumPoints, points.data(), actual_ids.data(), actual_distances.data() + ); + EXPECT_EQ(actual_ids, expected_ids); + EXPECT_EQ(actual_distances, expected_distances); + } + } +} + +// Every shared field gets a non-default value; QG graph settings must survive. +TEST(KMeansParametersTest, CopySharedParametersCopiesOnlySharedFields) { + rabitqkmeans::RaBitQKMeansParameters from; + from.early_stop_threshold = 0.25; + from.final_assignment = FinalAssignmentMode::Exact; + from.max_points_per_centroid = 17; + from.min_points_per_centroid = 3; + from.niter = 7; + from.num_threads = 5; + from.seed = 11; + from.spherical = true; + from.verbose = true; + qgkmeans::QGKMeansParameters to; + to.graph_degree = 64; + to.ef_build = 99; + to.ef_search = 33; + to.graph_build_iterations = 2; + to.quantization_bits = 4; + + clustering::detail::copy_shared_parameters(from, to); + EXPECT_EQ(to.early_stop_threshold, from.early_stop_threshold); + EXPECT_EQ(to.final_assignment, from.final_assignment); + EXPECT_EQ(to.max_points_per_centroid, from.max_points_per_centroid); + EXPECT_EQ(to.min_points_per_centroid, from.min_points_per_centroid); + EXPECT_EQ(to.niter, from.niter); + EXPECT_EQ(to.num_threads, from.num_threads); + EXPECT_EQ(to.seed, from.seed); + EXPECT_EQ(to.spherical, from.spherical); + EXPECT_EQ(to.verbose, from.verbose); + EXPECT_EQ(to.graph_degree, 64U); + EXPECT_EQ(to.ef_build, 99U); + EXPECT_EQ(to.ef_search, 33U); + EXPECT_EQ(to.graph_build_iterations, 2U); + EXPECT_EQ(to.quantization_bits, 4U); +} + +TEST(KMeansTrainCentroidsTest, RaBitQSkipsFinalAssignment) { + check_train_centroids_skips_final_assignment(); +} + +TEST(KMeansTrainCentroidsTest, QGSkipsFinalAssignment) { + check_train_centroids_skips_final_assignment(); +} + } // namespace } // namespace rabitqlib diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index b59546df..8835a2c4 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -21,10 +21,13 @@ #include #include +#include "rabitqlib/clustering/rabitqkmeans.hpp" #include "rabitqlib/core/encoder.hpp" #include "rabitqlib/core/metric.hpp" +#include "rabitqlib/index/detail/buffer.hpp" #include "rabitqlib/index/detail/hnsw_dispatch.hpp" #include "rabitqlib/index/hnsw/hnsw_quant.hpp" +#include "rabitqlib/index/training.hpp" #include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/cpu_features.hpp" #include "rabitqlib/utils/input_checks.hpp" @@ -61,6 +64,12 @@ struct HnswTestAccess { return {neighbors + 1, neighbors + 1 + index.link_count(neighbors)}; } + // Rotated centroids as the index quantizes with them. + static std::vector centroids(const HierarchicalNSW& index) { + const auto* first = reinterpret_cast(index.centroids_memory_.get()); + return {first, first + (index.num_clusters_ * index.padded_dim_)}; + } + static PID entry_point_label(const HierarchicalNSW& index) { return index.external_label(index.entry_point_id_); } @@ -875,7 +884,9 @@ TEST_F(HnswSaveTest, RejectedRebuildPreservesConstructedAndLoadedIndexes) { ); FAIL() << "A second construction must be rejected"; } catch (const std::logic_error& error) { - EXPECT_STREQ(error.what(), "HNSW index is already constructed or loaded"); + EXPECT_STREQ( + error.what(), "HNSW index is already trained, constructed or loaded" + ); } EXPECT_EQ(index->num_points(), kCount); EXPECT_EQ(index->num_clusters(), 2U); @@ -1476,6 +1487,242 @@ TEST(HnswSearchTest, ClampsEfToStoredPointsAndChecksClusterLookups) { } } +// Two centroids and alternating cluster IDs for AddFixture data. +struct TrainFixture : AddFixture { + std::vector centroids; + + TrainFixture(size_t total, unsigned seed) + : AddFixture(total, seed), centroids(2 * kDim, 0.0F) { + centroids[0] = 0.5F; + centroids[kDim] = -0.5F; + for (size_t i = 0; i < total; ++i) { + cluster_ids[i] = static_cast(i % 2); + } + } +}; + +TEST(HnswTrainTest, UntrainedIndexRejectsAddAndSaveWithoutChangingState) { + constexpr size_t kCount = 64; + TrainFixture fixture(kCount, 101); + HierarchicalNSW index(kCount, AddFixture::kDim, 4, 8, 50); + EXPECT_FALSE(index.is_trained()); + + try { + index.add(fixture.data.data(), kCount, nullptr, false, 1); + FAIL() << "adding to an untrained index succeeded"; + } catch (const std::logic_error& error) { + EXPECT_STREQ( + error.what(), "HNSW index must be trained, constructed or loaded before add" + ); + } + const rabitq_test::TempFile path_file("rabitq_hnsw_untrained.index"); + try { + index.save(path_file.path().c_str()); + FAIL() << "saving an index without points succeeded"; + } catch (const std::logic_error& error) { + EXPECT_STREQ(error.what(), "HNSW index has no points to save; add points first"); + } + EXPECT_FALSE(std::filesystem::exists(std::filesystem::u8path(path_file.path()))); + // Searching an empty index, trained or not, returns no results. + const auto results = index.search(fixture.data.data(), 1, 1, 10, 1); + ASSERT_EQ(results.size(), 1U); + EXPECT_TRUE(results[0].empty()); + + EXPECT_FALSE(index.is_trained()); + EXPECT_EQ(index.num_stored_points(), 0U); + EXPECT_EQ(index.max_elements(), kCount); + index.set_centroids(2, fixture.centroids.data()); + auto cluster_ids = fixture.cluster_ids; + try { + index.construct( + 2, fixture.centroids.data(), kCount, fixture.data.data(), cluster_ids.data(), 1 + ); + FAIL() << "constructing a trained index succeeded"; + } catch (const std::logic_error& error) { + EXPECT_STREQ(error.what(), "HNSW index is already trained, constructed or loaded"); + } + EXPECT_EQ(index.num_stored_points(), 0U); + index.add(fixture.data.data(), kCount, nullptr, false, 1); + EXPECT_EQ(index.num_points(), kCount); +} + +TEST(HnswTrainTest, SetCentroidsAndFirstAddMatchConstructByteForByte) { + constexpr size_t kCount = 160; + TrainFixture fixture(kCount, 103); + const rabitq_test::TempFile path_file("rabitq_hnsw_set_centroids.index"); + const std::string& path = path_file.path(); + + for (auto metric : {METRIC_L2, METRIC_IP}) { + for (size_t bits : {1UL, 4UL, 9UL}) { + for (const bool fast_quantization : {false, true}) { + SCOPED_TRACE( + ::testing::Message() + << "metric=" << (metric == METRIC_L2 ? "l2" : "ip") << " bits=" << bits + << " fast=" << fast_quantization + ); + auto cluster_ids = fixture.cluster_ids; + HierarchicalNSW constructed( + kCount, AddFixture::kDim, bits, 8, 50, 7, metric + ); + constructed.construct( + 2, + fixture.centroids.data(), + kCount, + fixture.data.data(), + cluster_ids.data(), + 1, + fast_quantization + ); + const auto expected = saved_bytes(constructed, path); + + HierarchicalNSW added(kCount, AddFixture::kDim, bits, 8, 50, 7, metric); + added.set_centroids(2, fixture.centroids.data()); + EXPECT_TRUE(added.is_trained()); + const auto labels = added.add( + fixture.data.data(), + kCount, + fixture.cluster_ids.data(), + fast_quantization, + 1 + ); + ASSERT_EQ(labels.size(), kCount); + for (size_t i = 0; i < kCount; ++i) { + EXPECT_EQ(labels[i], i); + } + EXPECT_TRUE(saved_bytes(added, path) == expected); + } + } + } +} + +TEST(HnswTrainTest, ReplacesCentroidsOnlyWhileEmpty) { + constexpr size_t kCount = 96; + TrainFixture fixture(kCount, 107); + const auto cp = default_training_parameters(METRIC_L2, 7); + auto second = fixture.centroids; + std::reverse(second.begin(), second.end()); + const rabitq_test::TempFile path_file("rabitq_hnsw_retrain.index"); + const std::string& path = path_file.path(); + + HierarchicalNSW index(kCount, AddFixture::kDim, 4, 8, 50); + index.set_centroids(2, fixture.centroids.data()); + index.train(4, kCount, fixture.data.data(), cp); + EXPECT_EQ(index.num_clusters(), 4U); + index.set_centroids(2, second.data()); + HierarchicalNSW reference(kCount, AddFixture::kDim, 4, 8, 50); + reference.set_centroids(2, second.data()); + index.add(fixture.data.data(), kCount, nullptr, false, 1); + reference.add(fixture.data.data(), kCount, nullptr, false, 1); + // Only the last centroids remain. + const auto saved = saved_bytes(index, path); + EXPECT_TRUE(saved_bytes(reference, path) == saved); + + HierarchicalNSW loaded; + loaded.load(path.c_str()); + for (HierarchicalNSW* holder : {&index, &loaded}) { + try { + holder->train(2, kCount, fixture.data.data(), cp); + FAIL() << "training an index with points succeeded"; + } catch (const std::logic_error& error) { + EXPECT_STREQ(error.what(), "HNSW cannot retrain an index with points"); + } + try { + holder->set_centroids(2, fixture.centroids.data()); + FAIL() << "replacing the centroids of an index with points succeeded"; + } catch (const std::logic_error& error) { + EXPECT_STREQ( + error.what(), "HNSW cannot replace the centroids of an index with points" + ); + } + EXPECT_TRUE(saved_bytes(*holder, path) == saved); + } +} + +TEST(HnswTrainTest, RejectsInvalidCentroidsAndTrainingDataWithoutChangingState) { + constexpr size_t kCount = 64; + TrainFixture fixture(kCount, 109); + auto bad_centroids = fixture.centroids; + bad_centroids[3] = std::numeric_limits::quiet_NaN(); + auto bad_points = fixture.data; + bad_points[5] = std::numeric_limits::infinity(); + const auto cp = default_training_parameters(METRIC_L2, 7); + auto bad_cp = cp; + bad_cp.niter = 0; + const float* data = fixture.data.data(); + + HierarchicalNSW index(kCount, AddFixture::kDim, 4, 8, 50); + const auto expect_rejected = [&] { + for (const size_t count : {size_t{0}, buffer::kSearchBufferMaxPointCount + 1}) { + try { + index.set_centroids(count, fixture.centroids.data()); + FAIL() << "setting " << count << " centroids succeeded"; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ(error.what(), "HNSW cluster count is out of range"); + } + EXPECT_THROW(index.train(count, kCount, data, cp), std::invalid_argument); + } + EXPECT_THROW(index.set_centroids(2, nullptr), std::invalid_argument); + EXPECT_THROW(index.set_centroids(2, bad_centroids.data()), std::invalid_argument); + EXPECT_THROW(index.train(2, kCount, nullptr, cp), std::invalid_argument); + try { + index.train(kCount + 1, kCount, data, cp); + FAIL() << "training with fewer points than clusters succeeded"; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ(error.what(), "HNSW training needs at least num_clusters points"); + } + EXPECT_THROW(index.train(2, kCount, bad_points.data(), cp), std::invalid_argument); + EXPECT_THROW(index.train(2, kCount, data, bad_cp), std::invalid_argument); + }; + expect_rejected(); + EXPECT_FALSE(index.is_trained()); + + index.set_centroids(2, fixture.centroids.data()); + const auto before = HnswTestAccess::centroids(index); + expect_rejected(); + // A rejected first add also keeps the centroids. + EXPECT_THROW( + index.add(bad_points.data(), kCount, nullptr, false, 1), std::invalid_argument + ); + EXPECT_EQ(index.num_clusters(), 2U); + EXPECT_EQ(index.num_stored_points(), 0U); + EXPECT_EQ(HnswTestAccess::centroids(index), before); + + HierarchicalNSW uninitialized; + EXPECT_THROW( + uninitialized.set_centroids(2, fixture.centroids.data()), std::logic_error + ); + EXPECT_THROW(uninitialized.train(2, kCount, data, cp), std::logic_error); +} + +TEST(HnswTrainTest, TrainsRaBitQKMeansCentroidsAndSearchesAfterAdd) { + constexpr size_t kCount = 256; + constexpr size_t kNumClusters = 4; + TrainFixture fixture(kCount, 113); + auto cp = default_training_parameters(METRIC_L2, 7); + cp.num_threads = 1; + + HierarchicalNSW index(kCount, AddFixture::kDim, 4, 8, 50); + index.train(kNumClusters, kCount, fixture.data.data(), cp); + EXPECT_TRUE(index.is_trained()); + EXPECT_EQ(index.num_points(), 0U); + + rabitqkmeans::RaBitQKMeans kmeans(AddFixture::kDim, kNumClusters, cp); + kmeans.train(kCount, fixture.data.data()); + HierarchicalNSW expected(kCount, AddFixture::kDim, 4, 8, 50); + expected.set_centroids(kNumClusters, kmeans.centroids.data()); + EXPECT_EQ(HnswTestAccess::centroids(index), HnswTestAccess::centroids(expected)); + + // Several threads link the first batch; results then need not be reproducible. + index.add(fixture.data.data(), kCount, nullptr, false, 4); + EXPECT_EQ(index.num_points(), kCount); + const auto results = index.search(fixture.data.data(), kCount, 1, 64, 1); + size_t found = 0; + for (size_t i = 0; i < kCount; ++i) { + found += !results[i].empty() && results[i][0].second == i; + } + EXPECT_GE(found, kCount * 95 / 100); +} + } // namespace } // namespace rabitqlib::hnsw diff --git a/tests/unit/rabitqlib/index/ivf_test.cpp b/tests/unit/rabitqlib/index/ivf_test.cpp index 0585e2ae..93939d5d 100644 --- a/tests/unit/rabitqlib/index/ivf_test.cpp +++ b/tests/unit/rabitqlib/index/ivf_test.cpp @@ -22,10 +22,13 @@ #include #include "ivf_search.hpp" +#include "rabitqlib/clustering/qgkmeans.hpp" +#include "rabitqlib/clustering/rabitqkmeans.hpp" #include "rabitqlib/core/metric.hpp" #include "rabitqlib/core/rotator.hpp" #include "rabitqlib/index/detail/buffer.hpp" #include "rabitqlib/index/ivf/initializer.hpp" +#include "rabitqlib/index/training.hpp" #include "rabitqlib/index/types.hpp" #include "rabitqlib/utils/input_checks.hpp" #include "temp_file.hpp" @@ -45,6 +48,12 @@ struct IvfTestAccess { } static std::unique_ptr>& rotator(IVF& index) { return index.rotator_; } + + // Rotated centroid as the initializer routes with it. + static std::vector centroid(const IVF& index, PID i) { + const float* centroid = index.initializer_->centroid(i); + return {centroid, centroid + index.padded_dim()}; + } }; namespace { @@ -1587,6 +1596,323 @@ TEST(IvfPersistenceTest, RejectsTrailingBytesAndReportsUnopenableDestination) { } } +std::vector all_centroids(const IVF& index) { + std::vector centroids; + for (size_t i = 0; i < index.num_clusters(); ++i) { + const auto centroid = IvfTestAccess::centroid(index, static_cast(i)); + centroids.insert(centroids.end(), centroid.begin(), centroid.end()); + } + return centroids; +} + +TEST(IvfTrainTest, UntrainedIndexRejectsAddSearchAndSaveWithoutChangingState) { + constexpr size_t kCount = 64; + const auto points = random_rows(kCount, 31); + IVF index(kDynDim, kDynNumClusters, 4, METRIC_L2, 7); + EXPECT_FALSE(index.is_trained()); + + try { + index.add(points.data(), kCount, nullptr, false, 1); + FAIL() << "adding to an untrained index succeeded"; + } catch (const std::logic_error& error) { + EXPECT_STREQ( + error.what(), "IVF index must be trained, constructed or loaded before add" + ); + } + PID id = 0; + try { + index.search(points.data(), 1, 1, &id); + FAIL() << "searching an untrained index succeeded"; + } catch (const std::logic_error& error) { + EXPECT_STREQ( + error.what(), "IVF index must be trained, constructed or loaded before search" + ); + } + const rabitq_test::TempFile path_file("rabitq_ivf_untrained.index"); + EXPECT_THROW(index.save(path_file.path().c_str()), std::logic_error); + EXPECT_FALSE(std::filesystem::exists(std::filesystem::u8path(path_file.path()))); + + EXPECT_FALSE(index.is_trained()); + EXPECT_EQ(index.max_elements(), 0U); + EXPECT_EQ(index.num_stored_points(), 0U); + // The rejected calls leave a usable index. + const auto centroids = random_rows(kDynNumClusters, 32); + index.set_centroids(centroids.data(), 1); + index.add(points.data(), kCount, nullptr, false, 1); + EXPECT_EQ(index.num_points(), kCount); +} + +TEST(IvfTrainTest, ConstructIsOnlyForPointCountIndexesUntilPointsAreAdded) { + constexpr size_t kCount = 64; + const auto points = random_rows(kCount, 33); + const auto centroids = random_rows(kDynNumClusters, 34); + const auto assigned = nearest_centroids(centroids, points.data(), kCount, METRIC_L2); + const auto cp = default_training_parameters(METRIC_L2, 7); + + IVF counted(kCount, kDynDim, kDynNumClusters, 4, METRIC_L2); + try { + counted.set_centroids(centroids.data(), 1); + FAIL() << "setting centroids on a point-count index succeeded"; + } catch (const std::logic_error& error) { + EXPECT_STREQ( + error.what(), "IVF index created with a point count must be built by construct" + ); + } + EXPECT_THROW(counted.train(kCount, points.data(), cp), std::logic_error); + EXPECT_FALSE(counted.is_trained()); + EXPECT_EQ(counted.max_elements(), kCount); + counted.construct(points.data(), centroids.data(), assigned.data(), false, 1); + EXPECT_EQ(counted.num_points(), kCount); + + IVF uncounted(kDynDim, kDynNumClusters, 4, METRIC_L2, 7); + uncounted.set_centroids(centroids.data(), 1); + try { + uncounted.construct(points.data(), centroids.data(), assigned.data(), false, 1); + FAIL() << "constructing an index without a point count succeeded"; + } catch (const std::logic_error& error) { + EXPECT_STREQ(error.what(), "IVF must be configured before construction"); + } + EXPECT_TRUE(uncounted.is_trained()); + EXPECT_EQ(uncounted.num_stored_points(), 0U); +} + +TEST(IvfTrainTest, SetCentroidsAndAddMatchConstructByteForByte) { + constexpr size_t kCount = 200; + constexpr uint32_t kSeed = 17; + const auto points = random_rows(kCount, 41); + const auto centroids = random_rows(kDynNumClusters, 42); + const rabitq_test::TempFile path_file("rabitq_ivf_set_centroids.index"); + const std::string& path = path_file.path(); + + for (auto metric : {METRIC_L2, METRIC_IP}) { + const auto assigned = nearest_centroids(centroids, points.data(), kCount, metric); + for (size_t bits : {1UL, 4UL, 9UL, 32UL}) { + SCOPED_TRACE( + ::testing::Message() + << "metric=" << (metric == METRIC_L2 ? "l2" : "ip") << " bits=" << bits + ); + IVF constructed(kCount, kDynDim, kDynNumClusters, bits, metric); + // Give the point-count constructor the rotation the seed selects. + IvfTestAccess::rotator(constructed) + .reset( + choose_rotator( + kDynDim, RotatorType::FhtKacRotator, constructed.padded_dim(), kSeed + ) + ); + constructed.construct( + points.data(), centroids.data(), assigned.data(), false, 1 + ); + const std::string expected = saved_bytes(constructed, path); + + IVF explicit_ids(kDynDim, kDynNumClusters, bits, metric, kSeed); + explicit_ids.set_centroids(centroids.data(), 1); + EXPECT_TRUE(explicit_ids.is_trained()); + EXPECT_EQ(explicit_ids.num_points(), 0U); + explicit_ids.add(points.data(), kCount, assigned.data(), false, 1); + EXPECT_TRUE(saved_bytes(explicit_ids, path) == expected) << "explicit IDs"; + + IVF routed(kDynDim, kDynNumClusters, bits, metric, kSeed); + routed.set_centroids(centroids.data(), 1); + routed.add(points.data(), kCount, nullptr, false, 1); + EXPECT_TRUE(saved_bytes(routed, path) == expected) << "routed"; + } + } +} + +TEST(IvfTrainTest, ReplacesCentroidsOnlyWhileEmpty) { + constexpr size_t kCount = 120; + constexpr uint32_t kSeed = 9; + const auto points = random_rows(kCount, 51); + const auto first = random_rows(kDynNumClusters, 52); + const auto second = random_rows(kDynNumClusters, 53); + const auto cp = default_training_parameters(METRIC_L2, kSeed); + const rabitq_test::TempFile path_file("rabitq_ivf_retrain.index"); + const std::string& path = path_file.path(); + + IVF index(kDynDim, kDynNumClusters, 4, METRIC_L2, kSeed); + index.set_centroids(first.data(), 1); + index.train(kCount, points.data(), cp); + index.set_centroids(second.data(), 1); + IVF reference(kDynDim, kDynNumClusters, 4, METRIC_L2, kSeed); + reference.set_centroids(second.data(), 1); + index.add(points.data(), kCount, nullptr, false, 1); + reference.add(points.data(), kCount, nullptr, false, 1); + // Only the last centroids remain. + const std::string saved = saved_bytes(index, path); + EXPECT_TRUE(saved_bytes(reference, path) == saved); + + IVF loaded; + loaded.load(path.c_str()); + for (IVF* holder : {&index, &loaded}) { + try { + holder->train(kCount, points.data(), cp); + FAIL() << "training an index with points succeeded"; + } catch (const std::logic_error& error) { + EXPECT_STREQ(error.what(), "IVF cannot retrain an index with points"); + } + try { + holder->set_centroids(first.data(), 1); + FAIL() << "replacing the centroids of an index with points succeeded"; + } catch (const std::logic_error& error) { + EXPECT_STREQ( + error.what(), "IVF cannot replace the centroids of an index with points" + ); + } + EXPECT_TRUE(saved_bytes(*holder, path) == saved); + } +} + +TEST(IvfTrainTest, RejectsInvalidCentroidsAndTrainingDataWithoutChangingState) { + constexpr size_t kCount = 64; + constexpr uint32_t kSeed = 11; + const auto points = random_rows(kCount, 61); + const auto centroids = random_rows(kDynNumClusters, 62); + auto bad_centroids = centroids; + bad_centroids[3] = std::numeric_limits::quiet_NaN(); + auto bad_points = points; + bad_points[5] = std::numeric_limits::infinity(); + const auto cp = default_training_parameters(METRIC_L2, kSeed); + auto bad_cp = cp; + bad_cp.niter = 0; + + IVF index(kDynDim, kDynNumClusters, 4, METRIC_L2, kSeed); + const auto expect_rejected = [&] { + EXPECT_THROW(index.set_centroids(nullptr, 1), std::invalid_argument); + EXPECT_THROW(index.set_centroids(bad_centroids.data(), 1), std::invalid_argument); + EXPECT_THROW(index.train(kCount, nullptr, cp), std::invalid_argument); + try { + index.train(kDynNumClusters - 1, points.data(), cp); + FAIL() << "training with fewer points than clusters succeeded"; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ(error.what(), "IVF training needs at least num_clusters points"); + } + EXPECT_THROW(index.train(kCount, bad_points.data(), cp), std::invalid_argument); + EXPECT_THROW(index.train(kCount, points.data(), bad_cp), std::invalid_argument); + }; + expect_rejected(); + EXPECT_FALSE(index.is_trained()); + + index.set_centroids(centroids.data(), 1); + const auto before = all_centroids(index); + expect_rejected(); + EXPECT_TRUE(index.is_trained()); + EXPECT_EQ(all_centroids(index), before); + + IVF unconfigured; + EXPECT_THROW(unconfigured.set_centroids(centroids.data(), 1), std::logic_error); + EXPECT_THROW(unconfigured.train(kCount, points.data(), cp), std::logic_error); +} + +TEST(IvfTrainTest, TrainedIndexSearchesAndSavesOnlyAfterAdd) { + constexpr size_t kCount = 300; + constexpr uint32_t kSeed = 21; + const auto points = random_rows(kCount, 71); + const rabitq_test::TempFile path_file("rabitq_ivf_trained_empty.index"); + const std::string& path = path_file.path(); + + for (auto metric : {METRIC_L2, METRIC_IP}) { + SCOPED_TRACE(metric == METRIC_L2 ? "l2" : "ip"); + std::remove(path.c_str()); + IVF index(kDynDim, kDynNumClusters, 4, metric, kSeed); + index.train(kCount, points.data(), default_training_parameters(metric, kSeed)); + EXPECT_TRUE(index.is_trained()); + EXPECT_EQ(index.num_points(), 0U); + PID id = 0; + try { + index.search(points.data(), 1, 1, &id); + FAIL() << "searching a trained index without points succeeded"; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ( + error.what(), "IVF search k must be between 1 and the point count" + ); + } + try { + index.save(path.c_str()); + FAIL() << "saving a trained index without points succeeded"; + } catch (const std::logic_error& error) { + EXPECT_STREQ(error.what(), "IVF index has no points to save; add points first"); + } + EXPECT_FALSE(std::filesystem::exists(std::filesystem::u8path(path))); + + index.add(points.data(), kCount, nullptr, false, 1); + EXPECT_EQ(index.num_points(), kCount); + const auto hits = search_everything(index, points.data()); + if (metric == METRIC_L2) { + EXPECT_EQ(hits.ids[0], 0U); + } + index.save(path.c_str()); + IVF loaded; + loaded.load(path.c_str()); + EXPECT_TRUE(loaded.is_trained()); + const auto reloaded = search_everything(loaded, points.data()); + EXPECT_EQ(reloaded.ids, hits.ids); + EXPECT_EQ(reloaded.distances, hits.distances); + } +} + +TEST(IvfTrainTest, ChoosesKMeansByClusterCount) { + for (const size_t num_clusters : + {kMaxRaBitQKMeansClusters, kMaxRaBitQKMeansClusters + 1}) { + SCOPED_TRACE(num_clusters); + const size_t count = num_clusters + 64; + const auto points = random_rows(count, 81); + auto cp = default_training_parameters(METRIC_L2, 23); + cp.niter = 2; + cp.num_threads = 1; // QGKMeans graph building is deterministic on one thread. + + IVF index(kDynDim, num_clusters, 1, METRIC_L2, 23); + index.train(count, points.data(), cp); + + rabitqkmeans::RaBitQKMeans rabitq(kDynDim, num_clusters, cp); + rabitq.train(count, points.data()); + qgkmeans::QGKMeansParameters qg_cp; + clustering::detail::copy_shared_parameters(cp, qg_cp); + qgkmeans::QGKMeans qg(kDynDim, num_clusters, qg_cp); + qg.train(count, points.data()); + ASSERT_NE(rabitq.centroids, qg.centroids); + + IVF expected(kDynDim, num_clusters, 1, METRIC_L2, 23); + expected.set_centroids( + (num_clusters <= kMaxRaBitQKMeansClusters ? rabitq.centroids : qg.centroids) + .data(), + 1 + ); + EXPECT_EQ(all_centroids(index), all_centroids(expected)); + } +} + +TEST(IvfTrainTest, DefaultsUseSphericalKMeansForInnerProduct) { + constexpr size_t kCount = 400; + auto points = random_rows(kCount, 91); + for (auto& value : points) { + value *= 10.0F; // Means of these points are far from unit norm. + } + for (auto metric : {METRIC_L2, METRIC_IP}) { + SCOPED_TRACE(metric == METRIC_L2 ? "l2" : "ip"); + const auto cp = default_training_parameters(metric, 5); + EXPECT_EQ(cp.niter, 25U); + EXPECT_EQ(cp.max_points_per_centroid, 256U); + EXPECT_EQ(cp.seed, 5U); + EXPECT_EQ(cp.spherical, metric == METRIC_IP); + + IVF index(kDynDim, kDynNumClusters, 4, metric, 5); + index.train(kCount, points.data(), cp); + for (size_t i = 0; i < kDynNumClusters; ++i) { + // The rotation preserves norms. + const auto centroid = IvfTestAccess::centroid(index, static_cast(i)); + double norm_sqr = 0; + for (const float value : centroid) { + norm_sqr += static_cast(value) * value; + } + if (metric == METRIC_IP) { + EXPECT_NEAR(std::sqrt(norm_sqr), 1.0, 1e-3); + } else { + EXPECT_GT(std::sqrt(norm_sqr), 2.0); + } + } + } +} + } // namespace } // namespace rabitqlib::ivf