diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 6d97338d..19d260b6 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -206,6 +206,7 @@ jobs: simd_filter+=':RaBitQKMeansTest.FinalModesKeepFlatLloydAssignmentsAcrossClusterCounts' simd_filter+=':IvfPersistenceTest.ReloadsRawAndQuantizedStorage' simd_filter+=':HnswSaveTest.SuccessfulSavePreservesSearchAfterLoading' + simd_filter+=':IvfAddTest.*:IvfRemoveTest.*' for cpu_model in hsw icl; do RABITQ_TEST_CPU="$cpu_model" "$RUNNER_TEMP/sde/sde64" "-$cpu_model" -- \ build-install/tests/rabitq_tests --gtest_filter="$simd_filter" diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index f25c54e8..2543c4b0 100644 --- a/DEVELOPMENT.md +++ b/DEVELOPMENT.md @@ -354,6 +354,18 @@ compatibility test fixture or an explicit rejection path. IVF, HNSW, and Symphon formats and must each be reviewed. SymphonyQG includes a versioned quantized format plus a legacy raw-format fallback; preserve both unless a breaking change is explicitly requested. +IVF `remove` stores a removed point as `f_add = +inf` inside the existing batch data, so it needs no +discriminator. No index written before `remove` existed has an infinite `f_add`, so nothing in an old +file is reinterpreted, and a file that has removals loads in release 0.5.0 too, where those points +never appear in results. Keep this convention when touching IVF code: + +- Every estimator and scan path must turn `f_add = +inf` into an infinite estimated distance and an + infinite lower bound, and never into NaN. A NaN distance would enter the result buffer. +- `fastscan::unpack_codes` must stay the exact inverse of `pack_codes`; `IVF::add` uses it to refill + the partial last batch of a cluster. +- `IvfRemoveTest.ExcludesRemovedPointsOnEveryScanPath`, `IvfRemoveTest.SurvivesSaveLoadAndLaterAdds`, + and `FastScanPackingTest.UnpackInvertsPackIncludingTailsAndDoesNotOverrun` guard these rules. + ### Change Python bindings Shared NumPy and string conversion helpers live in `python_bindings/bindings_common.hpp`. Register diff --git a/docs/docs/index/ivf.md b/docs/docs/index/ivf.md index 5ef43825..320ec48e 100644 --- a/docs/docs/index/ivf.md +++ b/docs/docs/index/ivf.md @@ -156,3 +156,85 @@ remaining quantized bits or the stored raw vectors. Raw reranking computes squared L2 or `1 - dot(query, vector)` in the original coordinates; the search API is unchanged. Cluster selection and filtering remain approximate in both modes. Search returns the top `k` results after scanning the selected clusters. + +## Updating an Index +A constructed or loaded index can gain and lose points without the original +data. The index file format does not change. + +```c++ +void IVF::add( + const float* data, + size_t n, + const PID* cluster_ids = nullptr, + bool faster = false, + size_t num_threads = std::numeric_limits::max() +); + +size_t IVF::remove(const PID* ids_to_remove, size_t n); +``` + +- **data**, **n**: `n` new vectors, quantized like `construct` does. Vector `i` + receives the PID `max_elements() + i`, so PIDs stay dense. +- **cluster_ids**: The cluster of each new vector, in `[0, cluster_num)`. When + it is `nullptr`, each vector goes to its nearest centroid, chosen the same way + a query is routed. +- **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 + repeated safely. + +In Python: + +```python +new_ids = index.add(vectors) # route to the nearest centroids +new_ids = index.add(vectors, cluster_ids=labels) # or choose the clusters +removed = index.remove(new_ids[:10]) +``` + +In C++, `construct` expects `max_elements()` rows once points have been added. It is +not safe to call `add` or `remove` while another thread searches the same index. + +### Cost of `add` + +`add` builds the grown index next to the old one and swaps it in, so the index is +unchanged if it throws. Every call copies the whole storage, so: + +- It needs memory for both copies while it runs. +- Each call has a fixed cost that grows with the size of the whole index, on top of a + cost per new vector. **Add many vectors per call rather than one at a time.** + +On an index of 1,000,000 vectors with 128 dimensions, 1,024 clusters and 1 bit +(8 threads, portable build, Intel i7-12700KF), adding 1,000 vectors in one call took +about 5 ms, and adding the same vectors one call each took about 3 s (2.8 s and 3.7 s +in two runs). For that 33 MB file, a process that had just loaded it held about +65 MB, and peaked at about 100 MB during one `add` of 1,000 vectors. + +There is no spare capacity: every `add` moves the data. + +### Recall after `add` + +The centroids and the rotation are fixed when the index is constructed, so +`add` never retrains them. If the added vectors come from a different distribution, +or the index grows many times beyond its original size, recall at a given +`nprobe` can drop. Rebuild from the original data with new centroids when +recall matters more than the cost of a rebuild. + +### How removal is stored + +`remove` hides points from search but keeps their storage. Removed points still +count in `max_elements()` and cannot be restored. + +A removed point has its `f_add` value set to `+inf`. That value is the constant +term of the distance estimate, so the estimated distance and the lower bound of the +point are `+inf` for every bit width, metric, and SIMD backend. The point never +enters the result buffer and never triggers reranking. Any code that estimates +IVF distances must keep this behavior and must never turn `+inf` into NaN. + +Search may return fewer than `k` results once points are removed: the missing slots +hold `kPidMax` (`2**32 - 1` in Python) with an infinite distance, as when the probed +clusters hold fewer than `k` points. + +The value lives in the ordinary batch data, so the file format does not change and +removal survives `save` and `load`. No file written before `remove` existed has an +infinite `f_add`, so nothing in an old file is reinterpreted. A file that has +removals also loads in release 0.5.0, where the removed points never appear in results. diff --git a/docs/docs/quick_start.md b/docs/docs/quick_start.md index 2e114a62..54320555 100644 --- a/docs/docs/quick_start.md +++ b/docs/docs/quick_start.md @@ -97,7 +97,16 @@ ids, distances = index.search(queries, k=10, nprobe=5, high_accuracy=False) # Omit high_accuracy, or pass None, to use automatic selection. ``` -See the [IVF guide](index/ivf.md) for storage costs and persistence compatibility. +An IVF index can also grow and shrink after it is built, without the original data: + +```python +new_vectors = rng.standard_normal((50, 64)).astype(np.float32) +new_ids = index.add(new_vectors) # nearest-centroid routing; ids 500..549 +index.remove(new_ids[:10]) # hide ten of them from later searches +``` + +See the [IVF guide](index/ivf.md) for storage costs, persistence compatibility, and +the limits of updating an index. The `metric` argument accepts `"l2"` and `"ip"` (also spelled `"innerproduct"`). To search by cosine similarity, normalize database and diff --git a/include/rabitqlib/fastscan/fastscan.hpp b/include/rabitqlib/fastscan/fastscan.hpp index 78d1253b..ed5757e8 100644 --- a/include/rabitqlib/fastscan/fastscan.hpp +++ b/include/rabitqlib/fastscan/fastscan.hpp @@ -117,6 +117,43 @@ inline void pack_codes( } } +/** + * @brief Inverse of pack_codes: recover the compact quantization codes of the first `num` + * vectors from their packed blocks. Padding lanes of a partial final batch are ignored. + * + * @param padded_dim dimension of quantized data (i.e., quantization code) + * @param blocks packed quantization code, as produced by pack_codes + * @param num number of quantization code to recover + * @param quantization_code recovered codes, num * padded_dim / 8 bytes + */ +inline void unpack_codes( + size_t padded_dim, const uint8_t* blocks, size_t num, uint8_t* quantization_code +) { + const size_t cols = padded_dim / 8; + + for (size_t row = 0; row < num; row += kBatchSize) { + for (size_t i = 0; i < cols; ++i) { + for (size_t j = 0; j < 16; ++j) { + // blocks[j] holds the high nibbles of vector kPerm0[j] and the low nibbles + // of vector kPerm0[j] + 16 packed together; blocks[j + 16] holds the rest. + const size_t first = row + kPerm0[j]; + const auto a = static_cast( + ((blocks[j] & 0x0F) << 4) | (blocks[j + 16] & 0x0F) + ); + const auto b = + static_cast((blocks[j] & 0xF0) | (blocks[j + 16] >> 4)); + if (first < num) { + quantization_code[first * cols + i] = a; + } + if (first + 16 < num) { + quantization_code[(first + 16) * cols + i] = b; + } + } + blocks += 32; + } + } +} + // Accumulate one batch into int32_t results. Throws if a result exceeds int32_t. // dim must be a positive multiple of 16. // Requires AVX2/FMA, AVX-512, or ARM NEON; otherwise throws std::runtime_error. diff --git a/include/rabitqlib/index/ivf/ivf.hpp b/include/rabitqlib/index/ivf/ivf.hpp index a2514f9e..157b86d1 100644 --- a/include/rabitqlib/index/ivf/ivf.hpp +++ b/include/rabitqlib/index/ivf/ivf.hpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -66,6 +67,20 @@ class IVF { void quantize_cluster(Cluster&, const std::vector&, const float*, const float*, float*, const quant::RabitqConfig&); + void + extend_cluster(const Cluster&, const std::vector&, const float*, const float*, PID, char*, char*, PID*, const quant::RabitqConfig&) + const; + + // True if `n` points cannot be addressed by the storage size computations + [[nodiscard]] bool exceeds_storage(size_t n) const { + const size_t max_size = std::numeric_limits::max(); + const size_t batch_bytes = BatchDataMap::data_bytes(padded_dim_); + const size_t rerank_bytes = rerank_vector_bytes(); + return n > max_size / sizeof(PID) || n > max_size / sizeof(float) / padded_dim_ || + n > max_size / batch_bytes || + (rerank_bytes != 0 && n > max_size / rerank_bytes); + } + [[nodiscard]] size_t ids_bytes() const { return sizeof(PID) * num_; } // get num of bytes used for 1-bit code and corresponding factors @@ -161,6 +176,7 @@ class IVF { ~IVF(); + // Number of stored points, including any removed with `remove`. `add` increases it. [[nodiscard]] size_t max_elements() const { return num_; } [[nodiscard]] size_t dimension() const { return dim_; } [[nodiscard]] size_t nbits() const { return raw_reranking_ ? 32 : ex_bits_ + 1; } @@ -169,6 +185,48 @@ class IVF { void construct(const float*, const float*, const PID*, bool, size_t); + /** + * @brief Append points to a constructed or loaded index without the original data. + * + * 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 + * data the centroids were trained on. Point `i` of `data` receives PID + * `max_elements() + i`. The index is unchanged if this throws. + * + * The grown index is built next to the old one, so a call needs memory for both and + * takes time proportional to the whole index, not just to `n`: prefer a few large + * calls to many small ones. `construct` expects `max_elements()` rows afterwards. + * Not safe to call while another thread searches, adds or removes. + * + * @param data New points (n * dim) + * @param n Number of new points + * @param cluster_ids Cluster of each new point, or nullptr to route each point to its + * nearest centroid the same way a query is routed + * @param faster Same meaning as in construct + * @param num_threads Threads used for routing and quantization + */ + void add( + const float* data, + size_t n, + const PID* cluster_ids = nullptr, + bool faster = false, + size_t num_threads = std::numeric_limits::max() + ); + + /** + * @brief Exclude points from all later search results. + * + * A removed point keeps its storage, still counts in `max_elements()`, and search may + * return fewer than k results (unfilled slots use kPidMax). Removal is stored in the + * index, so it survives save and load, and is idempotent. Removed points cannot be + * restored. Not safe to call while another thread searches or adds. + * + * @param ids_to_remove PIDs to remove + * @param n Number of PIDs + * @return Number of points newly removed + */ + size_t remove(const PID* ids_to_remove, size_t n); + void save(const char*) const; void load(const char*); @@ -216,11 +274,7 @@ inline IVF::IVF( } padded_dim_ = round_up_to_multiple(dim_, 64); const size_t max_size = std::numeric_limits::max(); - const size_t batch_bytes = BatchDataMap::data_bytes(padded_dim_); - const size_t rerank_bytes = rerank_vector_bytes(); - if (n > max_size / sizeof(PID) || n > max_size / sizeof(float) / padded_dim_ || - n > max_size / batch_bytes || (rerank_bytes != 0 && n > max_size / rerank_bytes) || - cluster_num > max_size / sizeof(float) / padded_dim_ || + if (exceeds_storage(n) || cluster_num > max_size / sizeof(float) / padded_dim_ || (type == RotatorType::MatrixRotator && dim > max_size / sizeof(float) / padded_dim_ )) { throw std::invalid_argument("IVF configuration exceeds addressable storage"); @@ -402,6 +456,277 @@ inline void IVF::quantize_cluster( } } +inline void IVF::add( + const float* data, size_t n, const PID* cluster_ids, bool faster, size_t num_threads +) { + if (!ready_ || initer_ == nullptr || rotator_ == nullptr || + cluster_lst_.size() != num_cluster_ || batch_storage_.empty() || + id_storage_.size() != num_) { + throw std::logic_error("IVF index must be constructed or loaded before add"); + } + if (n == 0) { + return; + } + if (data == nullptr) { + throw std::invalid_argument("IVF add data must not be null"); + } + if (n > buffer::kSearchBufferMaxPointCount - num_) { + throw std::invalid_argument("IVF point count exceeds the supported ID range"); + } + const size_t new_num = num_ + n; + if (exceeds_storage(new_num)) { + throw std::invalid_argument("IVF configuration exceeds addressable storage"); + } + + /* Assign every new point to a cluster */ + std::vector assigned(n); + if (cluster_ids != nullptr) { + for (size_t i = 0; i < n; ++i) { + if (cluster_ids[i] >= num_cluster_) { + throw std::invalid_argument("Cluster ID is out of range"); + } + assigned[i] = cluster_ids[i]; + } + } else { + parallel_for(0, n, num_threads, [&](size_t row, size_t /*thread_id*/) { + std::vector rotated(padded_dim_); + std::vector> nearest(1); + rotator_->rotate(data + (row * dim_), rotated.data()); + initer_->centroids_distances(rotated.data(), 1, nearest); + assigned[row] = nearest[0].id; + }); + } + + std::vector> rows_of(num_cluster_); + for (size_t i = 0; i < n; ++i) { + rows_of[assigned[i]].push_back(static_cast(i)); + } + std::vector new_sizes(num_cluster_); + for (size_t i = 0; i < num_cluster_; ++i) { + new_sizes[i] = cluster_lst_[i].num() + rows_of[i].size(); + } + + /* Build the grown storage on the side so the index is unchanged on failure */ + ByteStorage new_batch(batch_data_bytes(new_sizes)); + ByteStorage new_ex; + FloatStorage new_raw; + const size_t rerank_bytes = rerank_vector_bytes(); + if (raw_reranking_) { + new_raw = FloatStorage(new_num * dim_); + } else if (rerank_bytes > 0) { + new_ex = ByteStorage(new_num * rerank_bytes); + } + IdStorage new_ids(new_num); + char* new_rerank = raw_reranking_ ? reinterpret_cast(new_raw.data()) + : reinterpret_cast(new_ex.data()); + + std::vector batch_offset(num_cluster_); + std::vector vector_offset(num_cluster_); + size_t added_batches = 0; + size_t added_vectors = 0; + for (size_t i = 0; i < num_cluster_; ++i) { + batch_offset[i] = added_batches; + vector_offset[i] = added_vectors; + added_batches += div_round_up(new_sizes[i], fastscan::kBatchSize); + added_vectors += new_sizes[i]; + } + + quant::RabitqConfig config; + if (faster) { + config = quant::faster_config(padded_dim_, ex_bits_ + 1); + } + + const size_t batch_bytes = BatchDataMap::data_bytes(padded_dim_); + const PID first_pid = static_cast(num_); + // An exception must not leave an OpenMP region, so keep the first one and rethrow it. + std::exception_ptr failure; +#pragma omp parallel for schedule(dynamic) num_threads(resolve_num_threads(num_threads)) + for (std::ptrdiff_t index = 0; index < static_cast(num_cluster_); + ++index) { + const size_t i = static_cast(index); + try { + extend_cluster( + cluster_lst_[i], + rows_of[i], + data, + initer_->centroid(static_cast(i)), + first_pid, + reinterpret_cast(new_batch.data()) + (batch_offset[i] * batch_bytes), + rerank_bytes > 0 ? new_rerank + (vector_offset[i] * rerank_bytes) : nullptr, + new_ids.data() + vector_offset[i], + config + ); + } catch (...) { +#pragma omp critical(rabitq_ivf_add_failure) + if (!failure) { + failure = std::current_exception(); + } + } + } + if (failure) { + std::rethrow_exception(failure); + } + + /* Commit. Rebuilding the cluster list reuses its capacity, so it cannot throw. */ + batch_storage_.swap(new_batch); + ex_storage_.swap(new_ex); + raw_storage_.swap(new_raw); + id_storage_.swap(new_ids); + num_ = new_num; + cluster_lst_.clear(); + init_clusters(new_sizes); +} + +/** + * @brief Copy one cluster into its grown storage and append new points to it. + * + * Full batches, rerank data and IDs are copied unchanged. The old partial tail batch is + * unpacked and repacked together with the new points, so the result is byte-identical to + * quantizing the same points in the same order in one construct call. + */ +inline void IVF::extend_cluster( + const Cluster& old_cluster, + const std::vector& rows, + const float* data, + const float* rotated_centroid, + PID first_pid, + char* batch_dst, + char* rerank_dst, + PID* ids_dst, + const quant::RabitqConfig& config +) const { + const size_t old_num = old_cluster.num(); + const size_t batch_bytes = BatchDataMap::data_bytes(padded_dim_); + const size_t rerank_bytes = rerank_vector_bytes(); + const size_t cols = padded_dim_ / 8; + + if (rows.empty()) { + const size_t old_batches = div_round_up(old_num, fastscan::kBatchSize); + std::copy_n(old_cluster.batch_data(), old_batches * batch_bytes, batch_dst); + std::copy_n(old_cluster.ex_data(), old_num * rerank_bytes, rerank_dst); + std::copy_n(old_cluster.ids(), old_num, ids_dst); + return; + } + + const size_t full_batches = old_num / fastscan::kBatchSize; + const size_t kept = old_num - (full_batches * fastscan::kBatchSize); + const size_t tail_num = kept + rows.size(); + + std::copy_n(old_cluster.batch_data(), full_batches * batch_bytes, batch_dst); + std::copy_n(old_cluster.ex_data(), old_num * rerank_bytes, rerank_dst); + std::copy_n(old_cluster.ids(), old_num, ids_dst); + + /* Codes and factors of every point in the tail, old points first */ + std::vector codes(tail_num * cols); + std::vector f_add(tail_num); + std::vector f_rescale(tail_num); + std::vector f_error(tail_num); + if (kept > 0) { + ConstBatchDataMap old_tail( + old_cluster.batch_data() + (full_batches * batch_bytes), padded_dim_ + ); + fastscan::unpack_codes(padded_dim_, old_tail.bin_code(), kept, codes.data()); + old_tail.f_add().copy_to(f_add.data(), kept); + old_tail.f_rescale().copy_to(f_rescale.data(), kept); + old_tail.f_error().copy_to(f_error.data(), kept); + } + + std::vector rotated(padded_dim_); + for (size_t j = 0; j < rows.size(); ++j) { + const float* vector = data + (static_cast(rows[j]) * dim_); + const size_t lane = kept + j; + rotator_->rotate(vector, rotated.data()); + quant::quantize_compact_one_bit( + rotated.data(), + rotated_centroid, + padded_dim_, + codes.data() + (lane * cols), + f_add[lane], + f_rescale[lane], + f_error[lane], + metric_type_ + ); + char* rerank = rerank_dst + ((old_num + j) * rerank_bytes); + if (raw_reranking_) { + std::memcpy(rerank, vector, rerank_bytes); + } else if (ex_bits_ > 0) { + quant::quantize_compact_ex_bits( + rotated.data(), + rotated_centroid, + padded_dim_, + ex_bits_, + rerank, + metric_type_, + config + ); + } + ids_dst[old_num + j] = first_pid + rows[j]; + } + + /* Repack the tail into batches */ + char* batch = batch_dst + (full_batches * batch_bytes); + for (size_t start = 0; start < tail_num; start += fastscan::kBatchSize) { + const size_t count = std::min(fastscan::kBatchSize, tail_num - start); + BatchDataMap cur_batch(batch, padded_dim_); + fastscan::pack_codes( + padded_dim_, codes.data() + (start * cols), count, cur_batch.bin_code() + ); + for (size_t i = 0; i < count; ++i) { + cur_batch.f_add()[i] = f_add[start + i]; + cur_batch.f_rescale()[i] = f_rescale[start + i]; + cur_batch.f_error()[i] = f_error[start + i]; + } + batch += batch_bytes; + } +} + +inline size_t IVF::remove(const PID* ids_to_remove, size_t n) { + if (!ready_ || initer_ == nullptr || rotator_ == nullptr || + cluster_lst_.size() != num_cluster_ || batch_storage_.empty() || + id_storage_.size() != num_) { + throw std::logic_error("IVF index must be constructed or loaded before remove"); + } + if (n == 0) { + return 0; + } + if (ids_to_remove == nullptr) { + throw std::invalid_argument("IVF remove IDs must not be null"); + } + + // Validate every ID before changing anything + std::vector marked(num_, false); + for (size_t i = 0; i < n; ++i) { + if (ids_to_remove[i] >= num_) { + throw std::invalid_argument("IVF remove ID is out of range"); + } + marked[ids_to_remove[i]] = true; + } + + // A point whose f_add is +inf has +inf estimated and lower-bound distances in every + // scan path, so it never enters the result buffer or triggers reranking. + constexpr float kRemoved = std::numeric_limits::infinity(); + const size_t batch_bytes = BatchDataMap::data_bytes(padded_dim_); + size_t removed = 0; + for (const auto& cluster : cluster_lst_) { + const PID* cluster_ids = cluster.ids(); + for (size_t pos = 0; pos < cluster.num(); ++pos) { + if (!marked[cluster_ids[pos]]) { + continue; + } + BatchDataMap cur_batch( + cluster.batch_data() + ((pos / fastscan::kBatchSize) * batch_bytes), + padded_dim_ + ); + const size_t lane = pos % fastscan::kBatchSize; + if (static_cast(cur_batch.f_add()[lane]) != kRemoved) { + cur_batch.f_add()[lane] = kRemoved; + ++removed; + } + } + } + return removed; +} + inline void IVF::save(const char* filename) const { if (!ready_ || initer_ == nullptr || rotator_ == nullptr || cluster_lst_.size() != num_cluster_ || batch_storage_.empty() || diff --git a/python_bindings/bindings_common.hpp b/python_bindings/bindings_common.hpp index 6b8a5c61..50312123 100644 --- a/python_bindings/bindings_common.hpp +++ b/python_bindings/bindings_common.hpp @@ -68,4 +68,21 @@ inline py::array_t ensure_1d_array return array; } +// Like ensure_1d_array, but refuses input whose dtype is not an integer. forcecast +// would otherwise truncate floats and turn booleans or numeric strings into ids, and +// callers use these ids to remove points permanently. An empty sequence has no dtype to +// check. +inline py::array_t +ensure_1d_integer_array(py::handle value, const char* name) { + auto typed = py::array::ensure(value); + if (!typed) { + throw std::invalid_argument(std::string(name) + " must be a NumPy array"); + } + const char kind = typed.dtype().kind(); + if (typed.size() > 0 && kind != 'i' && kind != 'u') { + throw std::invalid_argument(std::string(name) + " must contain integers"); + } + return ensure_1d_array(value, name); +} + } // namespace rabitqlib::python_bindings \ No newline at end of file diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index 7442cb62..3d3a212d 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/python_bindings/ivf_bindings.cpp @@ -7,6 +7,7 @@ #include #include +#include #include #include #include @@ -91,6 +92,77 @@ class IvfIndex { built_ = true; } + py::array_t add( + py::handle data, + const py::object& cluster_ids, + size_t num_threads = 1, + bool fast_quantization = false + ) { + auto data_array = ensure_2d_array(data, "data"); + if (!built_) { + throw std::runtime_error("IvfIndex must be built or loaded before add"); + } + if (static_cast(data_array.shape(1)) != dim_) { + throw std::invalid_argument("data dimension does not match index dim"); + } + const auto rows = static_cast(data_array.shape(0)); + + // Validate cluster IDs as int64 first: casting straight to uint32 would wrap + // negative or oversized values onto valid clusters. + std::vector assigned; + if (!cluster_ids.is_none()) { + auto cluster_ids_array = ensure_1d_integer_array(cluster_ids, "cluster_ids"); + if (static_cast(cluster_ids_array.shape(0)) != rows) { + throw std::invalid_argument( + "cluster_ids length must match number of rows in data" + ); + } + assigned.reserve(rows); + for (py::ssize_t i = 0; i < cluster_ids_array.shape(0); ++i) { + const int64_t id = cluster_ids_array.data()[i]; + if (id < 0 || static_cast(id) >= num_clusters_) { + throw std::invalid_argument("cluster_ids must be in [0, num_clusters)"); + } + assigned.push_back(static_cast(id)); + } + } + + const size_t first = index_->max_elements(); + // Allocate the returned IDs before committing any changes to the index. + auto ids = py::array_t(static_cast(rows)); + auto* ids_data = ids.mutable_data(); + for (size_t i = 0; i < rows; ++i) { + ids_data[i] = static_cast(first + i); + } + + index_->add( + data_array.data(), + rows, + assigned.empty() ? nullptr : assigned.data(), + fast_quantization, + num_threads + ); + max_elements_ = index_->max_elements(); + return ids; + } + + size_t remove(py::handle ids) { + auto ids_array = ensure_1d_integer_array(ids, "ids"); + if (!built_) { + throw std::runtime_error("IvfIndex must be built or loaded before remove"); + } + std::vector to_remove; + to_remove.reserve(static_cast(ids_array.shape(0))); + for (py::ssize_t i = 0; i < ids_array.shape(0); ++i) { + const int64_t id = ids_array.data()[i]; + if (id < 0 || static_cast(id) >= max_elements_) { + throw std::invalid_argument("ids must be in [0, max_elements)"); + } + to_remove.push_back(static_cast(id)); + } + return index_->remove(to_remove.data(), to_remove.size()); + } + py::tuple search( py::handle queries, size_t k, @@ -210,6 +282,25 @@ void register_ivf(py::module_& m) { py::arg("num_threads") = 1, py::arg("fast_quantization") = false ) + .def( + "add", + &IvfIndex::add, + py::arg("data"), + py::arg("cluster_ids") = py::none(), + py::arg("num_threads") = 1, + py::arg("fast_quantization") = false, + "Append vectors without the original data and return their ids. Vectors " + "are quantized against the existing centroids, which do not move. Without " + "cluster_ids each vector goes to its nearest centroid. Each call copies " + "the whole index, so add many vectors per call rather than one at a time." + ) + .def( + "remove", + &IvfIndex::remove, + py::arg("ids"), + "Exclude ids from later search results and return how many were newly " + "removed. Removed vectors keep their storage and cannot be restored." + ) .def( "search", &IvfIndex::search, diff --git a/tests/python/test_ivf.py b/tests/python/test_ivf.py index 0f515470..0f59a2f5 100644 --- a/tests/python/test_ivf.py +++ b/tests/python/test_ivf.py @@ -437,3 +437,226 @@ def test_hnsw_inner_product_routing_roundtrip(tmp_path): ids, distances = loaded.search(query, k=1, nprobe=1, num_threads=4) assert np.all(ids == 0) np.testing.assert_allclose(distances, -99.0, rtol=0, atol=1e-5) + + +# ── add / remove ────────────────────────────────────────────────────────────── + + +def _dynamic_index(nbits=4, metric="l2", count=120, dim=65, clusters=3, seed=61): + """A small built index plus the points that were not built into it.""" + rng = np.random.default_rng(seed) + data = rng.standard_normal((count * 2, dim)).astype(np.float32) + centroids = rng.standard_normal((clusters, dim)).astype(np.float32) + if metric == "ip": + assigned = np.argmax(data @ centroids.T, axis=1) + else: + gap = data[:, None, :] - centroids[None, :, :] + assigned = np.argmin(np.einsum("ncd,ncd->nc", gap, gap), axis=1) + assigned = assigned.astype(np.uint32) + idx = IvfIndex(dim, count, clusters, nbits, metric) + idx.build(data[:count], centroids, assigned[:count]) + return idx, data[:count], data[count:], assigned[count:] + + +def test_add_returns_ids_and_makes_points_searchable(): + idx, old, new, _ = _dynamic_index(nbits=32) + added = idx.add(new) + assert added.dtype == np.uint32 + np.testing.assert_array_equal(added, np.arange(len(old), len(old) + len(new))) + assert idx.max_elements == len(old) + len(new) + # k may now exceed the original point count; raw storage makes the ranking exact. + ids, dists = idx.search(new[:3], k=len(old) + len(new), nprobe=3) + np.testing.assert_array_equal(ids[:, 0], added[:3]) + np.testing.assert_allclose(dists[:, 0], 0, atol=1e-4) + for row in ids: + np.testing.assert_array_equal(np.sort(row), np.arange(len(old) + len(new))) + + +@pytest.mark.parametrize("metric", ["l2", "ip"]) +@pytest.mark.parametrize("nbits", [1, 4, 32]) +def test_add_routing_matches_explicit_cluster_ids(tmp_path, metric, nbits): + idx, _, new, assigned = _dynamic_index(nbits=nbits, metric=metric) + base = tmp_path / "base.index" + idx.save(str(base)) + saved = {} + for name, kwargs in ( + ("routed", {}), + ("explicit", {"cluster_ids": assigned}), + ("threaded", {"num_threads": 4}), + ("int64", {"cluster_ids": assigned.astype(np.int64)}), + ): + # Reloading gives each copy the same rotator, so the bytes are comparable. + copy = IvfIndex.load(str(base)) + copy.add(new, **kwargs) + path = tmp_path / f"{name}.index" + copy.save(str(path)) + saved[name] = path.read_bytes() + assert len(set(saved.values())) == 1 + + +def test_add_chunks_match_a_single_add(tmp_path): + idx, _, new, assigned = _dynamic_index() + base = tmp_path / "base.index" + idx.save(str(base)) + whole = IvfIndex.load(str(base)) + whole.add(new, assigned) + parts = IvfIndex.load(str(base)) + for chunk in np.array_split(np.arange(len(new)), [1, 33, 34, 90]): + ids = parts.add(new[chunk], assigned[chunk]) + assert len(ids) == len(chunk) + a, b = tmp_path / "a.index", tmp_path / "b.index" + whole.save(str(a)) + parts.save(str(b)) + assert a.read_bytes() == b.read_bytes() + + +def test_add_accepts_fast_quantization_and_other_dtypes(): + idx, _, new, _ = _dynamic_index() + idx.add(new[:10].astype(np.float64), fast_quantization=True) + assert idx.max_elements == 130 + assert len(idx.add(new[:0])) == 0 + assert idx.max_elements == 130 + + +def test_add_rejects_invalid_arguments(): + idx, old, new, assigned = _dynamic_index() + with pytest.raises(ValueError, match="dimension"): + idx.add(new[:, :10]) + with pytest.raises(ValueError, match="2D"): + idx.add(new[0]) + with pytest.raises(ValueError, match="length"): + idx.add(new, assigned[:-1]) + for bad in (-1, 3, 2**32): # 2**32 would wrap to cluster 0 if cast to uint32 first + cluster_ids = assigned.astype(np.int64) + cluster_ids[0] = bad + with pytest.raises(ValueError, match="cluster_ids"): + idx.add(new, cluster_ids) + assert idx.max_elements == len(old) + + unbuilt = IvfIndex(65, 4, 1, 4) + with pytest.raises(RuntimeError, match="built or loaded"): + unbuilt.add(new[:2]) + + +@pytest.mark.parametrize("nbits", [1, 2, 3, 4, 5, 6, 7, 8, 9, 32]) +@pytest.mark.parametrize("metric", ["l2", "ip"]) +@pytest.mark.parametrize("high_accuracy", [False, True]) +def test_remove_excludes_points_and_survives_reload( + tmp_path, nbits, metric, high_accuracy +): + idx, old, new, _ = _dynamic_index(nbits=nbits, metric=metric) + count = len(old) + queries = new[:5] + before_ids, before_dists = idx.search( + queries, k=count, nprobe=3, high_accuracy=high_accuracy + ) + + removed = np.array([0, 5, 6, count - 1, 5]) + assert idx.remove(removed) == 4 + assert idx.remove(removed) == 0 + assert idx.remove([]) == 0 + assert idx.max_elements == count + + def check(index): + ids, dists = index.search( + queries, k=count, nprobe=3, high_accuracy=high_accuracy + ) + alive = count - 4 + for q in range(len(queries)): + keep = ~np.isin(before_ids[q], removed) + np.testing.assert_array_equal(ids[q, :alive], before_ids[q][keep]) + np.testing.assert_array_equal(dists[q, :alive], before_dists[q][keep]) + assert np.all(ids[q, alive:] == np.iinfo(np.uint32).max) + assert np.all(np.isinf(dists[q, alive:])) + top, top_dists = index.search( + queries, k=10, nprobe=3, high_accuracy=high_accuracy + ) + assert not np.isin(top, removed).any() + assert np.isfinite(top_dists).all() + + check(idx) + path = tmp_path / "removed.index" + idx.save(str(path)) + loaded = IvfIndex.load(str(path)) + check(loaded) + assert loaded.remove(removed) == 0 + + # Adding repacks partial batches that contain removed points; they must stay removed. + loaded.add(new) + idx.add(new) + ids, dists = loaded.search( + queries, k=count + len(new), nprobe=3, high_accuracy=high_accuracy + ) + assert not np.isin(ids, removed).any() + assert np.count_nonzero(ids[0] != np.iinfo(np.uint32).max) == count + len(new) - 4 + + # A reloaded index grows exactly like the original. + expected_ids, expected_dists = idx.search( + queries, k=count + len(new), nprobe=3, high_accuracy=high_accuracy + ) + np.testing.assert_array_equal(ids, expected_ids) + np.testing.assert_array_equal(dists, expected_dists) + + +def test_remove_rejects_invalid_ids_without_removing_anything(): + idx, old, new, _ = _dynamic_index() + count = len(old) + before = idx.search(new[:3], k=count, nprobe=3) + # 2**32 + 1 would wrap to point 1 if it were cast to uint32 first. + for bad in (-1, count, 2**32 + 1): + with pytest.raises(ValueError, match="ids"): + idx.remove([0, 1, bad]) + with pytest.raises(ValueError, match="1D"): + idx.remove(np.zeros((2, 2), dtype=np.int64)) + after = idx.search(new[:3], k=count, nprobe=3) + np.testing.assert_array_equal(before[0], after[0]) + np.testing.assert_array_equal(before[1], after[1]) + + unbuilt = IvfIndex(65, 4, 1, 4) + with pytest.raises(RuntimeError, match="built or loaded"): + unbuilt.remove([0]) + + +def test_removing_every_point_leaves_a_searchable_index(): + idx, old, new, _ = _dynamic_index() + assert idx.remove(np.arange(len(old))) == len(old) + ids, dists = idx.search(new[:2], k=5, nprobe=3) + assert np.all(ids == np.iinfo(np.uint32).max) + assert np.all(np.isinf(dists)) + + +@pytest.mark.parametrize( + "bad", + [ + [2.7], + [-0.5], + np.array( + [1.0, 2.0] + ), # integral floats are refused too: the dtype is the contract + np.array([True, False]), # a boolean mask is not a list of ids + ["3"], + np.array([1, 2], dtype=object), + ], +) +def test_remove_and_add_refuse_non_integer_ids(bad): + idx, old, new, assigned = _dynamic_index() + before = idx.search(new[:3], k=len(old), nprobe=3) + with pytest.raises(ValueError, match="integers"): + idx.remove(bad) + if len(bad) == 2: + with pytest.raises(ValueError, match="integers"): + idx.add(new[:2], cluster_ids=bad) + after = idx.search(new[:3], k=len(old), nprobe=3) + np.testing.assert_array_equal(before[0], after[0]) + assert idx.max_elements == len(old) + + +def test_remove_and_add_accept_any_integer_dtype_and_empty_input(): + idx, old, new, assigned = _dynamic_index() + assert idx.remove([]) == 0 + assert idx.remove(np.array([], dtype=np.float64)) == 0 # no dtype to check + assert idx.remove([1, 2]) == 2 + assert idx.remove(np.array([3, 4], dtype=np.uint8)) == 2 + assert idx.remove(np.array([5], dtype=np.int16)) == 1 + assert len(idx.add(new[:3], cluster_ids=[0, 1, 2])) == 3 + assert len(idx.add(new[:3], cluster_ids=np.array([0, 1, 2], dtype=np.uint16))) == 3 diff --git a/tests/unit/rabitqlib/fastscan/fastscan_test.cpp b/tests/unit/rabitqlib/fastscan/fastscan_test.cpp index 1375c8c6..1dcc705e 100644 --- a/tests/unit/rabitqlib/fastscan/fastscan_test.cpp +++ b/tests/unit/rabitqlib/fastscan/fastscan_test.cpp @@ -130,6 +130,32 @@ TEST(FastScanPackingTest, MatchesBitReferenceIncludingTailsAndUnalignedBuffers) } } +TEST(FastScanPackingTest, UnpackInvertsPackIncludingTailsAndDoesNotOverrun) { + for (size_t dim : {8U, 16U, 24U, 56U, 64U, 72U, 128U, 768U}) { + for (size_t num = 0; num <= 65; ++num) { + SCOPED_TRACE(::testing::Message() << "dim=" << dim << " num=" << num); + const size_t cols = dim / 8; + std::vector input(num * cols); + uint32_t state = 7; + for (auto& value : input) { + state = state * 1664525U + 1013904223U; + value = static_cast(state >> 24); + } + std::vector packed(((num + 31) / 32) * 32 * cols, 0xFF); + pack_codes(dim, input.data(), num, packed.data()); + const auto packed_original = packed; + + // Guard bytes on both sides detect writes outside the first `num` codes. + std::vector output(num * cols + 2, 0xA5); + unpack_codes(dim, packed.data(), num, output.data() + 1); + EXPECT_EQ(output.front(), 0xA5); + EXPECT_EQ(output.back(), 0xA5); + EXPECT_TRUE(std::equal(input.begin(), input.end(), output.begin() + 1)); + EXPECT_EQ(packed, packed_original); + } + } +} + constexpr size_t kMaxDim = 768; TEST(FastScanPackingTest, AccumulatesReferenceLutValuesOnEverySupportedBackend) { diff --git a/tests/unit/rabitqlib/index/ivf_test.cpp b/tests/unit/rabitqlib/index/ivf_test.cpp index 0b9d2f23..07e4d04c 100644 --- a/tests/unit/rabitqlib/index/ivf_test.cpp +++ b/tests/unit/rabitqlib/index/ivf_test.cpp @@ -6,11 +6,14 @@ #include #include #include +#include #include #include #include #include #include +#include +#include #include #include #include @@ -393,5 +396,500 @@ TEST(IvfPersistenceTest, RejectsPointIdsUsingTheSearchBufferMarker) { std::remove(path.c_str()); } +// ── add / remove ──────────────────────────────────────────────────────────────────────── + +constexpr size_t kDynDim = 65; // pads to 128, so padded and original dimensions differ +constexpr size_t kDynClusters = 4; + +std::vector random_rows(size_t rows, uint32_t seed) { + std::mt19937 gen(seed); + std::normal_distribution normal; + std::vector values(rows * kDynDim); + for (auto& value : values) { + value = normal(gen); + } + return values; +} + +// Nearest centroid in the original space: smallest L2 distance, or largest inner product. +std::vector nearest_centroids( + const std::vector& centroids, const float* rows, size_t count, MetricType metric +) { + std::vector assigned(count); + for (size_t i = 0; i < count; ++i) { + double best = std::numeric_limits::infinity(); + for (size_t c = 0; c < kDynClusters; ++c) { + double dist = 0; + for (size_t d = 0; d < kDynDim; ++d) { + const double x = rows[i * kDynDim + d]; + const double y = centroids[c * kDynDim + d]; + dist += metric == METRIC_L2 ? (x - y) * (x - y) : -x * y; + } + if (dist < best) { + best = dist; + assigned[i] = static_cast(c); + } + } + } + return assigned; +} + +std::string read_file(const std::string& path) { + std::ifstream input(path, std::ios::binary); + return {std::istreambuf_iterator(input), std::istreambuf_iterator()}; +} + +std::string saved_bytes(const IVF& index, const std::string& path) { + index.save(path.c_str()); + return read_file(path); +} + +// Exact distance as the index defines it: squared L2, or 1 - inner product. +double exact_distance(const float* a, const float* b, MetricType metric) { + double dist = metric == METRIC_L2 ? 0.0 : 1.0; + for (size_t d = 0; d < kDynDim; ++d) { + const double x = a[d]; + const double y = b[d]; + dist += metric == METRIC_L2 ? (x - y) * (x - y) : -x * y; + } + return dist; +} + +struct Hits { + std::vector ids; + std::vector distances; +}; + +// Ranks every point: k is the point count and every cluster is probed. +Hits search_everything( + const IVF& index, const float* query, std::optional hacc = std::nullopt +) { + const size_t k = index.max_elements(); + Hits hits{std::vector(k), std::vector(k)}; + if (hacc.has_value()) { + index.search(query, k, kDynClusters, hits.ids.data(), hits.distances.data(), *hacc); + } else { + index.search(query, k, kDynClusters, hits.ids.data(), hits.distances.data()); + } + return hits; +} + +TEST(IvfAddTest, EveryWayOfAddingTheSamePointsGivesByteIdenticalIndexes) { + constexpr size_t kOld = 200; + constexpr size_t kNew = 200; + const auto points = random_rows(kOld + kNew, 11); + const auto centroids = random_rows(kDynClusters, 12); + const float* extra = points.data() + (kOld * kDynDim); + const std::string base_path = ::testing::TempDir() + "rabitq_ivf_add_base.index"; + const std::string path = ::testing::TempDir() + "rabitq_ivf_add.index"; + + for (auto metric : {METRIC_L2, METRIC_IP}) { + const auto assigned = + nearest_centroids(centroids, points.data(), kOld + kNew, metric); + for (size_t bits : {1UL, 4UL, 9UL, 32UL}) { + SCOPED_TRACE( + ::testing::Message() + << "metric=" << (metric == METRIC_L2 ? "l2" : "ip") << " bits=" << bits + ); + IVF base(kOld, kDynDim, kDynClusters, bits, metric); + base.construct(points.data(), centroids.data(), assigned.data(), false, 1); + base.save(base_path.c_str()); + + // Every index below starts from the same file, so they share one rotator. + IVF single; + single.load(base_path.c_str()); + single.add(extra, kNew, assigned.data() + kOld, false, 1); + ASSERT_EQ(single.max_elements(), kOld + kNew); + const std::string expected = saved_bytes(single, path); + + IVF routed; + routed.load(base_path.c_str()); + routed.add(extra, kNew, nullptr, false, 1); + EXPECT_TRUE(saved_bytes(routed, path) == expected) << "routed"; + + // Chunk sizes put the tail batch of each cluster on and off a batch boundary. + IVF chunked; + chunked.load(base_path.c_str()); + size_t done = 0; + for (size_t chunk : std::array{1, 31, 32, 7, 1, kNew - 72}) { + chunked.add( + extra + (done * kDynDim), chunk, assigned.data() + kOld + done, false, 1 + ); + done += chunk; + } + ASSERT_EQ(done, kNew); + EXPECT_TRUE(saved_bytes(chunked, path) == expected) << "chunked"; + + IVF threaded; + threaded.load(base_path.c_str()); + threaded.add(extra, kNew, nullptr, false, 4); + EXPECT_TRUE(saved_bytes(threaded, path) == expected) << "threaded"; + } + } + std::remove(base_path.c_str()); + std::remove(path.c_str()); +} + +TEST(IvfAddTest, KeepsExistingDistancesAndQuantizesNewPointsAccurately) { + constexpr size_t kOld = 150; + constexpr size_t kNew = 150; + constexpr size_t kQueries = 8; + const auto points = random_rows(kOld + kNew, 21); + const auto centroids = random_rows(kDynClusters, 22); + const auto queries = random_rows(kQueries, 23); + + for (auto metric : {METRIC_L2, METRIC_IP}) { + const auto assigned = + nearest_centroids(centroids, points.data(), kOld + kNew, metric); + for (size_t bits : {1UL, 2UL, 4UL, 9UL, 32UL}) { + SCOPED_TRACE( + ::testing::Message() + << "metric=" << (metric == METRIC_L2 ? "l2" : "ip") << " bits=" << bits + ); + IVF index(kOld, kDynDim, kDynClusters, bits, metric); + index.construct(points.data(), centroids.data(), assigned.data(), false, 1); + std::vector before; + for (size_t q = 0; q < kQueries; ++q) { + before.push_back(search_everything(index, queries.data() + (q * kDynDim))); + } + + index.add(points.data() + (kOld * kDynDim), kNew, nullptr, false, 1); + + double old_error = 0; + double new_error = 0; + for (size_t q = 0; q < kQueries; ++q) { + const float* query = queries.data() + (q * kDynDim); + const auto after = search_everything(index, query); + std::vector by_id(kOld + kNew, std::nanf("")); + for (size_t r = 0; r < after.ids.size(); ++r) { + ASSERT_LT(after.ids[r], kOld + kNew) << "every point is still ranked"; + by_id[after.ids[r]] = after.distances[r]; + } + // A tail batch is unpacked and repacked; old points must not change at all. + for (size_t r = 0; r < before[q].ids.size(); ++r) { + EXPECT_EQ(by_id[before[q].ids[r]], before[q].distances[r]); + } + for (size_t id = 0; id < kOld + kNew; ++id) { + const double error = std::abs( + by_id[id] - + exact_distance(query, points.data() + (id * kDynDim), metric) + ); + (id < kOld ? old_error : new_error) += error; + } + } + old_error /= kQueries * kOld; + new_error /= kQueries * kNew; + // Mis-encoded points would be off by the spread of the data, not by a factor. + EXPECT_LT(new_error, 3 * old_error + 1e-4) << old_error << " vs " << new_error; + } + } +} + +TEST(IvfAddTest, AddedRawVectorsAreRerankedExactly) { + constexpr size_t kOld = 40; + constexpr size_t kNew = 60; + const auto points = random_rows(kOld + kNew, 31); + const auto centroids = random_rows(kDynClusters, 32); + const auto queries = random_rows(3, 33); + for (auto metric : {METRIC_L2, METRIC_IP}) { + IVF index(kOld, kDynDim, kDynClusters, 32, metric); + const auto assigned = nearest_centroids(centroids, points.data(), kOld, metric); + index.construct(points.data(), centroids.data(), assigned.data(), false, 1); + index.add(points.data() + (kOld * kDynDim), kNew, nullptr, false, 1); + for (size_t q = 0; q < 3; ++q) { + const float* query = queries.data() + (q * kDynDim); + const auto hits = search_everything(index, query); + for (size_t r = 0; r < hits.ids.size(); ++r) { + EXPECT_NEAR( + hits.distances[r], + exact_distance(query, points.data() + (hits.ids[r] * kDynDim), metric), + 1e-4 + ); + } + EXPECT_TRUE(std::is_sorted(hits.distances.begin(), hits.distances.end())); + } + } +} + +TEST(IvfAddTest, FasterQuantizationIsSupportedAndDeterministic) { + constexpr size_t kOld = 100; + constexpr size_t kNew = 50; + const auto points = random_rows(kOld + kNew, 41); + const auto centroids = random_rows(kDynClusters, 42); + const auto assigned = + nearest_centroids(centroids, points.data(), kOld + kNew, METRIC_L2); + const std::string base_path = ::testing::TempDir() + "rabitq_ivf_add_faster_base.index"; + const std::string path = ::testing::TempDir() + "rabitq_ivf_add_faster.index"; + + IVF base(kOld, kDynDim, kDynClusters, 5); + base.construct(points.data(), centroids.data(), assigned.data(), true, 1); + base.save(base_path.c_str()); + IVF first; + IVF second; + first.load(base_path.c_str()); + second.load(base_path.c_str()); + first.add(points.data() + (kOld * kDynDim), kNew, assigned.data() + kOld, true, 1); + second.add(points.data() + (kOld * kDynDim), kNew, assigned.data() + kOld, true, 2); + EXPECT_TRUE(saved_bytes(first, path) == saved_bytes(second, path)); + const auto hits = search_everything(first, points.data() + (kOld * kDynDim)); + EXPECT_EQ(hits.ids.size(), kOld + kNew); + EXPECT_TRUE(std::none_of(hits.distances.begin(), hits.distances.end(), [](float x) { + return std::isnan(x); + })); + std::remove(base_path.c_str()); + std::remove(path.c_str()); +} + +TEST(IvfAddTest, RoutesNewPointsWithHNSWCentroids) { + constexpr size_t kDim = 64; + constexpr size_t kCount = 20000; + // Centroid i has first coordinate i + 1, so an inner product with e0 is largest for the + // last centroid and the routing decision is unambiguous. + std::vector centroids(kCount * kDim, 0.0F); + for (size_t i = 0; i < kCount; ++i) { + centroids[i * kDim] = static_cast(i + 1); + } + std::array first{}; + first[0] = 1.0F; + const PID cluster = 0; + IVF index(1, kDim, kCount, 32, METRIC_IP); + index.construct(first.data(), centroids.data(), &cluster, false, 4); + + std::array second{}; + second[0] = 2.0F; + index.add(second.data(), 1, nullptr, false, 2); + ASSERT_EQ(index.max_elements(), 2U); + + // The query is routed by the same rule, so it reaches the cluster that received the + // point. + std::array result{kPidMax}; + std::array distance{}; + index.search(second.data(), 1, 1, result.data(), distance.data()); + EXPECT_EQ(result[0], 1U); + EXPECT_NEAR(distance[0], 1.0F - 4.0F, 1e-5F); +} + +TEST(IvfAddTest, RejectsBadInputAndLeavesTheIndexUnchanged) { + constexpr size_t kNum = 70; + const auto points = random_rows(kNum + 1, 51); + const auto centroids = random_rows(kDynClusters, 52); + const auto assigned = nearest_centroids(centroids, points.data(), kNum, METRIC_L2); + const std::string path = ::testing::TempDir() + "rabitq_ivf_add_reject.index"; + + IVF unbuilt(kNum, kDynDim, kDynClusters, 4); + EXPECT_THROW(unbuilt.add(points.data(), 1), std::logic_error); + IVF unconfigured; + EXPECT_THROW(unconfigured.add(points.data(), 1), std::logic_error); + + IVF index(kNum, kDynDim, kDynClusters, 4); + index.construct(points.data(), centroids.data(), assigned.data(), false, 1); + const std::string expected = saved_bytes(index, path); + + const std::array bad_clusters{0, static_cast(kDynClusters)}; + EXPECT_THROW(index.add(points.data(), 2, bad_clusters.data()), std::invalid_argument); + EXPECT_THROW(index.add(nullptr, 1), std::invalid_argument); + // The ID check happens before the data is read. + EXPECT_THROW( + index.add(points.data(), buffer::kSearchBufferMaxPointCount), std::invalid_argument + ); + EXPECT_NO_THROW(index.add(nullptr, 0)); + EXPECT_NO_THROW(index.add(points.data(), 0)); + EXPECT_EQ(index.max_elements(), kNum); + EXPECT_TRUE(saved_bytes(index, path) == expected); + std::remove(path.c_str()); +} + +TEST(IvfRemoveTest, ExcludesRemovedPointsOnEveryScanPath) { + constexpr size_t kNum = 200; + constexpr size_t kQueries = 5; + const auto points = random_rows(kNum, 61); + const auto centroids = random_rows(kDynClusters, 62); + const auto queries = random_rows(kQueries, 63); + + std::vector removed_ids; + std::vector is_removed(kNum, false); + for (PID id = 0; id < kNum; id += 3) { + removed_ids.push_back(id); + is_removed[id] = true; + } + removed_ids.push_back(kNum - 1); // the last point sits in a partial batch + is_removed[kNum - 1] = true; + removed_ids.push_back(0); // repeated IDs count once + const size_t unique_removed = + static_cast(std::count(is_removed.begin(), is_removed.end(), true)); + + for (auto metric : {METRIC_L2, METRIC_IP}) { + const auto assigned = nearest_centroids(centroids, points.data(), kNum, metric); + // 1-3 bits use standard FastScan, 4-9 bits and raw storage exercise other paths too + for (size_t bits : {1UL, 2UL, 3UL, 4UL, 5UL, 6UL, 7UL, 8UL, 9UL, 32UL}) { + for (bool hacc : {false, true}) { + SCOPED_TRACE( + ::testing::Message() << "metric=" << (metric == METRIC_L2 ? "l2" : "ip") + << " bits=" << bits << " hacc=" << hacc + ); + IVF index(kNum, kDynDim, kDynClusters, bits, metric); + index.construct(points.data(), centroids.data(), assigned.data(), false, 1); + std::vector before; + for (size_t q = 0; q < kQueries; ++q) { + before.push_back( + search_everything(index, queries.data() + (q * kDynDim), hacc) + ); + } + + EXPECT_EQ( + index.remove(removed_ids.data(), removed_ids.size()), unique_removed + ); + EXPECT_EQ(index.remove(removed_ids.data(), removed_ids.size()), 0U); + EXPECT_EQ(index.max_elements(), kNum); + + for (size_t q = 0; q < kQueries; ++q) { + const auto after = + search_everything(index, queries.data() + (q * kDynDim), hacc); + // Survivors keep their order and distances; the rest are sentinels. + size_t next = 0; + for (size_t r = 0; r < kNum; ++r) { + if (is_removed[before[q].ids[r]]) { + continue; + } + ASSERT_EQ(after.ids[next], before[q].ids[r]); + ASSERT_EQ(after.distances[next], before[q].distances[r]); + ++next; + } + ASSERT_EQ(next, kNum - unique_removed); + for (size_t r = next; r < kNum; ++r) { + EXPECT_EQ(after.ids[r], kPidMax); + EXPECT_EQ( + after.distances[r], std::numeric_limits::infinity() + ); + } + + // A small k must still be filled from the survivors only. + std::array top{}; + std::array top_distances{}; + index.search( + queries.data() + (q * kDynDim), + top.size(), + kDynClusters, + top.data(), + top_distances.data(), + hacc + ); + for (size_t r = 0; r < top.size(); ++r) { + ASSERT_LT(top[r], kNum); + EXPECT_FALSE(is_removed[top[r]]); + EXPECT_TRUE(std::isfinite(top_distances[r])); + } + } + } + } + } +} + +TEST(IvfRemoveTest, SurvivesSaveLoadAndLaterAdds) { + constexpr size_t kOld = 100; + constexpr size_t kNew = 60; + constexpr size_t kQueries = 4; + const auto points = random_rows(kOld + kNew, 71); + const auto centroids = random_rows(kDynClusters, 72); + const auto queries = random_rows(kQueries, 73); + const std::vector removed{1, 33, 34, 98, 99}; + const std::string path = ::testing::TempDir() + "rabitq_ivf_remove.index"; + + // Counts the live results, checking that no removed point appears and that the unfilled + // slots hold the sentinel with an infinite distance. + auto count_live = [&removed](const Hits& hits) { + size_t live = 0; + for (size_t r = 0; r < hits.ids.size(); ++r) { + if (hits.ids[r] == kPidMax) { + EXPECT_EQ(hits.distances[r], std::numeric_limits::infinity()); + continue; + } + ++live; + EXPECT_EQ(std::count(removed.begin(), removed.end(), hits.ids[r]), 0); + EXPECT_TRUE(std::isfinite(hits.distances[r])); + } + return live; + }; + + for (auto metric : {METRIC_L2, METRIC_IP}) { + const auto assigned = nearest_centroids(centroids, points.data(), kOld, metric); + for (size_t bits : {1UL, 2UL, 3UL, 4UL, 5UL, 6UL, 7UL, 8UL, 9UL, 32UL}) { + for (bool hacc : {false, true}) { + SCOPED_TRACE( + ::testing::Message() << "metric=" << (metric == METRIC_L2 ? "l2" : "ip") + << " bits=" << bits << " hacc=" << hacc + ); + IVF index(kOld, kDynDim, kDynClusters, bits, metric); + index.construct(points.data(), centroids.data(), assigned.data(), false, 1); + ASSERT_EQ(index.remove(removed.data(), removed.size()), removed.size()); + index.save(path.c_str()); + + IVF loaded; + loaded.load(path.c_str()); + ASSERT_EQ(loaded.max_elements(), kOld); + for (size_t q = 0; q < kQueries; ++q) { + const float* query = queries.data() + (q * kDynDim); + const auto expected = search_everything(index, query, hacc); + const auto actual = search_everything(loaded, query, hacc); + EXPECT_EQ(actual.ids, expected.ids); + EXPECT_EQ(actual.distances, expected.distances); + EXPECT_EQ(count_live(actual), kOld - removed.size()); + } + EXPECT_EQ(loaded.remove(removed.data(), removed.size()), 0U); + + // Adding repacks the tail batches that hold removed points; they must stay + // removed, and a reloaded index must grow exactly like the original. + const float* added = points.data() + (kOld * kDynDim); + index.add(added, kNew, nullptr, false, 1); + loaded.add(added, kNew, nullptr, false, 1); + ASSERT_EQ(loaded.max_elements(), kOld + kNew); + for (size_t q = 0; q < kQueries; ++q) { + const float* query = queries.data() + (q * kDynDim); + const auto expected = search_everything(index, query, hacc); + const auto actual = search_everything(loaded, query, hacc); + EXPECT_EQ(actual.ids, expected.ids); + EXPECT_EQ(actual.distances, expected.distances); + EXPECT_EQ(count_live(actual), kOld + kNew - removed.size()); + } + } + } + } + std::remove(path.c_str()); +} + +TEST(IvfRemoveTest, RejectsBadInputAndRemovesNothingOnFailure) { + constexpr size_t kNum = 70; + const auto points = random_rows(kNum, 81); + const auto centroids = random_rows(kDynClusters, 82); + const auto assigned = nearest_centroids(centroids, points.data(), kNum, METRIC_L2); + const std::string path = ::testing::TempDir() + "rabitq_ivf_remove_reject.index"; + + const PID id = 0; + IVF unbuilt(kNum, kDynDim, kDynClusters, 4); + EXPECT_THROW(unbuilt.remove(&id, 1), std::logic_error); + + IVF index(kNum, kDynDim, kDynClusters, 4); + index.construct(points.data(), centroids.data(), assigned.data(), false, 1); + const std::string expected = saved_bytes(index, path); + + const std::array out_of_range{3, static_cast(kNum)}; + EXPECT_THROW(index.remove(out_of_range.data(), 2), std::invalid_argument); + EXPECT_THROW(index.remove(nullptr, 1), std::invalid_argument); + EXPECT_EQ(index.remove(nullptr, 0), 0U); + EXPECT_TRUE(saved_bytes(index, path) == expected); + + // Removing everything leaves a searchable index that returns only sentinels. + std::vector everything(kNum); + for (size_t i = 0; i < kNum; ++i) { + everything[i] = static_cast(i); + } + EXPECT_EQ(index.remove(everything.data(), kNum), kNum); + const auto hits = search_everything(index, points.data()); + for (size_t r = 0; r < kNum; ++r) { + EXPECT_EQ(hits.ids[r], kPidMax); + } + std::remove(path.c_str()); +} + } // namespace } // namespace rabitqlib::ivf