From 43a2debfefa511fe00ad08def504d8ea20ab45e2 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sun, 1 Nov 2026 18:06:15 +0800 Subject: [PATCH 1/2] feat(python): train, set_centroids and add for IvfIndex and HnswIndex IvfIndex and HnswIndex now follow a train -> add -> search workflow, so callers no longer run k-means and pass centroids and assignments to build. - IvfIndex(dim, num_clusters, *, nbits, metric, initializer, seed=None) and HnswIndex(dim, *, nbits, num_clusters=16, M, ef_construction, metric, seed=100). nbits is required; seed=None draws one seed, which index.seed reports. - index.cp (RaBitQKMeansParameters: niter 25, max_points_per_centroid 256, spherical for inner product, seed from the index), train(x, *, num_threads=None), set_centroids(centroids), is_trained, len(index). - add(data, *, cluster_ids=None, num_threads=0, fast_quantization=False). train and add default to every CPU; search keeps one thread. HNSW capacity grows geometrically on add, and search bounds k by the stored point count. - A 0.8 compatibility layer in rabitqlib/__init__.py patches the native classes, as Faiss does for its SWIG classes: the 0.7 constructor forms, max_elements, random_seed, build(), is_built and positional add options still work, raise DeprecationWarning at the caller, and give the same results as 0.7. Stubs mark them @deprecated. - IVF set_centroids rotates centroids in parallel, as construct does. - Docs, samples and a 0.7 -> 0.8 migration page use the new API. --- README.md | 20 +- docs/docs/clustering.md | 15 +- docs/docs/index/hnsw.md | 45 ++- docs/docs/index/ivf.md | 62 ++- docs/docs/index/usage.md | 30 +- docs/docs/modular_migration.md | 2 +- docs/docs/python_migration.md | 76 ++++ docs/docs/quick_start.md | 38 +- docs/mkdocs.yml | 1 + include/rabitqlib/index/ivf/ivf.hpp | 3 + python_bindings/__init__.py | 238 ++++++++++++ python_bindings/__init__.pyi | 460 ++++++++++++++++------ python_bindings/bindings_common.hpp | 62 +++ python_bindings/hnsw_bindings.cpp | 495 +++++++++++++++++------- python_bindings/ivf_bindings.cpp | 383 +++++++++++++----- sample/python/README.md | 40 +- sample/python/hnsw_rabitq_indexing.py | 47 ++- sample/python/ivf_rabitq_indexing.py | 50 ++- src/index/ivf.cpp | 12 +- tests/python/test_compat_08.py | 214 ++++++++++ tests/python/test_editor_support.py | 78 +++- tests/python/test_examples.py | 26 ++ tests/python/test_hnsw.py | 19 +- tests/python/test_index_repr.py | 6 +- tests/python/test_input_boundaries.py | 5 +- tests/python/test_ivf.py | 8 +- tests/python/test_train_api.py | 395 +++++++++++++++++++ tests/python/test_txtai_compat.py | 95 +++++ tests/python/typing/public_api.py | 71 ++-- tests/test_mypy.py | 47 +++ tests/unit/rabitqlib/index/ivf_test.cpp | 2 +- 31 files changed, 2488 insertions(+), 557 deletions(-) create mode 100644 docs/docs/python_migration.md create mode 100644 tests/python/test_compat_08.py create mode 100644 tests/python/test_train_api.py create mode 100644 tests/python/test_txtai_compat.py diff --git a/README.md b/README.md index 70bea351..ebb47d6d 100644 --- a/README.md +++ b/README.md @@ -60,26 +60,26 @@ required. This example demonstrates the API, not performance. ```python import numpy as np -from rabitqlib import FinalAssignmentMode, IvfIndex, RaBitQKMeans +from rabitqlib import IvfIndex rng = np.random.default_rng(42) data = rng.standard_normal((500, 64)).astype(np.float32) queries = rng.standard_normal((5, 64)).astype(np.float32) -clustering = RaBitQKMeans( - 64, 5, num_threads=2, final_assignment=FinalAssignmentMode.Exact -) -clustering.train(data) -index = IvfIndex(dim=64, max_elements=len(data), num_clusters=5, nbits=4, metric="l2") -index.build(data, clustering.centroids, clustering.assignments) +index = IvfIndex(64, 5, nbits=4, metric="l2") +index.train(data) +index.add(data) ids, distances = index.search(queries, k=10, nprobe=5) print(ids.shape, distances.shape) # (5, 10) (5, 10) ``` -IDs identify input rows; distances here estimate squared L2 distance, with smaller -values first. Quantization remains approximate even when all clusters are probed. +`train` learns the centroids, and `add` quantizes the vectors into their +clusters. IDs identify added rows; distances here estimate squared L2 distance, +with smaller values first. Quantization remains approximate even when all clusters are probed. The [complete quick start](https://vectordb-ntu.github.io/RaBitQ-Library/quick_start/) explains result handling, -measures recall while varying `nprobe`, and shows save/load. +measures recall while varying `nprobe`, and shows save/load. Code written for +0.7 still runs in 0.8 with deprecation warnings; see the +[migration guide](https://vectordb-ntu.github.io/RaBitQ-Library/python_migration/). For other indexes, see the [HNSW-RaBitQ](https://vectordb-ntu.github.io/RaBitQ-Library/index/hnsw/#minimal-workflow) and [SymphonyQG](https://vectordb-ntu.github.io/RaBitQ-Library/index/qg/#minimal-workflow) examples. diff --git a/docs/docs/clustering.md b/docs/docs/clustering.md index 8ab5fc54..9aabbc5c 100644 --- a/docs/docs/clustering.md +++ b/docs/docs/clustering.md @@ -7,6 +7,14 @@ algorithms based on the cluster count. For installation, see [Quick start](quick_start.md#install). +You do not need these classes to build an index: `IvfIndex.train` and +`HnswIndex.train` run them for you (see +[IVF-RaBitQ training](index/ivf.md#training-centroids-in-python)). Use them +directly to inspect assignments, reuse one clustering across several indexes, +or pass the result to `set_centroids` and `add(..., cluster_ids=...)`. The +indexes train with their own defaults, including sampling at most 256 vectors +per centroid; the classes below default to training on every vector. + ## Choose a method | Class | Assignment | Recommended use | Cluster count | @@ -118,8 +126,11 @@ and `n > k * max_points_per_centroid`, training draws does; Faiss enables this by default with 256. Centroid initialization and all Lloyd iterations use only this sample. The final centroids then assign all `n` vectors with the selected `final_assignment` mode, as [`assign`](#assign-new-vectors) -does, so `assignments` and `distances` keep shape `(n,)` and can be passed -directly to index builders. +does, so `assignments` and `distances` keep shape `(n,)` and `assignments` can be +passed directly to an index's `add(..., cluster_ids=...)`. C++ callers that need +only the centroids can call `train_centroids(n, x)` instead of `train`, which +skips this final assignment and leaves `assignments`, `distances` and +`final_obj` empty or zero. `seed` selects the sample through a random stream separate from centroid initialization. Training that does not sample, because the parameter is `0` or diff --git a/docs/docs/index/hnsw.md b/docs/docs/index/hnsw.md index 16e82a0f..8294f60c 100644 --- a/docs/docs/index/hnsw.md +++ b/docs/docs/index/hnsw.md @@ -13,20 +13,25 @@ degree, and index parameters. ### Python Run the [Quick start](../quick_start.md#build-and-search-an-ivf-index) to prepare -`data`, `queries` and `clustering`, then use those same arrays with HNSW-RaBitQ: +`data` and `queries`, then use those same arrays with HNSW-RaBitQ: ```python from rabitqlib import HnswIndex -index = HnswIndex( - dim=data.shape[1], max_elements=len(data), M=16, ef_construction=100, nbits=4, -) -index.build(data, clustering.centroids, clustering.assignments) +index = HnswIndex(data.shape[1], nbits=4, M=16, ef_construction=100) +index.train(data) # 16 RaBitQKMeans centroids by default +index.add(data, num_threads=1) # one thread gives a reproducible graph ids, distances = index.search(queries, k=10, ef=100) index.save("hnsw.index") loaded = HnswIndex.load("hnsw.index") ``` +`train` and `set_centroids` work as for [IVF-RaBitQ](ivf.md#training-centroids-in-python), +with `num_clusters` centroids (default 16) and always RaBitQKMeans; `index.cp` +starts from the same defaults with `seed`'s low 32 bits. The first `add` into the +empty index links vectors with their exact coordinates, like the 0.7 `build`; +later adds link through the stored codes (see [Recall after `add`](#recall-after-add)). + ### C++ Link `rabitqlib::index`. Here `data`, `centroids` and `queries` are row-major @@ -58,16 +63,19 @@ see [C++ integration](../cpp_integration.md). | Parameter | Contract | | --- | --- | | `dim` | Original input dimension; Python supports 64–65,536 and pads internally. | -| `max_elements` | Capacity; build accepts 1 through this many rows. Reserve extra capacity or call `resize` before adding beyond it. | -| `nbits` / C++ `total_bits` | 1–9 total bits including the sign bit; Python defaults to 8. Raw-vector storage is not supported. | -| `M`, `ef_construction` | Graph connectivity and construction search budget; Python defaults are 16 and 200. `M` must be at least 2 and is capped at 10,000; `ef_construction` is raised to at least `M`. | -| `metric`, `random_seed` | Python defaults are `"l2"` and 100; `"ip"` is also supported. | -| `k`, `ef` | Positive result count and search window. Python requires `k <= max_elements`; omitted/zero `ef` uses `max(k, 10)`, positive `ef` is raised to at least `k`, and `ef` above the stored point count is clamped to it. | - -`random_seed` controls graph levels and the FHT/Kac rotation (using its low 32 bits). +| `nbits` / C++ `total_bits` | Required and, in Python, keyword-only: 1–9 total bits including the sign bit. Raw-vector storage is not supported. | +| `num_clusters` | Python keyword, default 16: centroids `train` learns or `set_centroids` supplies. | +| `M`, `ef_construction` | Graph connectivity and construction search budget; Python defaults are 16 and 200. `M` must be at least 2 and is capped at 10,000; `ef_construction` is raised to at least `M`. The base layer keeps up to `2 * M` edges per node. | +| `metric`, `seed` | Python defaults are `"l2"` and 100; `"ip"` is also supported. `seed` was called `random_seed` before 0.8. | +| `max_elements` | Python reports the current capacity, which `add` grows geometrically (at least doubling) and `resize` can preallocate. C++ `add` never grows it. | +| `k`, `ef` | Positive result count and search window. Python requires `k <= num_stored_points`, so a trained index without points rejects every `k`; omitted/zero `ef` uses `max(k, 10)`, positive `ef` is raised to at least `k`, and `ef` above the stored point count is clamped to it. | + +`seed` controls graph levels and the FHT/Kac rotation (using its low 32 bits). With identical inputs and parameters, single-threaded construction is reproducible -within the same build environment. Parallel insertion order and different toolchains -or SIMD backends can still change results. Loading restores the saved rotation. +within the same build environment. Python `train` and `add` default to every +available CPU, so **default builds are not reproducible**: pass `num_threads=1` to +`add` when you need the same graph again. Different toolchains or SIMD backends can +still change results. Loading restores the saved rotation. Results use squared L2 or `1 - dot` estimates, ordered from smallest to largest. Python fills missing results with ID `2**32 - 1` and distance `+inf`; C++ returns @@ -196,11 +204,9 @@ void HierarchicalNSW::resize(size_t new_max_elements); In C++, initialize the index through its parameterized constructor or `load` before calling `resize`. -In Python: +In Python, `add` grows the capacity when needed: ```python -# Grow before adding; preserve any capacity already available. -index.resize(max(index.max_elements, index.num_stored_points + len(vectors))) new_ids = index.add(vectors) # route to the nearest centroids # Alternatively, choose the clusters instead of the call above: # new_ids = index.add(vectors, cluster_ids=labels) @@ -221,9 +227,10 @@ retrying, or rebuild to restore full reachability. See the shared [index count contract](usage.md#index-counts). Removed records still occupy capacity. -`add` never grows the index: it rejects a whole batch whose size exceeds +C++ `add` never grows the index: it rejects a whole batch whose size exceeds `max_elements() - num_stored_points()`, and `resize` is the only way to raise -that. Reserving capacity avoids reallocating and copying the whole base-layer record storage during `add`. +that. Python `add` calls `resize` first when the batch does not fit, to at least +double the capacity. Reserving capacity avoids reallocating and copying the whole base-layer record storage during `add`. It does not make insertion allocation-free: batch scratch, upper-layer link lists, and label-map entries still require allocations, and the label map may rehash. diff --git a/docs/docs/index/ivf.md b/docs/docs/index/ivf.md index e6ec814a..cf0cf8cf 100644 --- a/docs/docs/index/ivf.md +++ b/docs/docs/index/ivf.md @@ -14,8 +14,9 @@ width, number of clusters, and `nprobe`. ### Python The [complete Quick Start](../quick_start.md) is the runnable IVF-RaBitQ example: -train centroids, call `IvfIndex.build`, search, then save and reload the index. -The update examples below reuse its built `index`. +create an index, `train` its centroids, `add` the vectors, search, then save and +reload it. [Training](#training-centroids-in-python) describes the clustering +options, and the update examples below reuse the Quick Start's `index`. ### C++ @@ -46,12 +47,12 @@ See [C++ integration](../cpp_integration.md) for build setup and | Parameter | Contract | | --- | --- | | `dim` | Original input dimension; Python supports 64–65,536 and pads internally. | -| `max_elements` | Python IVF-RaBitQ build requires exactly this many rows; `add` grows storage automatically. C++ calls this constructor argument `num_points`. | -| `num_clusters` | Must match the centroid count; every cluster ID must be in `[0, num_clusters)`. | -| `nbits` / C++ `total_bits` | 1–9 total bits including the sign bit, or 32 for owned raw-vector reranking. | +| `num_clusters` | Number of centroids `train` learns or `set_centroids` supplies; every cluster ID must be in `[0, num_clusters)`. | +| `nbits` / C++ `total_bits` | Required and, in Python, keyword-only: 1–9 total bits including the sign bit, or 32 for owned raw-vector reranking. | +| `seed` | Rotation seed in `[0, 2**32 - 1]` and the initial `cp.seed`. Python's default `None` draws one at random; `index.seed` reports it, and passing it back reproduces the index. | | `metric` | Python `"l2"` or `"ip"`; defaults to `"l2"`. | | `initializer` | Defaults to `"auto"`; see [centroid routing](#centroid-routing). | -| `k`, `nprobe` | Positive result count and cluster probe count; C++ and Python both require `k` not to exceed `max_elements`, which `add` increases. | +| `k`, `nprobe` | Positive result count and cluster probe count; C++ and Python both require `k` not to exceed the stored point count (`max_elements`), which `add` increases. | | `high_accuracy` | Python defaults to `None`: HACC for 4–9 bits, standard FastScan otherwise. | Results contain IDs and squared L2 or `1 - dot` distances, ordered from smallest @@ -125,8 +126,7 @@ The optional final C++ constructor argument selects `InitializerType::Auto`, `initializer="auto"`, `"flat"`, `"flat_rabitq"`, or `"hnsw"`: ```python -index = IvfIndex(dim, max_elements, num_clusters, nbits=4, - initializer="flat_rabitq") +index = IvfIndex(dim, num_clusters, nbits=4, initializer="flat_rabitq") print(index.initializer) # resolved routing type ``` @@ -199,6 +199,39 @@ For example: ivf.construct(data.data(), centroids.data(), cluster_ids.data(), true); ``` +### Training centroids in Python + +`train(x)` learns `num_clusters` centroids from the rows of `x`, which must have +at least `num_clusters` rows, and stores them in the index; `add` then assigns +and quantizes vectors. Training uses the settings in `index.cp`, a +[`RaBitQKMeansParameters`](../clustering.md#rabitqkmeans) whose defaults differ +from the k-means classes: 25 iterations, at most 256 sampled vectors per +centroid (`max_points_per_centroid`), spherical k-means for `metric="ip"`, and +`seed` equal to the index seed. Change fields before training: + +```python +index = IvfIndex(dim, 1024, nbits=4, metric="ip") +index.cp.niter = 10 +index.train(sample) # num_threads=None uses index.cp.num_threads (0: all CPUs) +index.add(vectors) +``` + +- Up to 2,000 clusters, training uses RaBitQKMeans; above that it uses + [QGKMeans](../clustering.md#qgkmeans), copying the shared fields of `cp` and + keeping its default graph settings. The threshold is provisional and has not + been measured. +- `set_centroids(centroids)` uses centroids of shape `(num_clusters, dim)` + computed elsewhere instead of training. +- Training or setting centroids again replaces them while the index stores no + vectors; once it stores any (including removed ones), both raise + `RuntimeError`. A failed call leaves the index unchanged. +- Before training, `add`, `search`, `remove`, `reserve` and `save` raise + `RuntimeError`. A trained index without vectors cannot be saved, and its + `search` rejects every `k` with `ValueError`; add vectors first. Loaded + indexes are trained, and their `seed` is `None`. +- `train` and `add` default to every available CPU; `search` defaults to one + thread. + ### Training centroids in C++ Instead of passing centroids to `construct`, create an index without a point @@ -216,7 +249,11 @@ 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 + rotation seed, followed by the optional rotator and initializer types. Pass + the seed: without it, four arguments such as `IVF(dim, num_clusters, 4, + METRIC_IP)` select the point-count constructor instead, reading them as + `(num_points, dim, num_clusters, total_bits)` with the metric converted to a + bit count, which may compile and run without an error. 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 @@ -358,10 +395,11 @@ Python `vectors` has shape `(n, index.dim)` and is converted to contiguous float one-dimensional integer array or sequence with one valid cluster ID per vector; `remove()` likewise accepts a one-dimensional integer array or sequence of point IDs and returns the number newly removed. Duplicate IDs count once. Python -`num_threads` defaults to 1; `fast_quantization` defaults to `False`. +`num_threads` defaults to 0 (every available CPU) and `fast_quantization` to +`False`; pass both by keyword. -After appending, rebuilding with `construct` or `build` requires -`num_stored_points` rows, including removed records. Updates require exclusive +After appending, rebuilding with C++ `construct` requires `num_stored_points` +rows, including removed records. Updates require exclusive access; see the [concurrency contract](usage.md#threading-and-file-paths). diff --git a/docs/docs/index/usage.md b/docs/docs/index/usage.md index cde6f5fb..b0516f6d 100644 --- a/docs/docs/index/usage.md +++ b/docs/docs/index/usage.md @@ -17,7 +17,8 @@ storage before native processing; inputs remain unchanged. Data, centroids and queries must be finite. Coordinates are also limited to the magnitude the clustering APIs accept, `sqrt(sqrt(FLT_MAX) / (64 * dim))` (about `5e7` for `dim = 128`), so that squared distances and their products stay inside -float32. `build`, `add` and `search` raise `ValueError` (C++: +float32. `train`, `set_centroids`, `add` and `search` (and SymphonyQG `build`) +raise `ValueError` (C++: `std::invalid_argument`) for NaN, infinite, or larger coordinates before changing the index. The C++ `quantize_*` encoders in `core/encoder.hpp` assume their callers have validated the input. @@ -42,22 +43,22 @@ IVF-RaBitQ, HNSW-RaBitQ, and SymphonyQG expose these read-only Python properties | `num_points` | Nonremoved points | | `num_stored_points` | Stored records, including removed ones | -`0 <= num_points <= num_stored_points <= max_elements`. Before building, both -point counts are zero; configured capacity may be nonzero. - -At build time, IVF-RaBitQ requires exactly `max_elements` input rows. HNSW-RaBitQ accepts -1 through `max_elements` rows, so it can reserve capacity for later additions. +`0 <= num_points <= num_stored_points <= max_elements`. Before vectors are +added, both point counts are zero; configured capacity may be nonzero. For +IVF-RaBitQ and HNSW-RaBitQ, `len(index)` equals `num_points`, and `is_trained` +reports whether centroids are set (by `train`, `set_centroids`, or `load`). Removal decreases only `num_points`; slots and IDs remain occupied. Append IDs start at the previous `num_stored_points`. Use this stored count for ID bounds, and `num_points` for a live count. Counts survive save/load. -IVF-RaBitQ grows automatically and has no spare logical capacity after build or append. -HNSW-RaBitQ's `resize` changes only capacity. SymphonyQG has no incremental updates; -all three counts are equal after build. +IVF-RaBitQ grows automatically and has no spare logical capacity after an append. +HNSW-RaBitQ's Python `add` grows its capacity geometrically, and `resize` changes +only capacity, for example to preallocate it. SymphonyQG has no incremental +updates; all three counts are equal after build. -`repr(index)` shows dimension, metric, quantization bits, live point count and build -state. IVF/HNSW also show stored count, capacity and centroid count; IVF includes +`repr(index)` shows dimension, metric, quantization bits and live point count, +with the training state for IVF/HNSW and the build state for SymphonyQG. IVF/HNSW also show stored count, capacity and centroid count; IVF includes the resolved initializer, and SymphonyQG includes its degree bound. The summary reads only metadata and follows the same concurrency rules as property reads. @@ -118,7 +119,7 @@ construction, updates, and file I/O. The following contract applies to | Operation on the same index | Concurrent access | | --- | --- | | `search`, `search_batch`, `save`, `repr`, and property reads | May run together | -| `build`, `add`, `remove`, `reserve`, and `resize`, where available | Require exclusive access | +| `train`, `set_centroids`, `build`, `add`, `remove`, `reserve`, `resize`, and setting `cp`, where available | Require exclusive access | A conflicting call immediately raises `RuntimeError("Index is busy: conflicting operation in progress")`. @@ -138,7 +139,10 @@ index updates themselves and use separate search output buffers. `num_threads=0` selects the detected available logical CPU count. Positive values set an upper limit, capped at that count; small workloads may use fewer workers. -Python index methods default to one thread; clustering defaults to `0`. +IVF-RaBitQ and HNSW-RaBitQ `train` and `add` default to `0`, as does clustering; +`search`, the deprecated `build`, and SymphonyQG methods default to one thread. +Multithreaded HNSW-RaBitQ `add` links points in a nondeterministic order, so pass +`num_threads=1` when a build must be reproducible. On Linux, CPU detection accounts for affinity and OpenMP binding. Apply binding before starting the program. diff --git a/docs/docs/modular_migration.md b/docs/docs/modular_migration.md index e0843b10..801282f9 100644 --- a/docs/docs/modular_migration.md +++ b/docs/docs/modular_migration.md @@ -175,7 +175,7 @@ ID bounds, append offsets, and minimum resize capacity. Use `num_points` for a live count, without subtracting a separate deletion counter. The IVF-RaBitQ constructor keyword `max_elements` still specifies the required initial -build size. SymphonyQG's C++ `num_vertices()` still returns the configured graph +build size; 0.8 deprecates it (see [Migrating to 0.8](python_migration.md)). SymphonyQG's C++ `num_vertices()` still returns the configured graph size for builders. Existing index files need no conversion: loading recovers counts from their removal markers. diff --git a/docs/docs/python_migration.md b/docs/docs/python_migration.md new file mode 100644 index 00000000..6f4d80dd --- /dev/null +++ b/docs/docs/python_migration.md @@ -0,0 +1,76 @@ +# Migrating Python code from 0.7 to 0.8 + +In 0.8, `IvfIndex` and `HnswIndex` follow a `train` → `add` → `search` workflow, +as Faiss indexes do: the index learns its own centroids, so you no longer run +k-means and pass the centroids and assignments to `build`. Code written for 0.7 +still runs in 0.8. Each old form raises a `DeprecationWarning` at your call site +and will be removed in **0.9**. Search calls are unchanged. + +## Old and new calls + +| 0.7 | 0.8 | +| --- | --- | +| `IvfIndex(dim, max_elements, num_clusters, nbits, metric, initializer)` | `IvfIndex(dim, num_clusters, nbits=..., metric=..., initializer=..., seed=None)`; storage grows on `add` | +| `IvfIndex(dim=..., max_elements=..., num_clusters=..., nbits=...)` | the same, without `max_elements` | +| `HnswIndex(dim, max_elements, M=16, ef_construction=200, nbits=8, metric="l2", random_seed=100)` | `HnswIndex(dim, nbits=..., num_clusters=16, M=16, ef_construction=200, metric="l2", seed=100)`; capacity grows on `add`, and `resize(n)` preallocates it | +| `HnswIndex(..., random_seed=s)` | `HnswIndex(..., seed=s)` | +| `kmeans = RaBitQKMeans(dim, k)`
`kmeans.train(x)`
`index.build(x, kmeans.centroids, kmeans.assignments)` | `index.train(x)`
`index.add(x)` | +| `index.build(x, centroids, cluster_ids)` | `index.set_centroids(centroids)`
`index.add(x, cluster_ids=cluster_ids)` | +| `index.add(x, cluster_ids, num_threads, fast_quantization)` | `index.add(x, cluster_ids=..., num_threads=..., fast_quantization=...)` | +| `index.is_built` | `index.is_trained`, or `len(index) > 0` for "holds vectors" | + +Except for `dim` and IVF-RaBitQ's `num_clusters`, constructor arguments are +keyword-only in 0.8, and `nbits` has no default. HNSW-RaBitQ used to default to +`nbits=8`; the old form keeps that default. + +`train` uses the settings in `index.cp`: 25 iterations, at most 256 sampled +vectors per centroid, spherical k-means for `metric="ip"`, and the index seed. +See [IVF-RaBitQ training](index/ivf.md#training-centroids-in-python) for the +details and for `set_centroids`. The [clustering classes](clustering.md) still +work if you want to run k-means yourself. + +## Behavior changes + +Unless an item says otherwise, a 0.7 call that raises a deprecation warning +keeps its 0.7 behavior. Calls that are valid in both versions, such as +`add(x)` and `search`, follow the 0.8 rules below. + +- **Thread defaults.** `train` and `add` now default to every available CPU + (`num_threads=0`); 0.7 `build` and `add` defaulted to one thread. This applies + to every `add` call that does not pass `num_threads`, including 0.7 code such as + `index.add(x)`, without a warning. `search` still defaults to one thread, and + the deprecated `build` keeps its default of 1. +- **HNSW-RaBitQ builds are not reproducible by default.** With more than one + thread, `add` links points in a nondeterministic order. Pass `num_threads=1` + to `add` for the same graph on every run. With one thread and the same seed, + `set_centroids` followed by `add` produces exactly the index that 0.7 `build` + produced from the same centroids and cluster IDs. +- **HNSW-RaBitQ `search` limits `k` to the stored point count** + (`num_stored_points`) instead of the capacity, since capacity now grows + geometrically. A `k` between the two used to return padding rows and now + raises `ValueError`, as it already did for IVF-RaBitQ. +- **HNSW-RaBitQ capacity grows.** `add` resizes the index, at least doubling it, + instead of raising `ValueError` when the batch does not fit. This also applies + to `add` on indexes created in the old form, and to the deprecated `build`, + which no longer rejects more rows than `max_elements`. +- **State errors.** Calling `add`, `search`, `remove`, `reserve` or `save` + before training raises `RuntimeError("IvfIndex must be trained or loaded + before ...")`; the 0.7 message said "built". A trained index without vectors + cannot be saved. `repr` shows `is_trained` instead of `is_built`. +- **Seeds.** `IvfIndex` still draws a random rotation seed by default, and now + reports it as `index.seed`; pass it back as `seed=` to reproduce an index. + Indexes loaded from files report `seed=None`. +- **HNSW-RaBitQ `num_clusters`** reports the configured count (16 by default) + before training, where 0.7 reported 0 before `build`. +- An `IvfIndex` created with `max_elements` is filled by `build`; `train` and + `set_centroids` raise `RuntimeError` for it. Create the index in the 0.8 form + to train it. + +Saved files, the `max_elements` property, and `search` results for the same +index are unchanged. Files written by 0.7 load as trained indexes. + +For C++, the existing constructors and `construct` are unchanged; see +[IVF-RaBitQ](index/ivf.md#training-centroids-in-c) and +[HNSW-RaBitQ](index/hnsw.md#training-centroids-in-c) for the new training entry +points. The one C++ behavior change is that saving an index without stored +points now throws `std::logic_error`. diff --git a/docs/docs/quick_start.md b/docs/docs/quick_start.md index b45f728a..939a4a61 100644 --- a/docs/docs/quick_start.md +++ b/docs/docs/quick_start.md @@ -1,7 +1,8 @@ # Quick start Build, search, save and reload a small IVF-RaBitQ index with Python. This guide describes -**0.7.2** and uses synthetic data, so no dataset download is required. +**0.7.2** and uses synthetic data, so no dataset download is required. Coming from 0.7? +See the [migration guide](python_migration.md). For C++, go directly to [C++ integration](cpp_integration.md). ## Install @@ -27,37 +28,34 @@ require a dataset download: ```python import numpy as np -from rabitqlib import FinalAssignmentMode, IvfIndex, RaBitQKMeans +from rabitqlib import IvfIndex rng = np.random.default_rng(42) data = rng.standard_normal((500, 64)).astype(np.float32) queries = rng.standard_normal((5, 64)).astype(np.float32) -clustering = RaBitQKMeans( - 64, 5, num_threads=2, final_assignment=FinalAssignmentMode.Exact -) -clustering.train(data) - -# IVF-RaBitQ requires exactly max_elements rows when building. -index = IvfIndex( - dim=64, - max_elements=len(data), - num_clusters=5, - nbits=4, - metric="l2", -) -index.build(data, clustering.centroids, clustering.assignments) +index = IvfIndex(64, 5, nbits=4, metric="l2", seed=42) +index.train(data) # learn 5 centroids +index.add(data) # quantize every vector into its nearest cluster ids, distances = index.search(queries, k=10, nprobe=5) print(ids.shape, distances.shape) # (5, 10) (5, 10) print(ids[0]) ``` +`train` clusters the data with the settings in `index.cp`: 25 k-means +iterations on at most 256 sampled vectors per centroid, spherical k-means for +`metric="ip"`, and the index `seed`. Change `index.cp` before training to tune +them, or call `index.set_centroids(centroids)` to use centroids you computed +elsewhere. `train` and `add` use every available CPU by default; pass +`num_threads=` to limit them. Without `seed`, each index draws a random rotation +seed and reports it as `index.seed`. + ## Read the results -Both outputs have shape `(number_of_queries, k)`. IDs are zero-based row numbers -in the data supplied to `build`; added vectors receive new IDs without renumbering -existing ones. Results are ordered by increasing distance: `metric="l2"` estimates +Both outputs have shape `(number_of_queries, k)`. IDs are zero-based and follow +the order vectors were added: here they are row numbers of `data`, and later +`add` calls continue the numbering without renumbering existing vectors. Results are ordered by increasing distance: `metric="l2"` estimates squared Euclidean distance, while `metric="ip"` estimates `1 - dot(data, query)`. Smaller values are better for both metrics. Quantized distances are estimates, not exact distances to the original vectors. @@ -73,7 +71,7 @@ cosine search, normalize database and query vectors first and use `metric="ip"`. ## Measure recall For this small L2 example, compute exact neighbors from the original vectors, -then compare several probe counts on the same built index. Recall@10 is the mean +then compare several probe counts on the same index. Recall@10 is the mean fraction of exact top-10 IDs returned for each query. ```python diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index dcb840a1..abb11e5b 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -9,6 +9,7 @@ copyright: Copyright © VectorDB Group, Nanyang Technological University nav: - Home: index.md - Quick start: quick_start.md + - Migrating to 0.8: python_migration.md - Installation: installation.md - C++ Integration: cpp_integration.md - Release notes: diff --git a/include/rabitqlib/index/ivf/ivf.hpp b/include/rabitqlib/index/ivf/ivf.hpp index ef62fa90..8a9c4de0 100644 --- a/include/rabitqlib/index/ivf/ivf.hpp +++ b/include/rabitqlib/index/ivf/ivf.hpp @@ -209,6 +209,9 @@ class IVF { * * The index starts with no point storage. construct rebuilds an index from * max_elements() rows, so it fails here until add has stored points. + * Pass the seed: four arguments such as (dim, num_clusters, 4, METRIC_IP) + * select the point-count constructor instead, which converts the metric to + * total_bits. * * @param dim Original input dimension * @param num_clusters Number of centroids that train or set_centroids provides diff --git a/python_bindings/__init__.py b/python_bindings/__init__.py index dca80bb8..f2d96f4d 100644 --- a/python_bindings/__init__.py +++ b/python_bindings/__init__.py @@ -4,6 +4,9 @@ and adds the clustering wrappers. Import public classes from ``rabitqlib``. """ +import inspect +import warnings + from ._rabitqlib import ( FinalAssignmentMode, HnswIndex, @@ -17,6 +20,241 @@ _RaBitQKMeans, ) +# --- 0.8 compatibility layer; remove in 0.9 -------------------------------- +# The extension classes expose only the 0.8 API. As Faiss does for its SWIG +# classes, the 0.7 call forms are restored by replacing methods on those classes +# at import time, so the public types stay the native ones. + +_MIGRATION = ( + "see https://vectordb-ntu.github.io/RaBitQ-Library/python_migration/; " + "the old form will be removed in 0.9" +) +_MISSING = object() + + +def _deprecated(message): + # Called from a replacement method, so the caller of that method is two up. + warnings.warn(f"{message}; {_MIGRATION}", DeprecationWarning, stacklevel=3) + + +def _bind(signature, name, args, kwargs, hint=""): + try: + bound = signature.bind(*args, **kwargs) + except TypeError as error: + raise TypeError(f"{name}() {error}{hint}") from None + bound.apply_defaults() + arguments = dict(bound.arguments) + del arguments["self"] + return arguments + + +def _ivf_signature( + self, dim, num_clusters, *, nbits, metric="l2", initializer="auto", seed=None +): ... + + +def _ivf_legacy_signature( + self, dim, max_elements, num_clusters, nbits, metric="l2", initializer="auto" +): ... + + +def _hnsw_signature( + self, + dim, + *, + nbits, + num_clusters=16, + M=16, + ef_construction=200, + metric="l2", + seed=100, +): ... + + +def _hnsw_legacy_signature( + self, + dim, + max_elements, + M=16, + ef_construction=200, + nbits=8, + metric="l2", + random_seed=100, +): ... + + +def _add_signature( + self, data, *, cluster_ids=None, num_threads=0, fast_quantization=False +): ... + + +def _patch_ivf_init(cls): + native_init = cls.__init__ + signature = inspect.signature(_ivf_signature) + legacy_signature = inspect.signature(_ivf_legacy_signature) + + def __init__(self, *args, **kwargs): + # A third positional argument, max_elements, or num_clusters by keyword + # after two positional arguments can only be the 0.7 form. + if ( + len(args) > 2 + or "max_elements" in kwargs + or (len(args) == 2 and "num_clusters" in kwargs) + ): + arguments = _bind( + legacy_signature, + "IvfIndex", + (self, *args), + kwargs, + " (read as the 0.7 form IvfIndex(dim, max_elements, num_clusters, " + "nbits, ...); the 0.8 form is IvfIndex(dim, num_clusters, nbits=...))", + ) + _deprecated( + "IvfIndex(dim, max_elements, num_clusters, nbits, ...) is deprecated; " + "use IvfIndex(dim, num_clusters, nbits=...) and train() or " + "set_centroids() before add()" + ) + native_init( + self, + arguments["dim"], + arguments["num_clusters"], + nbits=arguments["nbits"], + metric=arguments["metric"], + initializer=arguments["initializer"], + _max_elements=arguments["max_elements"], + ) + return + native_init(self, **_bind(signature, "IvfIndex", (self, *args), kwargs)) + + __init__.__doc__ = cls.__doc__ + __init__.__signature__ = signature + cls.__init__ = __init__ + + +def _patch_hnsw_init(cls): + native_init = cls.__init__ + signature = inspect.signature(_hnsw_signature) + legacy_signature = inspect.signature(_hnsw_legacy_signature) + + def __init__(self, *args, **kwargs): + if len(args) > 1 or "max_elements" in kwargs: + arguments = _bind( + legacy_signature, + "HnswIndex", + (self, *args), + kwargs, + " (read as the 0.7 form HnswIndex(dim, max_elements, ...); " + "the 0.8 form is HnswIndex(dim, nbits=...))", + ) + _deprecated( + "HnswIndex(dim, max_elements, ...) is deprecated; use " + "HnswIndex(dim, nbits=...), whose capacity grows on add" + ) + native_init( + self, + arguments["dim"], + _max_elements=arguments["max_elements"], + M=arguments["M"], + ef_construction=arguments["ef_construction"], + nbits=arguments["nbits"], + metric=arguments["metric"], + random_seed=arguments["random_seed"], + ) + return + random_seed = kwargs.pop("random_seed", _MISSING) + if random_seed is not _MISSING: + if "seed" in kwargs: + raise TypeError( + "HnswIndex() got both seed and its old name random_seed" + ) + _deprecated("HnswIndex(random_seed=...) is deprecated; use seed=") + kwargs["seed"] = random_seed + native_init(self, **_bind(signature, "HnswIndex", (self, *args), kwargs)) + + __init__.__doc__ = cls.__doc__ + __init__.__signature__ = signature + cls.__init__ = __init__ + + +def _patch_add(cls): + native_add = cls.add + name = cls.__name__ + options = ("cluster_ids", "num_threads", "fast_quantization") + + def add(self, data, *args, **kwargs): + if args: + if len(args) > len(options): + raise TypeError( + f"{name}.add() takes at most {len(options) + 1} positional " + f"arguments ({len(args) + 1} given)" + ) + _deprecated( + f"passing {', '.join(options[: len(args)])} to {name}.add() by " + "position is deprecated; pass them by keyword" + ) + for option, value in zip(options, args, strict=False): + if option in kwargs: + raise TypeError( + f"{name}.add() got multiple values for argument {option!r}" + ) + kwargs[option] = value + return native_add(self, data, **kwargs) + + add.__doc__ = native_add.__doc__ + add.__signature__ = inspect.signature(_add_signature) + cls.add = add + + +def _patch_build(cls): + name = cls.__name__ + + def build( + self, data, centroids, cluster_ids, num_threads=1, fast_quantization=False + ): + """Build the index from vectors and supplied centroid assignments. + + Deprecated since 0.8 and removed in 0.9: use set_centroids() (or + train()) and then add(). The defaults match 0.7, including a single + thread. + + Parameters + ---------- + data : array_like of float, shape (n, dim) + Vectors in the original, unpadded coordinates. Converted to + contiguous float32; a copy may be made. + centroids : array_like of float, shape (num_clusters, dim) + Vectors in the original, unpadded coordinates. Converted to + contiguous float32; a copy may be made. + cluster_ids : array_like of int, shape (n,) + Centroid ID per vector, in [0, num_clusters). Floating-point and + boolean IDs are rejected. + num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. + fast_quantization : bool, optional + Use faster quantizer configuration. Default is False. + """ + _deprecated( + f"{name}.build() is deprecated; use set_centroids() or train(), then add()" + ) + self._build(data, centroids, cluster_ids, num_threads, fast_quantization) + + def is_built(self): + """Deprecated since 0.8: whether the index stores points; use len().""" + _deprecated(f"{name}.is_built is deprecated; use is_trained or len()") + return self.num_stored_points != 0 + + cls.build = build + cls.is_built = property(is_built, doc=is_built.__doc__) + + +_patch_ivf_init(IvfIndex) +_patch_hnsw_init(HnswIndex) +for _cls in (IvfIndex, HnswIndex): + _patch_add(_cls) + _patch_build(_cls) +del _cls +# --- end of the 0.8 compatibility layer ------------------------------------- + class _KMeans: """Shared result handling for the two clustering methods.""" diff --git a/python_bindings/__init__.pyi b/python_bindings/__init__.pyi index 5cd2b5d3..7c5a677e 100644 --- a/python_bindings/__init__.pyi +++ b/python_bindings/__init__.pyi @@ -5,10 +5,11 @@ Array shapes and numeric ranges are documented runtime constraints. from collections.abc import Sequence from os import PathLike -from typing import Any, ClassVar, Literal, SupportsIndex, TypeAlias +from typing import Any, ClassVar, Literal, SupportsIndex, TypeAlias, overload import numpy as np import numpy.typing as npt +from typing_extensions import deprecated __all__ = [ "FinalAssignmentMode", @@ -76,47 +77,54 @@ class IvfIndex: ---------- dim : int Original vector dimension, from 64 through 65536. - max_elements : int - Number of vectors required by the initial build. Later adds grow - storage. num_clusters : int - Number of supplied centroids. + Number of centroids, from train or set_centroids. nbits : int - Total quantization bits, 1 through 9, or 32 for owned float32 vectors - with one-bit filtering and raw reranking. + Keyword-only. Total quantization bits, 1 through 9, or 32 for owned + float32 vectors with one-bit filtering and raw reranking. metric : {'l2', 'ip', 'innerproduct'}, optional Squared L2 or inner-product distance. Default is 'l2'. Vectors are not automatically normalized for cosine similarity. initializer : {'auto', 'flat', 'flat_rabitq', 'hnsw'}, optional Centroid routing strategy. Default is 'auto'. + seed : int, optional + Rotation seed in [0, 2**32 - 1], also the initial cp.seed. None + (default) draws one at random; the seed property reports it. Notes ----- - Build or load the index before searching or saving. Concurrent reads are - supported; an overlapping mutation raises RuntimeError. + Train the index or set its centroids, then add vectors before searching or + saving; loaded indexes are already trained. Concurrent reads are supported; + an overlapping mutation raises RuntimeError. Data, centroids and queries must be finite float32 values within the index coordinate range; other values raise ValueError. + The 0.7 form IvfIndex(dim, max_elements, num_clusters, nbits, ...) and + build() still work in 0.8, raise DeprecationWarning, and will be removed in + 0.9. Examples -------- >>> import numpy as np >>> from rabitqlib import IvfIndex >>> x = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) - >>> index = IvfIndex(64, len(x), num_clusters=1, nbits=4) - >>> index.build(x, x.mean(axis=0, keepdims=True), np.zeros(len(x), dtype=np.uint32)) - >>> ids, distances = index.search(x[:2], k=3, nprobe=1) + >>> index = IvfIndex(64, 2, nbits=4, seed=7) + >>> index.train(x) + >>> ids = index.add(x) + >>> ids, distances = index.search(x[:2], k=3, nprobe=2) >>> ids.shape, distances.shape ((2, 3), (2, 3)) """ + @overload def __init__( self, dim: SupportsIndex, - max_elements: SupportsIndex, num_clusters: SupportsIndex, + *, nbits: SupportsIndex, metric: _Metric = "l2", initializer: _Initializer = "auto", + seed: SupportsIndex | None = None, ) -> None: """An approximate nearest-neighbor index using IVF. @@ -124,41 +132,50 @@ class IvfIndex: ---------- dim : int Original vector dimension, from 64 through 65536. - max_elements : int - Number of vectors required by the initial build. Later adds grow - storage. num_clusters : int - Number of supplied centroids. + Number of centroids, from train or set_centroids. nbits : int - Total quantization bits, 1 through 9, or 32 for owned float32 vectors - with one-bit filtering and raw reranking. + Keyword-only. Total quantization bits, 1 through 9, or 32 for owned + float32 vectors with one-bit filtering and raw reranking. metric : {'l2', 'ip', 'innerproduct'}, optional Squared L2 or inner-product distance. Default is 'l2'. Vectors are not automatically normalized for cosine similarity. initializer : {'auto', 'flat', 'flat_rabitq', 'hnsw'}, optional Centroid routing strategy. Default is 'auto'. + seed : int, optional + Rotation seed in [0, 2**32 - 1], also the initial cp.seed. None + (default) draws one at random; the seed property reports it. Notes ----- - Build or load the index before searching or saving. Concurrent reads are - supported; an overlapping mutation raises RuntimeError. + Train the index or set its centroids, then add vectors before searching or + saving; loaded indexes are already trained. Concurrent reads are supported; + an overlapping mutation raises RuntimeError. Data, centroids and queries must be finite float32 values within the index coordinate range; other values raise ValueError. - - Examples - -------- - >>> import numpy as np - >>> from rabitqlib import IvfIndex - >>> x = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) - >>> index = IvfIndex(64, len(x), num_clusters=1, nbits=4) - >>> index.build(x, x.mean(axis=0, keepdims=True), np.zeros(len(x), dtype=np.uint32)) - >>> ids, distances = index.search(x[:2], k=3, nprobe=1) - >>> ids.shape, distances.shape - ((2, 3), (2, 3)) + The 0.7 form IvfIndex(dim, max_elements, num_clusters, nbits, ...) and + build() still work in 0.8, raise DeprecationWarning, and will be removed in + 0.9. """ - + @overload + @deprecated( + "IvfIndex(dim, max_elements, num_clusters, nbits, ...) is deprecated since " + "0.8; use IvfIndex(dim, num_clusters, nbits=...)" + ) + def __init__( + self, + dim: SupportsIndex, + max_elements: SupportsIndex, + num_clusters: SupportsIndex, + nbits: SupportsIndex, + metric: _Metric = "l2", + initializer: _Initializer = "auto", + ) -> None: ... def __repr__(self) -> str: - """Return a one-line summary of configuration, counts and build state.""" + """Return a one-line summary of configuration, counts and training state.""" + + def __len__(self) -> int: + """Return the number of live, searchable points.""" @property def dim(self) -> int: @@ -166,7 +183,7 @@ class IvfIndex: @property def max_elements(self) -> int: - """Current storage capacity.""" + """Logical capacity: stored points, or the build row count of a 0.7-form index before build.""" @property def num_points(self) -> int: @@ -177,9 +194,24 @@ class IvfIndex: """Number of stored points, including removed points.""" @property + def is_trained(self) -> bool: + """Whether centroids are set, by train, set_centroids, build or load.""" + + @property + @deprecated("IvfIndex.is_built is deprecated since 0.8; use is_trained or len()") def is_built(self) -> bool: - """Whether the index has been built or loaded.""" + """Deprecated since 0.8: whether the index stores points; use len().""" + @property + def seed(self) -> int | None: + """Rotation seed, or None for an index loaded from disk or created in the 0.7 form.""" + + @property + def cp(self) -> RaBitQKMeansParameters: + """Clustering parameters train uses, a RaBitQKMeansParameters. Defaults: niter 25, max_points_per_centroid 256, spherical for inner product, and seed equal to the index seed.""" + + @cp.setter + def cp(self, value: RaBitQKMeansParameters) -> None: ... @property def metric(self) -> Literal["l2", "ip"]: """Canonical metric name: 'l2' or 'ip'.""" @@ -192,6 +224,49 @@ class IvfIndex: def num_clusters(self) -> int: """Number of centroids used for quantization.""" + def train( + self, x: npt.ArrayLike, *, num_threads: SupportsIndex | None = None + ) -> None: + """Train num_clusters centroids on vectors and set them. + + Parameters + ---------- + x : array_like of float, shape (n, dim) + Training vectors in the original, unpadded coordinates, with + n >= num_clusters. Converted to contiguous float32; a copy may be made. + num_threads : int, optional + Keyword-only. Worker count for this call. None (default) uses + cp.num_threads; 0 selects the available CPU count. + + Notes + ----- + Clusters with the settings in cp: RaBitQKMeans up to 2000 clusters, and + QGKMeans with its default graph settings above that (a provisional, + unmeasured threshold). Only centroids are trained; add assigns vectors. + Retraining replaces the centroids while the index stores no vectors, and + raises RuntimeError afterwards. A failed call leaves the index unchanged. + """ + + def set_centroids(self, centroids: npt.ArrayLike) -> None: + """Use the given centroids instead of training. + + Parameters + ---------- + centroids : array_like of float, shape (num_clusters, dim) + Centroids in the original, unpadded coordinates. Converted to + contiguous float32; a copy may be made. + + Notes + ----- + The index counts as trained afterwards. Replacing centroids is allowed + while the index stores no vectors and raises RuntimeError afterwards. A + failed call leaves the index unchanged. + """ + + @deprecated( + "IvfIndex.build() is deprecated since 0.8; use set_centroids() or " + "train(), then add()" + ) def build( self, data: npt.ArrayLike, @@ -202,14 +277,18 @@ class IvfIndex: ) -> None: """Build the index from vectors and supplied centroid assignments. + Deprecated since 0.8 and removed in 0.9: use set_centroids() (or + train()) and then add(). The defaults match 0.7, including a single + thread. + Parameters ---------- data : array_like of float, shape (n, dim) - Vectors in the original, unpadded coordinates. Converted to contiguous - float32; a copy may be made. + Vectors in the original, unpadded coordinates. Converted to + contiguous float32; a copy may be made. centroids : array_like of float, shape (num_clusters, dim) - Vectors in the original, unpadded coordinates. Converted to contiguous - float32; a copy may be made. + Vectors in the original, unpadded coordinates. Converted to + contiguous float32; a copy may be made. cluster_ids : array_like of int, shape (n,) Centroid ID per vector, in [0, num_clusters). Floating-point and boolean IDs are rejected. @@ -217,22 +296,18 @@ class IvfIndex: Worker count. Default is 1; 0 selects the available CPU count. fast_quantization : bool, optional Use faster quantizer configuration. Default is False. - - Notes - ----- - The initial data must satisfy n == max_elements. Centroids and data use - original coordinates. The input arrays need not be retained after this - synchronous call. """ + @overload def add( self, data: npt.ArrayLike, + *, cluster_ids: _IntegerIds | None = None, - num_threads: SupportsIndex = 1, + num_threads: SupportsIndex = 0, fast_quantization: bool | np.bool_ = False, ) -> npt.NDArray[np.uint32]: - """Append vectors using the existing centroids. + """Add vectors using the trained centroids. Parameters ---------- @@ -240,13 +315,14 @@ class IvfIndex: Vectors in the original, unpadded coordinates. Converted to contiguous float32; a copy may be made. cluster_ids : array_like of int, shape (n,), optional - Existing centroid IDs. None routes each vector using the index's - configured centroid initializer; routing may be approximate. + Keyword-only. Existing centroid IDs. None routes each vector using the + index's configured centroid initializer; routing may be approximate. Centroids do not move. num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. + Keyword-only. Worker count. Default is 0, which selects the available + CPU count. fast_quantization : bool, optional - Use faster quantizer configuration. Default is False. + Keyword-only. Use faster quantizer configuration. Default is False. Returns ------- @@ -255,10 +331,22 @@ class IvfIndex: Notes ----- - Batch additions to reduce repeated cluster copying; reserve can preallocate - room. + Requires a trained or loaded index. Batch additions to reduce repeated + cluster copying; reserve can preallocate room. Passing the options + positionally is deprecated. """ - + @overload + @deprecated( + "Passing options to IvfIndex.add() by position is deprecated since 0.8; " + "pass them by keyword" + ) + def add( + self, + data: npt.ArrayLike, + cluster_ids: _IntegerIds | None, + num_threads: SupportsIndex = ..., + fast_quantization: bool | np.bool_ = ..., + ) -> npt.NDArray[np.uint32]: ... def remove(self, ids: _IntegerIds) -> int: """Exclude point IDs from subsequent search results. @@ -313,13 +401,13 @@ class IvfIndex: Notes ----- - Returns (ids, distances), in that order. Missing neighbors have ID - 4294967295 and distance inf. Array shapes and value ranges are checked at - runtime, not by static type checkers. + Requires a trained or loaded index. Returns (ids, distances), in that + order. Missing neighbors have ID 4294967295 and distance inf. Array shapes + and value ranges are checked at runtime, not by static type checkers. """ def save(self, path: _Path) -> None: - """Save a built index to disk. + """Save an index with stored vectors to disk. Parameters ---------- @@ -329,6 +417,7 @@ class IvfIndex: Notes ----- + A trained index without vectors raises RuntimeError; add vectors first. IVF with HNSW centroid routing also writes its required sidecar file. """ @@ -345,7 +434,7 @@ class IvfIndex: Returns ------- index : IvfIndex - A built index ready for search. + A trained index ready for search. """ @property @@ -364,7 +453,7 @@ class IvfIndex: Notes ----- - Requires a built index. Reserved room is not saved. + Requires a trained or loaded index. Reserved room is not saved. """ class HnswIndex: @@ -374,37 +463,60 @@ class HnswIndex: ---------- dim : int Original vector dimension, from 64 through 65536. - max_elements : int - Storage capacity, at least 1, including removed points. Use resize to grow it. + nbits : int + Keyword-only. Total quantization bits, 1 through 9. + num_clusters : int, optional + Keyword-only. Number of centroids train or set_centroids provides. + Default is 16. M : int, optional - Graph connectivity, at least 2. Default is 16. + Keyword-only. Graph connectivity, at least 2; the base layer keeps up to + 2 * M edges. Default is 16. ef_construction : int, optional - Construction search breadth. Default is 200. - nbits : int, optional - Total quantization bits, 1 through 9. Default is 8. + Keyword-only. Construction search breadth. Default is 200. metric : {'l2', 'ip', 'innerproduct'}, optional - Squared L2 or inner-product distance. Default is 'l2'. Vectors are not - automatically normalized for cosine similarity. - random_seed : int, optional - Graph random seed. Default is 100. + Keyword-only. Squared L2 or inner-product distance. Default is 'l2'. + Vectors are not automatically normalized for cosine similarity. + seed : int, optional + Keyword-only. Graph and rotation seed, also the initial cp.seed (its low + 32 bits). Default is 100. Notes ----- - Build or load the index before searching or saving. Concurrent reads are - supported; an overlapping mutation raises RuntimeError. + Train the index or set its centroids, then add vectors before searching or + saving; loaded indexes are already trained. Capacity grows automatically + on add; resize can preallocate it. Concurrent reads are supported; an + overlapping mutation raises RuntimeError. Data, centroids and queries must be finite float32 values within the index coordinate range; other values raise ValueError. + With more than one thread, add links points in a nondeterministic order; + pass num_threads=1 for a reproducible graph. + The 0.7 forms HnswIndex(dim, max_elements, ...), random_seed= and build() + still work in 0.8, raise DeprecationWarning, and will be removed in 0.9. + + Examples + -------- + >>> import numpy as np + >>> from rabitqlib import HnswIndex + >>> x = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) + >>> index = HnswIndex(64, nbits=4, num_clusters=2) + >>> index.train(x) + >>> ids = index.add(x, num_threads=1) + >>> ids, distances = index.search(x[:2], k=3) + >>> ids.shape, distances.shape + ((2, 3), (2, 3)) """ + @overload def __init__( self, dim: SupportsIndex, - max_elements: SupportsIndex, + *, + nbits: SupportsIndex, + num_clusters: SupportsIndex = 16, M: SupportsIndex = 16, ef_construction: SupportsIndex = 200, - nbits: SupportsIndex = 8, metric: _Metric = "l2", - random_seed: SupportsIndex = 100, + seed: SupportsIndex = 100, ) -> None: """An approximate nearest-neighbor index using HNSW. @@ -412,30 +524,69 @@ class HnswIndex: ---------- dim : int Original vector dimension, from 64 through 65536. - max_elements : int - Storage capacity, at least 1, including removed points. Use resize to grow it. + nbits : int + Keyword-only. Total quantization bits, 1 through 9. + num_clusters : int, optional + Keyword-only. Number of centroids train or set_centroids provides. + Default is 16. M : int, optional - Graph connectivity, at least 2. Default is 16. + Keyword-only. Graph connectivity, at least 2; the base layer keeps up to + 2 * M edges. Default is 16. ef_construction : int, optional - Construction search breadth. Default is 200. - nbits : int, optional - Total quantization bits, 1 through 9. Default is 8. + Keyword-only. Construction search breadth. Default is 200. metric : {'l2', 'ip', 'innerproduct'}, optional - Squared L2 or inner-product distance. Default is 'l2'. Vectors are not - automatically normalized for cosine similarity. - random_seed : int, optional - Graph random seed. Default is 100. + Keyword-only. Squared L2 or inner-product distance. Default is 'l2'. + Vectors are not automatically normalized for cosine similarity. + seed : int, optional + Keyword-only. Graph and rotation seed, also the initial cp.seed (its low + 32 bits). Default is 100. Notes ----- - Build or load the index before searching or saving. Concurrent reads are - supported; an overlapping mutation raises RuntimeError. + Train the index or set its centroids, then add vectors before searching or + saving; loaded indexes are already trained. Capacity grows automatically + on add; resize can preallocate it. Concurrent reads are supported; an + overlapping mutation raises RuntimeError. Data, centroids and queries must be finite float32 values within the index coordinate range; other values raise ValueError. + With more than one thread, add links points in a nondeterministic order; + pass num_threads=1 for a reproducible graph. + The 0.7 forms HnswIndex(dim, max_elements, ...), random_seed= and build() + still work in 0.8, raise DeprecationWarning, and will be removed in 0.9. """ - + @overload + @deprecated( + "HnswIndex(dim, max_elements, ...) is deprecated since 0.8; use " + "HnswIndex(dim, nbits=...), whose capacity grows on add" + ) + def __init__( + self, + dim: SupportsIndex, + max_elements: SupportsIndex, + M: SupportsIndex = 16, + ef_construction: SupportsIndex = 200, + nbits: SupportsIndex = 8, + metric: _Metric = "l2", + random_seed: SupportsIndex = 100, + ) -> None: ... + @overload + @deprecated("HnswIndex(random_seed=...) is deprecated since 0.8; use seed=") + def __init__( + self, + dim: SupportsIndex, + *, + nbits: SupportsIndex, + num_clusters: SupportsIndex = 16, + M: SupportsIndex = 16, + ef_construction: SupportsIndex = 200, + metric: _Metric = "l2", + random_seed: SupportsIndex, + ) -> None: ... def __repr__(self) -> str: - """Return a one-line summary of configuration, counts and build state.""" + """Return a one-line summary of configuration, counts and training state.""" + + def __len__(self) -> int: + """Return the number of live, searchable points.""" @property def dim(self) -> int: @@ -454,9 +605,24 @@ class HnswIndex: """Number of stored points, including removed points.""" @property + def is_trained(self) -> bool: + """Whether centroids are set, by train, set_centroids, build or load.""" + + @property + @deprecated("HnswIndex.is_built is deprecated since 0.8; use is_trained or len()") def is_built(self) -> bool: - """Whether the index has been built or loaded.""" + """Deprecated since 0.8: whether the index stores points; use len().""" + + @property + def seed(self) -> int | None: + """Graph and rotation seed, or None for an index loaded from disk.""" + + @property + def cp(self) -> RaBitQKMeansParameters: + """Clustering parameters train uses, a RaBitQKMeansParameters. Defaults: niter 25, max_points_per_centroid 256, spherical for inner product, and seed equal to the low 32 bits of the index seed.""" + @cp.setter + def cp(self, value: RaBitQKMeansParameters) -> None: ... @property def metric(self) -> Literal["l2", "ip"]: """Canonical metric name: 'l2' or 'ip'.""" @@ -467,8 +633,50 @@ class HnswIndex: @property def num_clusters(self) -> int: - """Number of centroids used for quantization.""" + """Number of centroids: the configured count until centroids are set.""" + def train( + self, x: npt.ArrayLike, *, num_threads: SupportsIndex | None = None + ) -> None: + """Train num_clusters centroids on vectors with RaBitQKMeans and set them. + + Parameters + ---------- + x : array_like of float, shape (n, dim) + Training vectors in the original, unpadded coordinates, with + n >= num_clusters. Converted to contiguous float32; a copy may be made. + num_threads : int, optional + Keyword-only. Worker count for this call. None (default) uses + cp.num_threads; 0 selects the available CPU count. + + Notes + ----- + Clusters with the settings in cp. Only centroids are trained; add assigns + vectors. Retraining replaces the centroids while the index stores no + vectors, and raises RuntimeError afterwards. A failed call leaves the index + unchanged. + """ + + def set_centroids(self, centroids: npt.ArrayLike) -> None: + """Use the given centroids instead of training. + + Parameters + ---------- + centroids : array_like of float, shape (num_clusters, dim) + Centroids in the original, unpadded coordinates. Converted to + contiguous float32; a copy may be made. + + Notes + ----- + The index counts as trained afterwards. Replacing centroids is allowed + while the index stores no vectors and raises RuntimeError afterwards. A + failed call leaves the index unchanged. + """ + + @deprecated( + "HnswIndex.build() is deprecated since 0.8; use set_centroids() or " + "train(), then add()" + ) def build( self, data: npt.ArrayLike, @@ -479,14 +687,18 @@ class HnswIndex: ) -> None: """Build the index from vectors and supplied centroid assignments. + Deprecated since 0.8 and removed in 0.9: use set_centroids() (or + train()) and then add(). The defaults match 0.7, including a single + thread. + Parameters ---------- data : array_like of float, shape (n, dim) - Vectors in the original, unpadded coordinates. Converted to contiguous - float32; a copy may be made. + Vectors in the original, unpadded coordinates. Converted to + contiguous float32; a copy may be made. centroids : array_like of float, shape (num_clusters, dim) - Vectors in the original, unpadded coordinates. Converted to contiguous - float32; a copy may be made. + Vectors in the original, unpadded coordinates. Converted to + contiguous float32; a copy may be made. cluster_ids : array_like of int, shape (n,) Centroid ID per vector, in [0, num_clusters). Floating-point and boolean IDs are rejected. @@ -494,22 +706,18 @@ class HnswIndex: Worker count. Default is 1; 0 selects the available CPU count. fast_quantization : bool, optional Use faster quantizer configuration. Default is False. - - Notes - ----- - The initial data must satisfy 1 <= n <= max_elements. Centroids and data - use original coordinates. The input arrays need not be retained after this - synchronous call. """ + @overload def add( self, data: npt.ArrayLike, + *, cluster_ids: _IntegerIds | None = None, - num_threads: SupportsIndex = 1, + num_threads: SupportsIndex = 0, fast_quantization: bool | np.bool_ = False, ) -> npt.NDArray[np.uint32]: - """Append vectors using the existing centroids. + """Add vectors using the trained centroids. Parameters ---------- @@ -517,12 +725,14 @@ class HnswIndex: Vectors in the original, unpadded coordinates. Converted to contiguous float32; a copy may be made. cluster_ids : array_like of int, shape (n,), optional - Existing centroid IDs. None assigns each vector to its nearest - centroid. Centroids do not move. + Keyword-only. Existing centroid IDs. None assigns each vector to its + nearest centroid. Centroids do not move. num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. + Keyword-only. Worker count. Default is 0, which selects the available + CPU count; more than one thread links nondeterministically, so pass 1 + for a reproducible graph. fast_quantization : bool, optional - Use faster quantizer configuration. Default is False. + Keyword-only. Use faster quantizer configuration. Default is False. Returns ------- @@ -531,10 +741,23 @@ class HnswIndex: Notes ----- - The stored point count plus n must fit max_elements; resize first if - needed. Removed IDs are not reused. + Requires a trained or loaded index. The first add into an empty index links + vectors with their exact coordinates; later adds link through the stored + codes. Capacity grows geometrically when needed. Removed IDs are not + reused. Passing the options positionally is deprecated. """ - + @overload + @deprecated( + "Passing options to HnswIndex.add() by position is deprecated since 0.8; " + "pass them by keyword" + ) + def add( + self, + data: npt.ArrayLike, + cluster_ids: _IntegerIds | None, + num_threads: SupportsIndex = ..., + fast_quantization: bool | np.bool_ = ..., + ) -> npt.NDArray[np.uint32]: ... def remove(self, ids: _IntegerIds) -> int: """Exclude point IDs from subsequent search results. @@ -569,7 +792,7 @@ class HnswIndex: Vectors in the original, unpadded coordinates. Converted to contiguous float32; a copy may be made. k : int - Number of neighbors requested, between 1 and max_elements. + Number of neighbors requested, between 1 and num_stored_points. ef : int, optional Search breadth. 0 (default) selects max(k, 10); values above the stored point count are clamped. @@ -586,13 +809,13 @@ class HnswIndex: Notes ----- - Returns (ids, distances), in that order. Missing neighbors have ID - 4294967295 and distance inf. Array shapes and value ranges are checked at - runtime, not by static type checkers. + Requires a trained or loaded index. Returns (ids, distances), in that + order. Missing neighbors have ID 4294967295 and distance inf. Array shapes + and value ranges are checked at runtime, not by static type checkers. """ def save(self, path: _Path) -> None: - """Save a built index to disk. + """Save an index with stored vectors to disk. Parameters ---------- @@ -602,6 +825,7 @@ class HnswIndex: Notes ----- + A trained index without vectors raises RuntimeError; add vectors first. The saved state includes the rotation used to build the index. """ @@ -618,11 +842,11 @@ class HnswIndex: Returns ------- index : HnswIndex - A built index ready for search. + A trained index ready for search. """ def resize(self, max_elements: SupportsIndex) -> None: - """Resize the storage capacity of a built index. + """Change the storage capacity, for example to preallocate before add. Parameters ---------- diff --git a/python_bindings/bindings_common.hpp b/python_bindings/bindings_common.hpp index 75cfbafa..9c025b1c 100644 --- a/python_bindings/bindings_common.hpp +++ b/python_bindings/bindings_common.hpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -125,6 +126,67 @@ inline void ensure_dimension(const py::array& array, size_t expected, const char } } +// A seed is an integer in [0, max_seed]; bool is rejected like other integer options. +inline uint64_t checked_seed(py::handle value, uint64_t max_seed) { + const std::string range = + "seed must be an integer in [0, " + std::to_string(max_seed) + "]"; + if (PyBool_Check(value.ptr())) { + throw py::type_error(range); + } + auto integer = py::reinterpret_steal(PyNumber_Index(value.ptr())); + if (!integer) { + PyErr_Clear(); + throw py::type_error(range); + } + const unsigned long long seed = PyLong_AsUnsignedLongLong(integer.ptr()); + if (PyErr_Occurred() != nullptr) { + PyErr_Clear(); + throw std::invalid_argument(range); + } + if (seed > max_seed) { + throw std::invalid_argument(range); + } + return seed; +} + +// Call only after rank validation. +inline void check_centroid_shape( + const py::array& centroids, size_t num_clusters, size_t dim +) { + if (static_cast(centroids.shape(0)) != num_clusters || + static_cast(centroids.shape(1)) != dim) { + throw std::invalid_argument( + "centroids must have shape (num_clusters, dim); expected (" + + std::to_string(num_clusters) + ", " + std::to_string(dim) + "), got shape " + + array_shape(centroids) + ); + } +} + +// Training input rows, which must cover num_clusters centroids. Call only after +// rank validation. +inline size_t training_rows(const py::array& x, size_t num_clusters) { + const auto rows = static_cast(x.shape(0)); + if (rows < num_clusters) { + throw std::invalid_argument( + "x must have at least num_clusters rows; expected at least " + + std::to_string(num_clusters) + ", got shape " + array_shape(x) + ); + } + return rows; +} + +// The parameters for one train call: an explicit num_threads overrides cp's. +template +Parameters training_parameters(Parameters cp, std::optional num_threads) { + if (num_threads.has_value()) { + cp.num_threads = static_cast( + std::min(*num_threads, std::numeric_limits::max()) + ); + } + return cp; +} + inline MetricType metric_from_string(const std::string& metric) { if (metric == "l2") { return METRIC_L2; diff --git a/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index d36ff43f..e30fd2bd 100644 --- a/python_bindings/hnsw_bindings.cpp +++ b/python_bindings/hnsw_bindings.cpp @@ -4,18 +4,26 @@ #include #include #include +#include // IWYU pragma: keep; registers std::optional casters #include #include +#include +#include #include +#include #include #include #include #include #include "bindings_common.hpp" +#include "rabitqlib/clustering/rabitqkmeans.hpp" +#include "rabitqlib/index/detail/buffer.hpp" #include "rabitqlib/index/hnsw/hnsw.hpp" +#include "rabitqlib/index/training.hpp" #include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/input_checks.hpp" namespace py = pybind11; @@ -24,26 +32,84 @@ namespace rabitqlib::python_bindings { class HnswIndex { public: HnswIndex( + size_t dim, + size_t nbits, + size_t num_clusters, + size_t M, + size_t ef_construction, + const std::string& metric, + const py::object& seed + ) { + if (num_clusters == 0 || num_clusters > buffer::kSearchBufferMaxPointCount) { + throw std::invalid_argument( + "num_clusters must be in [1, " + + std::to_string(buffer::kSearchBufferMaxPointCount) + "]" + ); + } + const auto metric_type = metric_from_string(metric); + const auto graph_seed = checked_seed(seed, std::numeric_limits::max()); + // Capacity grows on add; resize preallocates. + index_ = std::make_unique( + 1, dim, nbits, M, ef_construction, graph_seed, metric_type + ); + num_clusters_ = num_clusters; + seed_ = graph_seed; + cp_ = default_training_parameters(metric_type, static_cast(graph_seed)); + } + + // The 0.7 form, whose max_elements becomes the initial capacity. + static HnswIndex with_point_count( size_t dim, size_t max_elements, size_t M, size_t ef_construction, size_t nbits, const std::string& metric, - size_t random_seed - ) - : index_( - std::make_unique( - max_elements, - dim, - nbits, - M, - ef_construction, - random_seed, - metric_from_string(metric) - ) - ) {} + const py::object& random_seed + ) { + const auto metric_type = metric_from_string(metric); + const auto graph_seed = + checked_seed(random_seed, std::numeric_limits::max()); + HnswIndex wrapper; + wrapper.index_ = std::make_unique( + max_elements, dim, nbits, M, ef_construction, graph_seed, metric_type + ); + wrapper.num_clusters_ = kDefaultNumClusters; + wrapper.seed_ = graph_seed; + wrapper.cp_ = + default_training_parameters(metric_type, static_cast(graph_seed)); + return wrapper; + } + void train(py::handle x, std::optional num_threads) { + const auto operation = access_.write(); + auto x_array = ensure_2d_array(x, "x"); + check_trainable(); + ensure_dimension( + x_array, index_->dimension(), "x dimension does not match index dim" + ); + const auto rows = training_rows(x_array, num_clusters_); + const auto cp = training_parameters(cp_, num_threads); + const auto* x_ptr = x_array.data(); + const size_t num_clusters = num_clusters_; + py::gil_scoped_release release; + index_->train(num_clusters, rows, x_ptr, cp); + } + + void set_centroids(py::handle centroids) { + const auto operation = access_.write(); + auto centroids_array = ensure_2d_array(centroids, "centroids"); + check_trainable(); + check_centroid_shape(centroids_array, num_clusters_, index_->dimension()); + const auto* centroids_ptr = centroids_array.data(); + const size_t num_clusters = num_clusters_; + py::gil_scoped_release release; + index_->set_centroids(num_clusters, centroids_ptr); + } + + // Deprecated: the 0.7 build, which takes any number of centroids. An untrained + // index is constructed; a trained one without points takes the centroids and + // adds the data, which links it the same way. void build( py::handle data, py::handle centroids, @@ -58,11 +124,9 @@ class HnswIndex { ensure_dimension( data_array, index_->dimension(), "data dimension does not match index dim" ); - if (data_array.shape(0) == 0 || - static_cast(data_array.shape(0)) > index_->max_elements()) { + if (data_array.shape(0) == 0) { throw std::invalid_argument( - "number of data rows must be between 1 and max_elements; expected 1 to " + - std::to_string(index_->max_elements()) + ", got shape " + + "number of data rows must be at least 1; got shape " + array_shape(data_array) ); } @@ -75,29 +139,42 @@ class HnswIndex { if (num_clusters == 0) { throw std::invalid_argument("at least one centroid is required"); } + const auto rows = static_cast(data_array.shape(0)); auto cluster_ids_vec = checked_cluster_ids( - cluster_ids, - static_cast(data_array.shape(0)), - num_clusters, - "cluster_ids contains an out-of-range value" + cluster_ids, rows, num_clusters, "cluster_ids contains an out-of-range value" ); const auto* data_ptr = data_array.data(); const auto* centroids_ptr = centroids_array.data(); - const auto rows = static_cast(data_array.shape(0)); + const size_t dim = index_->dimension(); { py::gil_scoped_release release; - index_->construct( - num_clusters, - centroids_ptr, - rows, - data_ptr, - cluster_ids_vec.data(), - num_threads, - fast_quantization - ); + const bool empty = index_->num_stored_points() == 0; + if (empty) { + // Invalid centroids must not grow or retrain the index first. + check_coordinates( + centroids_ptr, num_clusters * dim, dim, "HNSW construction centroids" + ); + grow_capacity(rows, data_ptr, "HNSW construction data"); + } + if (empty && index_->is_trained()) { + index_->set_centroids(num_clusters, centroids_ptr); + index_->add( + data_ptr, rows, cluster_ids_vec.data(), fast_quantization, num_threads + ); + } else { + index_->construct( + num_clusters, + centroids_ptr, + rows, + data_ptr, + cluster_ids_vec.data(), + num_threads, + fast_quantization + ); + } } - built_ = true; + num_clusters_ = num_clusters; } py::array_t add( @@ -108,9 +185,7 @@ class HnswIndex { ) { const auto operation = access_.write(); auto data_array = ensure_2d_array(data, "data"); - if (!built_) { - throw std::runtime_error("HnswIndex must be built or loaded before add"); - } + check_trained("add"); ensure_dimension( data_array, index_->dimension(), "data dimension does not match index dim" ); @@ -126,6 +201,7 @@ class HnswIndex { const auto* data_ptr = data_array.data(); const auto labels = [&] { py::gil_scoped_release release; + grow_capacity(rows, data_ptr, "HNSW add data"); return index_->add( data_ptr, rows, @@ -140,9 +216,6 @@ class HnswIndex { void resize(size_t new_max_elements) { const auto operation = access_.write(); - if (!built_) { - throw std::runtime_error("HnswIndex must be built or loaded before resize"); - } py::gil_scoped_release release; index_->resize(new_max_elements); } @@ -150,9 +223,7 @@ class HnswIndex { size_t remove(py::handle ids) { const auto operation = access_.write(); auto ids_array = ensure_1d_integer_array(ids, "ids"); - if (!built_) { - throw std::runtime_error("HnswIndex must be built or loaded before remove"); - } + check_trained("remove"); auto to_remove = checked_removal_ids(ids_array, index_->num_stored_points()); py::gil_scoped_release release; return index_->remove(to_remove.data(), to_remove.size()); @@ -160,29 +231,26 @@ class HnswIndex { [[nodiscard]] size_t num_points() const { const auto operation = access_.read(); - return built_ ? index_->num_points() : 0; + return index_->num_points(); } [[nodiscard]] size_t num_stored_points() const { const auto operation = access_.read(); - return built_ ? index_->num_stored_points() : 0; + return index_->num_stored_points(); } py::tuple search(py::handle queries, size_t k, size_t ef, size_t num_threads) { const auto operation = access_.read(); auto query_array = ensure_2d_array(queries, "queries"); - if (!built_) { - throw std::runtime_error("HnswIndex must be built or loaded before search"); - } - if (index_->dimension() != 0) { - ensure_dimension( - query_array, index_->dimension(), "query dimension does not match index dim" - ); - } + check_trained("search"); + ensure_dimension( + query_array, index_->dimension(), "query dimension does not match index dim" + ); if (ef == 0) { ef = std::max(k, 10); } - if (k == 0 || k > index_->max_elements()) { - throw std::invalid_argument("k must be between 1 and max_elements"); + // Capacity grows geometrically, so bound k by what is stored, as IVF does. + if (k == 0 || k > index_->num_stored_points()) { + throw std::invalid_argument("k must be between 1 and num_stored_points"); } auto output = unfilled_search_results(static_cast(query_array.shape(0)), k); @@ -214,8 +282,9 @@ class HnswIndex { void save(py::handle path) const { const auto filename = path_string(path); const auto operation = access_.read(); - if (!built_) { - throw std::runtime_error("HnswIndex must be built or loaded before save"); + check_trained("save"); + if (index_->num_stored_points() == 0) { + throw std::runtime_error("HnswIndex has no points to save; add points first"); } py::gil_scoped_release release; index_->save(filename.c_str()); @@ -229,7 +298,13 @@ class HnswIndex { py::gil_scoped_release release; wrapper.index_->load(filename.c_str()); } - wrapper.built_ = true; + // Untrained, empty files from 0.7 have no clusters; train them with the default. + wrapper.num_clusters_ = wrapper.index_->is_trained() + ? wrapper.index_->num_clusters() + : kDefaultNumClusters; + wrapper.cp_ = default_training_parameters( + wrapper.index_->metric_type(), rabitqkmeans::RaBitQKMeansParameters{}.seed + ); return wrapper; } @@ -238,12 +313,11 @@ class HnswIndex { return "HnswIndex(dim=" + std::to_string(index_->dimension()) + ", metric='" + metric_to_string(index_->metric_type()) + "', nbits=" + std::to_string(index_->nbits()) + - ", num_points=" + std::to_string(built_ ? index_->num_points() : 0) + - ", num_stored_points=" + - std::to_string(built_ ? index_->num_stored_points() : 0) + + ", num_points=" + std::to_string(index_->num_points()) + + ", num_stored_points=" + std::to_string(index_->num_stored_points()) + ", max_elements=" + std::to_string(index_->max_elements()) + - ", num_clusters=" + std::to_string(index_->num_clusters()) + - ", is_built=" + (built_ ? "True" : "False") + ")"; + ", num_clusters=" + std::to_string(current_num_clusters()) + + ", is_trained=" + (index_->is_trained() ? "True" : "False") + ")"; } [[nodiscard]] size_t dim() const { @@ -258,27 +332,77 @@ class HnswIndex { const auto operation = access_.read(); return index_->nbits(); } - [[nodiscard]] bool is_built() const { + [[nodiscard]] bool is_trained() const { const auto operation = access_.read(); - return built_; + return index_->is_trained(); } [[nodiscard]] size_t num_clusters() const { const auto operation = access_.read(); - return index_->num_clusters(); + return current_num_clusters(); + } + [[nodiscard]] std::optional seed() const { + const auto operation = access_.read(); + return seed_; } [[nodiscard]] std::string metric() const { const auto operation = access_.read(); return metric_to_string(index_->metric_type()); } + // Parameters are wrapper state that train copies with the GIL held. + rabitqkmeans::RaBitQKMeansParameters& cp() { return cp_; } + void set_cp(const rabitqkmeans::RaBitQKMeansParameters& cp) { + const auto operation = access_.write(); + cp_ = cp; + } + private: + static constexpr size_t kDefaultNumClusters = 16; + HnswIndex() = default; - // The public constructor creates index_; built_ is set only by a successful - // build or load. - bool built_ = false; + // The configured count until centroids are set, then the trained count. + [[nodiscard]] size_t current_num_clusters() const { + return index_->is_trained() ? index_->num_clusters() : num_clusters_; + } + + void check_trained(const char* operation) const { + if (!index_->is_trained()) { + throw std::runtime_error( + std::string("HnswIndex must be trained or loaded before ") + operation + ); + } + } + + void check_trainable() const { + if (index_->num_stored_points() != 0) { + throw std::runtime_error("HnswIndex cannot be trained once it stores points"); + } + } + + // Grow capacity geometrically to fit `rows` more points from `data`, checked + // first so that invalid input leaves the capacity unchanged. Leaves IDs beyond + // the supported range for add to reject. Called without the GIL. + void grow_capacity(size_t rows, const float* data, const char* what) { + const size_t capacity = index_->max_elements(); + const size_t stored = index_->num_stored_points(); + if (rows <= capacity - stored || + rows > buffer::kSearchBufferMaxPointCount - stored) { + return; + } + const size_t dim = index_->dimension(); + check_coordinates(data, rows * dim, dim, what); + const size_t needed = stored + rows; + const size_t doubled = + std::min(capacity * 2, static_cast(buffer::kSearchBufferMaxPointCount)); + index_->resize(std::max(needed, doubled)); + } + mutable IndexAccess access_; std::unique_ptr index_; + rabitqkmeans::RaBitQKMeansParameters cp_; + size_t num_clusters_ = 0; // Centroid count train and set_centroids expect. + std::optional seed_; // Unknown for loaded indexes. }; } // namespace rabitqlib::python_bindings @@ -294,70 +418,159 @@ Parameters ---------- dim : int Original vector dimension, from 64 through 65536. -max_elements : int - Storage capacity, at least 1, including removed points. Use resize to grow it. +nbits : int + Keyword-only. Total quantization bits, 1 through 9. +num_clusters : int, optional + Keyword-only. Number of centroids train or set_centroids provides. + Default is 16. M : int, optional - Graph connectivity, at least 2. Default is 16. + Keyword-only. Graph connectivity, at least 2; the base layer keeps up to + 2 * M edges. Default is 16. ef_construction : int, optional - Construction search breadth. Default is 200. -nbits : int, optional - Total quantization bits, 1 through 9. Default is 8. + Keyword-only. Construction search breadth. Default is 200. metric : {'l2', 'ip', 'innerproduct'}, optional - Squared L2 or inner-product distance. Default is 'l2'. Vectors are not - automatically normalized for cosine similarity. -random_seed : int, optional - Graph random seed. Default is 100. + Keyword-only. Squared L2 or inner-product distance. Default is 'l2'. + Vectors are not automatically normalized for cosine similarity. +seed : int, optional + Keyword-only. Graph and rotation seed, also the initial cp.seed (its low + 32 bits). Default is 100. Notes ----- -Build or load the index before searching or saving. Concurrent reads are -supported; an overlapping mutation raises RuntimeError. +Train the index or set its centroids, then add vectors before searching or +saving; loaded indexes are already trained. Capacity grows automatically +on add; resize can preallocate it. Concurrent reads are supported; an +overlapping mutation raises RuntimeError. Data, centroids and queries must be finite float32 values within the index -coordinate range; other values raise ValueError.)doc"; +coordinate range; other values raise ValueError. +With more than one thread, add links points in a nondeterministic order; +pass num_threads=1 for a reproducible graph. +The 0.7 forms HnswIndex(dim, max_elements, ...), random_seed= and build() +still work in 0.8, raise DeprecationWarning, and will be removed in 0.9. + +Examples +-------- +>>> import numpy as np +>>> from rabitqlib import HnswIndex +>>> x = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) +>>> index = HnswIndex(64, nbits=4, num_clusters=2) +>>> index.train(x) +>>> ids = index.add(x, num_threads=1) +>>> ids, distances = index.search(x[:2], k=3) +>>> ids.shape, distances.shape +((2, 3), (2, 3)))doc"; + + static constexpr char kAddDoc[] = R"doc(Add vectors using the trained centroids. + +Parameters +---------- +data : array_like of float, shape (n, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. +cluster_ids : array_like of int, shape (n,), optional + Keyword-only. Existing centroid IDs. None assigns each vector to its + nearest centroid. Centroids do not move. +num_threads : int, optional + Keyword-only. Worker count. Default is 0, which selects the available + CPU count; more than one thread links nondeterministically, so pass 1 + for a reproducible graph. +fast_quantization : bool, optional + Keyword-only. Use faster quantizer configuration. Default is False. + +Returns +------- +ids : ndarray of uint32, shape (n,) + Assigned point IDs. + +Notes +----- +Requires a trained or loaded index. The first add into an empty index links +vectors with their exact coordinates; later adds link through the stored +codes. Capacity grows geometrically when needed. Removed IDs are not +reused. Passing the options positionally is deprecated.)doc"; py::class_(m, "HnswIndex", kIndexDoc) .def( - py::init(), + py::init< + size_t, + size_t, + size_t, + size_t, + size_t, + const std::string&, + const py::object&>(), py::arg("dim"), - py::arg("max_elements"), + py::kw_only(), + py::arg("nbits"), + py::arg("num_clusters") = 16, py::arg("M") = 16, py::arg("ef_construction") = 200, - py::arg("nbits") = 8, py::arg("metric") = "l2", - py::arg("random_seed") = 100, + py::arg("seed") = 100, kIndexDoc ) + // Private to the 0.8 compatibility layer in rabitqlib/__init__.py. .def( - "build", - &HnswIndex::build, - py::arg("data"), + py::init(&HnswIndex::with_point_count), + py::arg("dim"), + py::kw_only(), + py::arg("_max_elements"), + py::arg("M"), + py::arg("ef_construction"), + py::arg("nbits"), + py::arg("metric"), + py::arg("random_seed") + ) + .def( + "train", + &HnswIndex::train, + py::arg("x"), + py::kw_only(), + py::arg("num_threads") = py::none(), + R"doc(Train num_clusters centroids on vectors with RaBitQKMeans and set them. + +Parameters +---------- +x : array_like of float, shape (n, dim) + Training vectors in the original, unpadded coordinates, with + n >= num_clusters. Converted to contiguous float32; a copy may be made. +num_threads : int, optional + Keyword-only. Worker count for this call. None (default) uses + cp.num_threads; 0 selects the available CPU count. + +Notes +----- +Clusters with the settings in cp. Only centroids are trained; add assigns +vectors. Retraining replaces the centroids while the index stores no +vectors, and raises RuntimeError afterwards. A failed call leaves the index +unchanged.)doc" + ) + .def( + "set_centroids", + &HnswIndex::set_centroids, py::arg("centroids"), - py::arg("cluster_ids"), - py::arg("num_threads") = 1, - py::arg("fast_quantization") = false, - R"doc(Build the index from vectors and supplied centroid assignments. + R"doc(Use the given centroids instead of training. Parameters ---------- -data : array_like of float, shape (n, dim) - Vectors in the original, unpadded coordinates. Converted to contiguous - float32; a copy may be made. centroids : array_like of float, shape (num_clusters, dim) - Vectors in the original, unpadded coordinates. Converted to contiguous - float32; a copy may be made. -cluster_ids : array_like of int, shape (n,) - Centroid ID per vector, in [0, num_clusters). Floating-point and - boolean IDs are rejected. -num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. -fast_quantization : bool, optional - Use faster quantizer configuration. Default is False. + Centroids in the original, unpadded coordinates. Converted to + contiguous float32; a copy may be made. Notes ----- -The initial data must satisfy 1 <= n <= max_elements. Centroids and data -use original coordinates. The input arrays need not be retained after this -synchronous call.)doc" +The index counts as trained afterwards. Replacing centroids is allowed +while the index stores no vectors and raises RuntimeError afterwards. A +failed call leaves the index unchanged.)doc" + ) + .def( + "_build", + &HnswIndex::build, + py::arg("data"), + py::arg("centroids"), + py::arg("cluster_ids"), + py::arg("num_threads"), + py::arg("fast_quantization") ) .def( "search", @@ -374,7 +587,7 @@ queries : array_like of float, shape (n_queries, dim) Vectors in the original, unpadded coordinates. Converted to contiguous float32; a copy may be made. k : int - Number of neighbors requested, between 1 and max_elements. + Number of neighbors requested, between 1 and num_stored_points. ef : int, optional Search breadth. 0 (default) selects max(k, 10); values above the stored point count are clamped. @@ -391,47 +604,25 @@ distances : ndarray of float32, shape (n_queries, k) Notes ----- -Returns (ids, distances), in that order. Missing neighbors have ID -4294967295 and distance inf. Array shapes and value ranges are checked at -runtime, not by static type checkers.)doc" +Requires a trained or loaded index. Returns (ids, distances), in that +order. Missing neighbors have ID 4294967295 and distance inf. Array shapes +and value ranges are checked at runtime, not by static type checkers.)doc" ) .def( "add", &HnswIndex::add, py::arg("data"), + py::kw_only(), py::arg("cluster_ids") = py::none(), - py::arg("num_threads") = 1, + py::arg("num_threads") = 0, py::arg("fast_quantization") = false, - R"doc(Append vectors using the existing centroids. - -Parameters ----------- -data : array_like of float, shape (n, dim) - Vectors in the original, unpadded coordinates. Converted to contiguous - float32; a copy may be made. -cluster_ids : array_like of int, shape (n,), optional - Existing centroid IDs. None assigns each vector to its nearest - centroid. Centroids do not move. -num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. -fast_quantization : bool, optional - Use faster quantizer configuration. Default is False. - -Returns -------- -ids : ndarray of uint32, shape (n,) - Assigned point IDs. - -Notes ------ -The stored point count plus n must fit max_elements; resize first if -needed. Removed IDs are not reused.)doc" + kAddDoc ) .def( "resize", &HnswIndex::resize, py::arg("max_elements"), - R"doc(Resize the storage capacity of a built index. + R"doc(Change the storage capacity, for example to preallocate before add. Parameters ---------- @@ -463,7 +654,7 @@ Removal does not reclaim storage and cannot be undone.)doc" "save", &HnswIndex::save, py::arg("path"), - R"doc(Save a built index to disk. + R"doc(Save an index with stored vectors to disk. Parameters ---------- @@ -473,6 +664,7 @@ path : str, bytes, bytearray, or os.PathLike Notes ----- +A trained index without vectors raises RuntimeError; add vectors first. The saved state includes the rotation used to build the index.)doc" ) .def_static( @@ -490,12 +682,17 @@ path : str, bytes, bytearray, or os.PathLike Returns ------- index : HnswIndex - A built index ready for search.)doc" + A trained index ready for search.)doc" ) .def( "__repr__", &HnswIndex::repr, - "Return a one-line summary of configuration, counts and build state." + "Return a one-line summary of configuration, counts and training state." + ) + .def( + "__len__", + &HnswIndex::num_points, + "Return the number of live, searchable points." ) .def_property_readonly( "dim", &HnswIndex::dim, R"doc(Original vector dimension.)doc" @@ -519,12 +716,24 @@ index : HnswIndex .def_property_readonly( "num_clusters", &HnswIndex::num_clusters, - R"doc(Number of centroids used for quantization.)doc" + R"doc(Number of centroids: the configured count until centroids are set.)doc" + ) + .def_property_readonly( + "is_trained", + &HnswIndex::is_trained, + R"doc(Whether centroids are set, by train, set_centroids, build or load.)doc" ) .def_property_readonly( - "is_built", - &HnswIndex::is_built, - R"doc(Whether the index has been built or loaded.)doc" + "seed", + &HnswIndex::seed, + R"doc(Graph and rotation seed, or None for an index loaded from disk.)doc" + ) + .def_property( + "cp", + &HnswIndex::cp, + &HnswIndex::set_cp, + py::return_value_policy::reference_internal, + R"doc(Clustering parameters train uses, a RaBitQKMeansParameters. Defaults: niter 25, max_points_per_centroid 256, spherical for inner product, and seed equal to the low 32 bits of the index seed.)doc" ) .def_property_readonly( "metric", &HnswIndex::metric, R"doc(Canonical metric name: 'l2' or 'ip'.)doc" diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index b2bf292d..aa04a1d9 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/python_bindings/ivf_bindings.cpp @@ -8,17 +8,22 @@ #include #include +#include #include #include +#include #include #include #include #include "bindings_common.hpp" +#include "rabitqlib/clustering/rabitqkmeans.hpp" #include "rabitqlib/core/rotator.hpp" #include "rabitqlib/index/ivf/initializer.hpp" #include "rabitqlib/index/ivf/ivf.hpp" +#include "rabitqlib/index/training.hpp" #include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/input_checks.hpp" namespace py = pybind11; @@ -61,6 +66,35 @@ std::string initializer_to_string(ivf::InitializerType type) { class IvfIndex { public: IvfIndex( + size_t dim, + size_t num_clusters, + size_t nbits, + const std::string& metric, + const std::string& initializer, + const py::object& seed + ) { + const auto metric_type = metric_from_string(metric); + const auto initializer_type = initializer_from_string(initializer); + const auto rotation_seed = + seed.is_none() ? std::random_device{}() + : static_cast( + checked_seed(seed, std::numeric_limits::max()) + ); + index_ = std::make_unique( + dim, + num_clusters, + nbits, + metric_type, + rotation_seed, + RotatorType::FhtKacRotator, + initializer_type + ); + seed_ = rotation_seed; + cp_ = default_training_parameters(metric_type, rotation_seed); + } + + // The 0.7 point-count form, which build fills with construct. + static IvfIndex with_point_count( size_t dim, size_t max_elements, size_t num_clusters, @@ -69,7 +103,8 @@ class IvfIndex { const std::string& initializer ) { const auto metric_type = metric_from_string(metric); - index_ = std::make_unique( + IvfIndex wrapper; + wrapper.index_ = std::make_unique( max_elements, dim, num_clusters, @@ -78,8 +113,37 @@ class IvfIndex { RotatorType::FhtKacRotator, initializer_from_string(initializer) ); + wrapper.cp_ = untrainable_parameters(metric_type); + wrapper.builds_by_construct_ = true; + return wrapper; } + void train(py::handle x, std::optional num_threads) { + const auto operation = access_.write(); + auto x_array = ensure_2d_array(x, "x"); + check_trainable(); + ensure_dimension( + x_array, index_->dimension(), "x dimension does not match index dim" + ); + const auto rows = training_rows(x_array, index_->num_clusters()); + const auto cp = training_parameters(cp_, num_threads); + const auto* x_ptr = x_array.data(); + py::gil_scoped_release release; + index_->train(rows, x_ptr, cp); + } + + void set_centroids(py::handle centroids) { + const auto operation = access_.write(); + auto centroids_array = ensure_2d_array(centroids, "centroids"); + check_trainable(); + check_centroid_shape(centroids_array, index_->num_clusters(), index_->dimension()); + const auto* centroids_ptr = centroids_array.data(); + py::gil_scoped_release release; + index_->set_centroids(centroids_ptr); + } + + // Deprecated: the 0.7 build. A point-count index is constructed; any other + // index sets the centroids and adds the data. void build( py::handle data, py::handle centroids, @@ -94,7 +158,8 @@ class IvfIndex { ensure_dimension( data_array, index_->dimension(), "data dimension does not match index dim" ); - if (static_cast(data_array.shape(0)) != index_->max_elements()) { + if (builds_by_construct_ && + static_cast(data_array.shape(0)) != index_->max_elements()) { throw std::invalid_argument( "number of data rows must match index max_elements; expected " + std::to_string(index_->max_elements()) + ", got shape " + @@ -113,17 +178,29 @@ class IvfIndex { array_shape(centroids_array) ); } - auto cluster_ids_vec = checked_cluster_ids( - cluster_ids, static_cast(data_array.shape(0)), index_->num_clusters() - ); + const auto rows = static_cast(data_array.shape(0)); + auto cluster_ids_vec = + checked_cluster_ids(cluster_ids, rows, index_->num_clusters()); + if (!builds_by_construct_) { + check_trainable(); + } const auto* data_ptr = data_array.data(); const auto* centroids_ptr = centroids_array.data(); const auto* cluster_ids_ptr = cluster_ids_vec.data(); py::gil_scoped_release release; - index_->construct( - data_ptr, centroids_ptr, cluster_ids_ptr, fast_quantization, num_threads - ); + if (builds_by_construct_) { + index_->construct( + data_ptr, centroids_ptr, cluster_ids_ptr, fast_quantization, num_threads + ); + } else { + // Check the data before replacing the centroids, so that invalid input + // leaves the index unchanged; add repeats the check. + const size_t dim = index_->dimension(); + check_coordinates(data_ptr, rows * dim, dim, "IVF add data"); + index_->set_centroids(centroids_ptr, num_threads); + index_->add(data_ptr, rows, cluster_ids_ptr, fast_quantization, num_threads); + } } py::array_t add( @@ -134,9 +211,7 @@ class IvfIndex { ) { const auto operation = access_.write(); auto data_array = ensure_2d_array(data, "data"); - if (!built()) { - throw std::runtime_error("IvfIndex must be built or loaded before add"); - } + check_trained("add"); ensure_dimension( data_array, index_->dimension(), "data dimension does not match index dim" ); @@ -171,9 +246,7 @@ class IvfIndex { void reserve(int64_t cluster_id, int64_t num_points) { const auto operation = access_.write(); - if (!built()) { - throw std::runtime_error("IvfIndex must be built or loaded before reserve"); - } + check_trained("reserve"); if (cluster_id < 0 || static_cast(cluster_id) >= index_->num_clusters()) { throw std::invalid_argument("cluster_id must be in [0, num_clusters)"); } @@ -187,9 +260,7 @@ class IvfIndex { size_t remove(py::handle ids) { const auto operation = access_.write(); auto ids_array = ensure_1d_integer_array(ids, "ids"); - if (!built()) { - throw std::runtime_error("IvfIndex must be built or loaded before remove"); - } + check_trained("remove"); auto to_remove = checked_removal_ids(ids_array, index_->num_stored_points()); py::gil_scoped_release release; return index_->remove(to_remove.data(), to_remove.size()); @@ -204,9 +275,7 @@ class IvfIndex { ) { const auto operation = access_.read(); auto query_array = ensure_2d_array(queries, "queries"); - if (!built()) { - throw std::runtime_error("IvfIndex must be built or loaded before search"); - } + check_trained("search"); ensure_dimension( query_array, index_->dimension(), "query dimension does not match index dim" ); @@ -243,8 +312,9 @@ class IvfIndex { void save(py::handle path) const { const auto filename = path_string(path); const auto operation = access_.read(); - if (!built()) { - throw std::runtime_error("IvfIndex must be built or loaded before save"); + check_trained("save"); + if (index_->num_stored_points() == 0) { + throw std::runtime_error("IvfIndex has no points to save; add points first"); } py::gil_scoped_release release; index_->save(filename.c_str()); @@ -258,6 +328,8 @@ class IvfIndex { py::gil_scoped_release release; wrapper.index_->load(filename.c_str()); } + wrapper.cp_ = untrainable_parameters(wrapper.index_->metric_type()); + wrapper.builds_by_construct_ = true; return wrapper; } @@ -271,7 +343,7 @@ class IvfIndex { ", max_elements=" + std::to_string(index_->max_elements()) + ", num_clusters=" + std::to_string(index_->num_clusters()) + ", initializer='" + initializer_to_string(index_->initializer_type()) + - "', is_built=" + (built() ? "True" : "False") + ")"; + "', is_trained=" + (index_->is_trained() ? "True" : "False") + ")"; } [[nodiscard]] size_t dim() const { @@ -298,9 +370,13 @@ class IvfIndex { const auto operation = access_.read(); return index_->nbits(); } - [[nodiscard]] bool is_built() const { + [[nodiscard]] bool is_trained() const { + const auto operation = access_.read(); + return index_->is_trained(); + } + [[nodiscard]] std::optional seed() const { const auto operation = access_.read(); - return built(); + return seed_; } [[nodiscard]] std::string initializer() const { const auto operation = access_.read(); @@ -311,15 +387,49 @@ class IvfIndex { return metric_to_string(index_->metric_type()); } + // Parameters are wrapper state that train copies with the GIL held. + rabitqkmeans::RaBitQKMeansParameters& cp() { return cp_; } + void set_cp(const rabitqkmeans::RaBitQKMeansParameters& cp) { + const auto operation = access_.write(); + cp_ = cp; + } + private: IvfIndex() = default; - // Construction or loading succeeded. A failed construct keeps the old index, or - // leaves it empty if allocation had begun. Removed points remain stored. - [[nodiscard]] bool built() const { return index_->num_stored_points() != 0; } + // Parameters of an index that cannot be trained; the k-means default seed. + static rabitqkmeans::RaBitQKMeansParameters untrainable_parameters(MetricType metric) { + return default_training_parameters( + metric, rabitqkmeans::RaBitQKMeansParameters{}.seed + ); + } + + void check_trained(const char* operation) const { + if (!index_->is_trained()) { + throw std::runtime_error( + std::string("IvfIndex must be trained or loaded before ") + operation + ); + } + } + + void check_trainable() const { + if (index_->num_stored_points() != 0) { + throw std::runtime_error("IvfIndex cannot be trained once it stores points"); + } + if (builds_by_construct_) { + throw std::runtime_error( + "IvfIndex created with max_elements is filled by build; create " + "IvfIndex(dim, num_clusters, nbits=...) to train or set centroids" + ); + } + } mutable IndexAccess access_; std::unique_ptr index_; + rabitqkmeans::RaBitQKMeansParameters cp_; + std::optional seed_; // Unknown for point-count and loaded indexes. + // Point-count and loaded indexes build by construct, never by training. + bool builds_by_construct_ = false; }; } // namespace rabitqlib::python_bindings @@ -335,119 +445,170 @@ Parameters ---------- dim : int Original vector dimension, from 64 through 65536. -max_elements : int - Number of vectors required by the initial build. Later adds grow - storage. num_clusters : int - Number of supplied centroids. + Number of centroids, from train or set_centroids. nbits : int - Total quantization bits, 1 through 9, or 32 for owned float32 vectors - with one-bit filtering and raw reranking. + Keyword-only. Total quantization bits, 1 through 9, or 32 for owned + float32 vectors with one-bit filtering and raw reranking. metric : {'l2', 'ip', 'innerproduct'}, optional Squared L2 or inner-product distance. Default is 'l2'. Vectors are not automatically normalized for cosine similarity. initializer : {'auto', 'flat', 'flat_rabitq', 'hnsw'}, optional Centroid routing strategy. Default is 'auto'. +seed : int, optional + Rotation seed in [0, 2**32 - 1], also the initial cp.seed. None + (default) draws one at random; the seed property reports it. Notes ----- -Build or load the index before searching or saving. Concurrent reads are -supported; an overlapping mutation raises RuntimeError. +Train the index or set its centroids, then add vectors before searching or +saving; loaded indexes are already trained. Concurrent reads are supported; +an overlapping mutation raises RuntimeError. Data, centroids and queries must be finite float32 values within the index coordinate range; other values raise ValueError. +The 0.7 form IvfIndex(dim, max_elements, num_clusters, nbits, ...) and +build() still work in 0.8, raise DeprecationWarning, and will be removed in +0.9. Examples -------- >>> import numpy as np >>> from rabitqlib import IvfIndex >>> x = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) ->>> index = IvfIndex(64, len(x), num_clusters=1, nbits=4) ->>> index.build(x, x.mean(axis=0, keepdims=True), np.zeros(len(x), dtype=np.uint32)) ->>> ids, distances = index.search(x[:2], k=3, nprobe=1) +>>> index = IvfIndex(64, 2, nbits=4, seed=7) +>>> index.train(x) +>>> ids = index.add(x) +>>> ids, distances = index.search(x[:2], k=3, nprobe=2) >>> ids.shape, distances.shape ((2, 3), (2, 3)))doc"; + static constexpr char kAddDoc[] = R"doc(Add vectors using the trained centroids. + +Parameters +---------- +data : array_like of float, shape (n, dim) + Vectors in the original, unpadded coordinates. Converted to contiguous + float32; a copy may be made. +cluster_ids : array_like of int, shape (n,), optional + Keyword-only. Existing centroid IDs. None routes each vector using the + index's configured centroid initializer; routing may be approximate. + Centroids do not move. +num_threads : int, optional + Keyword-only. Worker count. Default is 0, which selects the available + CPU count. +fast_quantization : bool, optional + Keyword-only. Use faster quantizer configuration. Default is False. + +Returns +------- +ids : ndarray of uint32, shape (n,) + Assigned point IDs. + +Notes +----- +Requires a trained or loaded index. Batch additions to reduce repeated +cluster copying; reserve can preallocate room. Passing the options +positionally is deprecated.)doc"; + py::class_(m, "IvfIndex", kIndexDoc) .def( py::init< size_t, size_t, size_t, - size_t, const std::string&, - const std::string&>(), + const std::string&, + const py::object&>(), py::arg("dim"), - py::arg("max_elements"), py::arg("num_clusters"), + py::kw_only(), py::arg("nbits"), py::arg("metric") = "l2", py::arg("initializer") = "auto", + py::arg("seed") = py::none(), kIndexDoc ) + // Private to the 0.8 compatibility layer in rabitqlib/__init__.py. .def( - "build", - &IvfIndex::build, - py::arg("data"), - py::arg("centroids"), - py::arg("cluster_ids"), - py::arg("num_threads") = 1, - py::arg("fast_quantization") = false, - R"doc(Build the index from vectors and supplied centroid assignments. + py::init([](size_t dim, + size_t num_clusters, + size_t nbits, + const std::string& metric, + const std::string& initializer, + size_t max_elements) { + return IvfIndex::with_point_count( + dim, max_elements, num_clusters, nbits, metric, initializer + ); + }), + py::arg("dim"), + py::arg("num_clusters"), + py::kw_only(), + py::arg("nbits"), + py::arg("metric"), + py::arg("initializer"), + py::arg("_max_elements") + ) + .def( + "train", + &IvfIndex::train, + py::arg("x"), + py::kw_only(), + py::arg("num_threads") = py::none(), + R"doc(Train num_clusters centroids on vectors and set them. Parameters ---------- -data : array_like of float, shape (n, dim) - Vectors in the original, unpadded coordinates. Converted to contiguous - float32; a copy may be made. -centroids : array_like of float, shape (num_clusters, dim) - Vectors in the original, unpadded coordinates. Converted to contiguous - float32; a copy may be made. -cluster_ids : array_like of int, shape (n,) - Centroid ID per vector, in [0, num_clusters). Floating-point and - boolean IDs are rejected. +x : array_like of float, shape (n, dim) + Training vectors in the original, unpadded coordinates, with + n >= num_clusters. Converted to contiguous float32; a copy may be made. num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. -fast_quantization : bool, optional - Use faster quantizer configuration. Default is False. + Keyword-only. Worker count for this call. None (default) uses + cp.num_threads; 0 selects the available CPU count. Notes ----- -The initial data must satisfy n == max_elements. Centroids and data use -original coordinates. The input arrays need not be retained after this -synchronous call.)doc" +Clusters with the settings in cp: RaBitQKMeans up to 2000 clusters, and +QGKMeans with its default graph settings above that (a provisional, +unmeasured threshold). Only centroids are trained; add assigns vectors. +Retraining replaces the centroids while the index stores no vectors, and +raises RuntimeError afterwards. A failed call leaves the index unchanged.)doc" ) .def( - "add", - &IvfIndex::add, - py::arg("data"), - py::arg("cluster_ids") = py::none(), - py::arg("num_threads") = 1, - py::arg("fast_quantization") = false, - R"doc(Append vectors using the existing centroids. + "set_centroids", + &IvfIndex::set_centroids, + py::arg("centroids"), + R"doc(Use the given centroids instead of training. Parameters ---------- -data : array_like of float, shape (n, dim) - Vectors in the original, unpadded coordinates. Converted to contiguous - float32; a copy may be made. -cluster_ids : array_like of int, shape (n,), optional - Existing centroid IDs. None routes each vector using the index's - configured centroid initializer; routing may be approximate. - Centroids do not move. -num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. -fast_quantization : bool, optional - Use faster quantizer configuration. Default is False. - -Returns -------- -ids : ndarray of uint32, shape (n,) - Assigned point IDs. +centroids : array_like of float, shape (num_clusters, dim) + Centroids in the original, unpadded coordinates. Converted to + contiguous float32; a copy may be made. Notes ----- -Batch additions to reduce repeated cluster copying; reserve can preallocate -room.)doc" +The index counts as trained afterwards. Replacing centroids is allowed +while the index stores no vectors and raises RuntimeError afterwards. A +failed call leaves the index unchanged.)doc" + ) + .def( + "_build", + &IvfIndex::build, + py::arg("data"), + py::arg("centroids"), + py::arg("cluster_ids"), + py::arg("num_threads"), + py::arg("fast_quantization") + ) + .def( + "add", + &IvfIndex::add, + py::arg("data"), + py::kw_only(), + py::arg("cluster_ids") = py::none(), + py::arg("num_threads") = 0, + py::arg("fast_quantization") = false, + kAddDoc ) .def( "reserve", @@ -465,7 +626,7 @@ num_points : int Notes ----- -Requires a built index. Reserved room is not saved.)doc" +Requires a trained or loaded index. Reserved room is not saved.)doc" ) .def( "remove", @@ -523,15 +684,15 @@ distances : ndarray of float32, shape (n_queries, k) Notes ----- -Returns (ids, distances), in that order. Missing neighbors have ID -4294967295 and distance inf. Array shapes and value ranges are checked at -runtime, not by static type checkers.)doc" +Requires a trained or loaded index. Returns (ids, distances), in that +order. Missing neighbors have ID 4294967295 and distance inf. Array shapes +and value ranges are checked at runtime, not by static type checkers.)doc" ) .def( "save", &IvfIndex::save, py::arg("path"), - R"doc(Save a built index to disk. + R"doc(Save an index with stored vectors to disk. Parameters ---------- @@ -541,6 +702,7 @@ path : str, bytes, bytearray, or os.PathLike Notes ----- +A trained index without vectors raises RuntimeError; add vectors first. IVF with HNSW centroid routing also writes its required sidecar file.)doc" ) .def_static( @@ -558,16 +720,23 @@ path : str, bytes, bytearray, or os.PathLike Returns ------- index : IvfIndex - A built index ready for search.)doc" + A trained index ready for search.)doc" ) .def( "__repr__", &IvfIndex::repr, - "Return a one-line summary of configuration, counts and build state." + "Return a one-line summary of configuration, counts and training state." + ) + .def( + "__len__", + &IvfIndex::num_points, + "Return the number of live, searchable points." ) .def_property_readonly("dim", &IvfIndex::dim, R"doc(Original vector dimension.)doc") .def_property_readonly( - "max_elements", &IvfIndex::max_elements, R"doc(Current storage capacity.)doc" + "max_elements", + &IvfIndex::max_elements, + R"doc(Logical capacity: stored points, or the build row count of a 0.7-form index before build.)doc" ) .def_property_readonly( "num_points", @@ -588,9 +757,21 @@ index : IvfIndex "nbits", &IvfIndex::nbits, R"doc(Total storage bits per coordinate.)doc" ) .def_property_readonly( - "is_built", - &IvfIndex::is_built, - R"doc(Whether the index has been built or loaded.)doc" + "is_trained", + &IvfIndex::is_trained, + R"doc(Whether centroids are set, by train, set_centroids, build or load.)doc" + ) + .def_property_readonly( + "seed", + &IvfIndex::seed, + R"doc(Rotation seed, or None for an index loaded from disk or created in the 0.7 form.)doc" + ) + .def_property( + "cp", + &IvfIndex::cp, + &IvfIndex::set_cp, + py::return_value_policy::reference_internal, + R"doc(Clustering parameters train uses, a RaBitQKMeansParameters. Defaults: niter 25, max_points_per_centroid 256, spherical for inner product, and seed equal to the index seed.)doc" ) .def_property_readonly( "metric", &IvfIndex::metric, R"doc(Canonical metric name: 'l2' or 'ip'.)doc" diff --git a/sample/python/README.md b/sample/python/README.md index 0112f0ac..a7a523f4 100644 --- a/sample/python/README.md +++ b/sample/python/README.md @@ -6,13 +6,24 @@ with `rabitqlib>=0.7.0` and NumPy installed. When working from source, so the examples and native extension use matching APIs. Clustering and indexing both use RaBitQ Library; FAISS is needed only for `compare_with_faiss.py`. -## IVF-RaBitQ: cluster, then index +## IVF-RaBitQ -`kmeans_clustering.py` uses `--method qg` (the default) for QGKMeans graph +The indexing script trains the centroids itself unless you pass a clustering +file: + +```sh +python sample/python/ivf_rabitq_indexing.py --total-bits 5 --num-clusters 4096 \ + data/gist/gist_base.fvecs data/gist/ivf_4096_5.index + +python sample/python/ivf_rabitq_querying.py \ + data/gist/ivf_4096_5.index data/gist/gist_query.fvecs data/gist/gist_groundtruth.ivecs +``` + +To reuse one clustering across index configurations, cluster once and pass the +file with `--clusters`. `kmeans_clustering.py` uses `--method qg` (the default) for QGKMeans graph assignment or `--method rabitq` for RaBitQKMeans flat assignment. Both make exact final assignments. The saved `.npz` contains float32 `centroids`, uint32 -`cluster_ids`, and `metric` (`l2` or `ip`), and can be reused across index -configurations. Choose the method explicitly; the script does not switch based +`cluster_ids`, and `metric` (`l2` or `ip`). Choose the method explicitly; the script does not switch based on the cluster count. ```sh @@ -22,9 +33,6 @@ python sample/python/kmeans_clustering.py --method qg --num-clusters 4096 \ python sample/python/ivf_rabitq_indexing.py --total-bits 5 \ --clusters data/gist/clusters_4096_l2.npz \ data/gist/gist_base.fvecs data/gist/ivf_4096_5.index - -python sample/python/ivf_rabitq_querying.py \ - data/gist/ivf_4096_5.index data/gist/gist_query.fvecs data/gist/gist_groundtruth.ivecs ``` Use the same base vectors in the same row order for clustering and indexing. @@ -33,26 +41,26 @@ detect replaced or reordered vectors. See [QGKMeans](../../docs/docs/clustering. and [RaBitQKMeans](../../docs/docs/clustering.md#rabitqkmeans) for input limits and assignment behavior. -For inner product, pass `--metric ip` to both stages and provide normalized -training vectors. Use queries and ground truth prepared for the same metric. +For inner product, pass `--metric ip` to every stage and provide normalized +vectors. Use queries and ground truth prepared for the same metric. `--num-threads 0` uses the library's available logical CPU count; each stage accepts its own limit. See [threading](../../docs/docs/index/usage.md#threading-and-file-paths) for affinity and OpenMP placement behavior. ## HNSW-RaBitQ -Use [RaBitQKMeans](../../docs/docs/clustering.md#rabitqkmeans) for this small-cluster -example, with the same clustering file format: +The indexing script trains 16 centroids by default (`--num-clusters`): ```sh -python sample/python/kmeans_clustering.py --method rabitq --num-clusters 16 \ - data/gist/gist_base.fvecs data/gist/clusters_16_l2.npz - python sample/python/hnsw_rabitq_indexing.py --total-bits 5 --degree 16 \ - --ef-construction 200 --clusters data/gist/clusters_16_l2.npz \ - data/gist/gist_base.fvecs data/gist/hnsw_5.index + --ef-construction 200 data/gist/gist_base.fvecs data/gist/hnsw_5.index ``` +It also accepts `--clusters` with a file from `kmeans_clustering.py`; use +`--method rabitq` there for a small cluster count. With more than one thread +(`--num-threads`), the graph links vectors in a nondeterministic order; pass +`--num-threads 1` for a reproducible index. + Query with `hnsw_rabitq_querying.py`. SymphonyQG builds directly from vectors; use `symqg_indexing.py` and `symqg_querying.py` without a clustering stage. diff --git a/sample/python/hnsw_rabitq_indexing.py b/sample/python/hnsw_rabitq_indexing.py index b0ba734a..208801be 100644 --- a/sample/python/hnsw_rabitq_indexing.py +++ b/sample/python/hnsw_rabitq_indexing.py @@ -14,6 +14,7 @@ METRIC = "l2" # "l2" or "ip" FASTER_QUANT = False # use faster quantization NUM_THREADS = 16 # number of threads for build +NUM_CLUSTERS = 16 # centroids to train when no clustering file is given # ────────────────────────────────────────────── @@ -22,31 +23,38 @@ def main(args=None) -> None: data = read_fvecs(args.data_file) print(f"Data shape: {data.shape}") - # 2. Load clusters computed by RaBitQKMeans or QGKMeans. - centroids, cluster_ids = load_clusters(args.clusters, data.shape, args.metric) - print(f"Centroids: {centroids.shape}, cluster_ids: {cluster_ids.shape}") - - # 3. Build HNSW index + # 2. Train centroids in the index, or reuse a kmeans_clustering.py file. n, dim = data.shape print( f"\nBuilding HNSW index: n={n}, dim={dim}, M={args.degree}, " f"ef={args.ef_construction}, bits={args.total_bits}, metric={args.metric}" ) - + if args.clusters is None: + centroids, cluster_ids = None, None + num_clusters = args.num_clusters + else: + centroids, cluster_ids = load_clusters(args.clusters, data.shape, args.metric) + print(f"Centroids: {centroids.shape}, cluster_ids: {cluster_ids.shape}") + num_clusters = len(centroids) idx = HnswIndex( - dim=dim, - max_elements=n, + dim, + nbits=args.total_bits, + num_clusters=num_clusters, M=args.degree, ef_construction=args.ef_construction, - nbits=args.total_bits, metric=args.metric, ) t0 = time() - idx.build( + if centroids is None: + idx.train(data, num_threads=args.num_threads) + else: + idx.set_centroids(centroids) + # The first add into the empty index links vectors with their exact + # coordinates. More than one thread links them in a nondeterministic order. + idx.add( data, - centroids, - cluster_ids, + cluster_ids=cluster_ids, num_threads=args.num_threads, fast_quantization=args.faster_quant, ) @@ -63,8 +71,19 @@ def main(args=None) -> None: parser.add_argument("index_file", type=str, help="Path to save the index") parser.add_argument( "--clusters", - required=True, - help="Clustering file produced by kmeans_clustering.py for this data and metric", + default=None, + help=( + "Clustering file produced by kmeans_clustering.py for this data and " + "metric; without it, the index trains its own centroids" + ), + ) + parser.add_argument( + "--num-clusters", + dest="num_clusters", + type=int, + metavar="INT", + default=NUM_CLUSTERS, + help="Number of centroids to train when --clusters is not given", ) parser.add_argument( "--degree", diff --git a/sample/python/ivf_rabitq_indexing.py b/sample/python/ivf_rabitq_indexing.py index 9942499b..3fdbebd2 100644 --- a/sample/python/ivf_rabitq_indexing.py +++ b/sample/python/ivf_rabitq_indexing.py @@ -12,6 +12,7 @@ METRIC = "l2" # "l2" or "ip" FASTER_QUANT = True # use faster quantization NUM_THREADS = 16 # number of threads for building the index +NUM_CLUSTERS = 4096 # centroids to train when no clustering file is given # ────────────────────────────────────────────── @@ -23,29 +24,29 @@ def main(args=None) -> None: print(f"\tN: {n}") print(f"\tDIM: {dim}") - # 2. Load clusters computed by RaBitQKMeans or QGKMeans. - centroids, cluster_ids = load_clusters(args.clusters, data.shape, args.metric) - print(f"Centroids: {centroids.shape}, cluster_ids: {cluster_ids.shape}") - - # 3. Build IVF index + # 2. Train centroids in the index, or reuse a kmeans_clustering.py file. print( f"\nBuilding IVF index: bits={args.total_bits}, metric={args.metric}, " f"num_threads={args.num_threads}, faster_quant={args.faster_quant}" ) - - idx = IvfIndex( - dim=dim, - max_elements=n, - num_clusters=len(centroids), - nbits=args.total_bits, - metric=args.metric, - ) - t0 = time() - idx.build( + if args.clusters is None: + idx = IvfIndex( + dim, args.num_clusters, nbits=args.total_bits, metric=args.metric + ) + idx.train(data, num_threads=args.num_threads) + cluster_ids = None + print(f"Trained {args.num_clusters} centroids") + else: + centroids, cluster_ids = load_clusters(args.clusters, data.shape, args.metric) + print(f"Centroids: {centroids.shape}, cluster_ids: {cluster_ids.shape}") + idx = IvfIndex(dim, len(centroids), nbits=args.total_bits, metric=args.metric) + idx.set_centroids(centroids) + + # 3. Add the vectors + idx.add( data, - centroids, - cluster_ids, + cluster_ids=cluster_ids, num_threads=args.num_threads, fast_quantization=args.faster_quant, ) @@ -64,8 +65,19 @@ def main(args=None) -> None: parser.add_argument("index_file", type=str, help="Path to save the index") parser.add_argument( "--clusters", - required=True, - help="Clustering file produced by kmeans_clustering.py for this data and metric", + default=None, + help=( + "Clustering file produced by kmeans_clustering.py for this data and " + "metric; without it, the index trains its own centroids" + ), + ) + parser.add_argument( + "--num-clusters", + dest="num_clusters", + type=int, + metavar="INT", + default=NUM_CLUSTERS, + help="Number of centroids to train when --clusters is not given", ) parser.add_argument( "--total-bits", diff --git a/src/index/ivf.cpp b/src/index/ivf.cpp index 0984d5bd..fa7e69ed 100644 --- a/src/index/ivf.cpp +++ b/src/index/ivf.cpp @@ -248,9 +248,15 @@ void IVF::set_centroids(const float* centroids, size_t num_threads) { // 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_]); - } + index_detail::parallel_for( + 0, + num_clusters_, + num_threads, + [&](size_t i, size_t /*thread*/) { + rotator_->rotate(centroids + (i * dim_), &rotated_centroids[i * padded_dim_]); + }, + /*reuse_workers=*/true + ); auto initializer = make_initializer(); initializer->add_vectors( rotated_centroids.data(), std::min(resolve_num_threads(num_threads), num_clusters_) diff --git a/tests/python/test_compat_08.py b/tests/python/test_compat_08.py new file mode 100644 index 00000000..5cdcda14 --- /dev/null +++ b/tests/python/test_compat_08.py @@ -0,0 +1,214 @@ +"""The 0.7 call forms the 0.8 compatibility layer keeps; removed in 0.9.""" + +import inspect +import warnings + +import numpy as np +import pytest +from conftest import DIM, N_VECTORS +from rabitqlib import HnswIndex, IvfIndex + + +def single_deprecation(call): + """Run call and return its result and its one DeprecationWarning.""" + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + result = call() + assert [w.category for w in caught] == [DeprecationWarning], [ + str(w.message) for w in caught + ] + # Attributed to the caller, so users see their own line. + assert caught[0].filename == __file__ + assert "python_migration" in str(caught[0].message) + return result, str(caught[0].message) + + +def quietly(call): + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + return call() + + +@pytest.mark.parametrize( + "make", + [ + lambda: IvfIndex(DIM, N_VECTORS, 4, 4), + lambda: IvfIndex(DIM, N_VECTORS, 4, nbits=4), + lambda: IvfIndex(DIM, N_VECTORS, num_clusters=4, nbits=4), + lambda: IvfIndex( + dim=DIM, max_elements=N_VECTORS, num_clusters=4, nbits=4, metric="ip" + ), + ], +) +def test_ivf_point_count_forms(make, base_data, clusters): + index, message = single_deprecation(make) + assert message.startswith( + "IvfIndex(dim, max_elements, num_clusters, nbits, ...) is deprecated" + ) + assert index.max_elements == N_VECTORS + assert index.num_clusters == 4 + assert index.seed is None + assert not index.is_trained + + +@pytest.mark.parametrize( + "make", + [ + lambda: HnswIndex(DIM, N_VECTORS), + lambda: HnswIndex(DIM, N_VECTORS, 8, 50, 4, "ip", 5), + lambda: HnswIndex(dim=DIM, max_elements=N_VECTORS, M=8, nbits=4, random_seed=5), + ], +) +def test_hnsw_max_elements_forms(make): + index, message = single_deprecation(make) + assert message.startswith("HnswIndex(dim, max_elements, ...) is deprecated") + assert index.max_elements == N_VECTORS + + +def test_old_forms_keep_old_defaults(): + hnsw = quietly(lambda: HnswIndex(DIM, N_VECTORS)) + assert hnsw.nbits == 8 + assert hnsw.seed == 100 + assert hnsw.num_clusters == 16 + build = inspect.signature(IvfIndex.build).parameters + assert build["num_threads"].default == 1 + assert inspect.signature(HnswIndex.build).parameters["num_threads"].default == 1 + + +def test_random_seed_keyword(): + index, message = single_deprecation(lambda: HnswIndex(DIM, nbits=4, random_seed=7)) + assert message.startswith("HnswIndex(random_seed=...) is deprecated; use seed=") + assert index.seed == 7 + with pytest.raises( + TypeError, match="^HnswIndex\\(\\) got both seed and its old name random_seed$" + ): + HnswIndex(DIM, nbits=4, seed=1, random_seed=2) + + +@pytest.mark.parametrize( + "make,message", + [ + # Mixing forms: the old form has no seed, and the new one no max_elements. + ( + lambda: IvfIndex(DIM, N_VECTORS, 4, nbits=4, seed=1), + "IvfIndex() got an unexpected keyword argument 'seed' (read as the 0.7 form", + ), + ( + lambda: IvfIndex(DIM, 4, nbits=4, max_elements=N_VECTORS), + "IvfIndex() multiple values for argument 'max_elements' (read as the 0.7 form", + ), + ( + lambda: HnswIndex(DIM, max_elements=10, num_clusters=4), + "HnswIndex() got an unexpected keyword argument 'num_clusters' (read as the 0.7 form", + ), + ( + lambda: HnswIndex(DIM, 10, seed=4), + "HnswIndex() got an unexpected keyword argument 'seed' (read as the 0.7 form", + ), + ], +) +def test_mixed_forms_raise_instead_of_guessing(make, message): + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + with pytest.raises(TypeError) as error: + make() + assert str(error.value).startswith(message) + + +@pytest.mark.parametrize("kind", ["ivf", "hnsw"]) +def test_build_and_is_built(kind, base_data, clusters): + centroids, cluster_ids = clusters + if kind == "ivf": + index = quietly(lambda: IvfIndex(DIM, N_VECTORS, len(centroids), 4)) + else: + index = quietly(lambda: HnswIndex(DIM, N_VECTORS, nbits=4)) + built, _ = single_deprecation(lambda: index.is_built) + assert built is False + _, message = single_deprecation( + lambda: index.build(base_data, centroids, cluster_ids) + ) + assert message.startswith(f"{type(index).__name__}.build() is deprecated") + assert index.is_trained + assert len(index) == N_VECTORS + built, message = single_deprecation(lambda: index.is_built) + assert built is True + assert message.startswith(f"{type(index).__name__}.is_built is deprecated") + + +def test_build_on_new_form_indexes(base_data, clusters, tmp_path): + centroids, cluster_ids = clusters + ivf = IvfIndex(DIM, len(centroids), nbits=4, seed=3) + single_deprecation(lambda: ivf.build(base_data, centroids, cluster_ids)) + reference = IvfIndex(DIM, len(centroids), nbits=4, seed=3) + reference.set_centroids(centroids) + reference.add(base_data, cluster_ids=cluster_ids, num_threads=1) + ivf.save(tmp_path / "a.index") + reference.save(tmp_path / "b.index") + assert (tmp_path / "a.index").read_bytes() == (tmp_path / "b.index").read_bytes() + + # HNSW build takes any number of centroids, as in 0.7. + hnsw = HnswIndex(DIM, nbits=4) + assert hnsw.num_clusters == 16 + single_deprecation(lambda: hnsw.build(base_data, centroids, cluster_ids)) + assert hnsw.num_clusters == len(centroids) + assert hnsw.max_elements == N_VECTORS + + +@pytest.mark.parametrize("kind", ["ivf", "hnsw"]) +def test_positional_add_options(kind, base_data): + index = ( + IvfIndex(DIM, 2, nbits=4) + if kind == "ivf" + else HnswIndex(DIM, nbits=4, num_clusters=2) + ) + index.set_centroids(base_data[:2]) + labels = np.zeros(10, dtype=np.uint32) + ids, message = single_deprecation(lambda: index.add(base_data[:10], labels, 1)) + assert message.startswith( + f"passing cluster_ids, num_threads to {type(index).__name__}.add() by position is deprecated" + ) + np.testing.assert_array_equal(ids, np.arange(10)) + with pytest.raises( + TypeError, match="got multiple values for argument 'cluster_ids'" + ): + quietly(lambda: index.add(base_data[:10], labels, cluster_ids=labels)) + with pytest.raises( + TypeError, match=r"takes at most 4 positional arguments \(5 given\)" + ): + index.add(base_data[:10], labels, 1, False, 3) + assert len(index) == 10 + + +def test_failed_build_leaves_a_new_form_index_unchanged(base_data, clusters): + centroids, cluster_ids = clusters + bad = base_data.copy() + bad[7, 3] = np.nan + ivf = IvfIndex(DIM, len(centroids), nbits=4) + hnsw = HnswIndex(DIM, nbits=4) + for index in (ivf, hnsw): + with pytest.raises(ValueError, match="must contain only finite values"): + quietly(lambda: index.build(bad, centroids, cluster_ids)) + assert not index.is_trained + assert len(index) == 0 + assert hnsw.max_elements == 1 + bad_centroids = centroids.copy() + bad_centroids[0, 0] = np.inf + with pytest.raises(ValueError, match="must contain only finite values"): + quietly(lambda: hnsw.build(base_data, bad_centroids, cluster_ids)) + assert hnsw.max_elements == 1 + + +def test_hnsw_build_on_a_trained_index_replaces_centroids( + base_data, clusters, tmp_path +): + centroids, cluster_ids = clusters + trained = HnswIndex(DIM, nbits=4, num_clusters=4) + trained.train(base_data) + single_deprecation(lambda: trained.build(base_data, centroids, cluster_ids)) + assert trained.num_clusters == len(centroids) + # The same index the 0.7 form builds: both link with the exact vectors. + built = quietly(lambda: HnswIndex(DIM, N_VECTORS, nbits=4)) + quietly(lambda: built.build(base_data, centroids, cluster_ids)) + trained.save(tmp_path / "a.index") + built.save(tmp_path / "b.index") + assert (tmp_path / "a.index").read_bytes() == (tmp_path / "b.index").read_bytes() diff --git a/tests/python/test_editor_support.py b/tests/python/test_editor_support.py index 0e2d0e91..eaccc5d1 100644 --- a/tests/python/test_editor_support.py +++ b/tests/python/test_editor_support.py @@ -2,6 +2,7 @@ import ast import doctest +import inspect import re from pathlib import Path @@ -84,26 +85,72 @@ def test_editor_docs_match_runtime_help(): assert normalized(documentation) in normalized(actual), (name, node.name) +def is_deprecated(node): + return any( + isinstance(d, ast.Call) + and isinstance(d.func, ast.Name) + and d.func.id == "deprecated" + for d in node.decorator_list + ) + + +def stub_parameters(node): + """Names, keyword-only names and literal defaults declared in the stub.""" + args = node.args + positional = args.posonlyargs + args.args + defaults = { + arg.arg: ast.literal_eval(value) + for arg, value in zip( + positional[len(positional) - len(args.defaults) :], args.defaults + ) + } + defaults.update( + (arg.arg, ast.literal_eval(value)) + for arg, value in zip(args.kwonlyargs, args.kw_defaults) + if value is not None + ) + names = [arg.arg for arg in positional + args.kwonlyargs] + return names, [arg.arg for arg in args.kwonlyargs], defaults + + +def runtime_parameters(member): + """The same, from a Python replacement's signature or pybind11's first doc line.""" + if inspect.isfunction(member): + parameters = inspect.signature(member).parameters.values() + return ( + [p.name for p in parameters], + [p.name for p in parameters if p.kind is inspect.Parameter.KEYWORD_ONLY], + {p.name: p.default for p in parameters if p.default is not p.empty}, + ) + signature = member.__doc__.splitlines()[0] + names = re.findall(r"(?:\(|, )([A-Za-z_]\w*):", signature) + _, star, keyword_part = signature.partition(", *, ") + keyword = re.findall(r"(?:^|, )([A-Za-z_]\w*):", keyword_part) if star else [] + defaults = { + found[1]: ast.literal_eval(found[2]) + for found in re.finditer(r"\b([A-Za-z_]\w*): [^=,]+ = ([^,)]+)", signature) + } + return names, keyword, defaults + + @pytest.mark.parametrize("name", ["IvfIndex", "HnswIndex", "SymqgIndex"]) def test_native_signature_names_and_defaults(name): cls = declarations()[name] runtime = getattr(rabitqlib, name) for node in cls.body: - if not isinstance(node, ast.FunctionDef): + if not isinstance(node, ast.FunctionDef) or is_deprecated(node): + # Deprecated overloads describe the 0.7 forms the compatibility layer + # accepts; the runtime signature documents only the current form. continue if any( - isinstance(d, ast.Name) and d.id == "property" for d in node.decorator_list + (isinstance(d, ast.Name) and d.id == "property") + or (isinstance(d, ast.Attribute) and d.attr == "setter") + for d in node.decorator_list ): continue - signature = getattr(runtime, node.name).__doc__.splitlines()[0] - actual_names = re.findall(r"(?:\(|, )([A-Za-z_]\w*):", signature) - expected_names = [a.arg for a in node.args.args] - assert actual_names == expected_names, (name, node.name, signature) - default_args = node.args.args[len(node.args.args) - len(node.args.defaults) :] - for arg, value in zip(default_args, node.args.defaults): - actual = re.search(rf"\b{arg.arg}: [^=]+ = ([^,)]+)", signature) - assert actual is not None, (name, node.name, arg.arg) - assert ast.literal_eval(actual[1]) == ast.literal_eval(value) + expected = stub_parameters(node) + actual = runtime_parameters(getattr(runtime, node.name)) + assert actual == expected, (name, node.name) @pytest.mark.parametrize("name", ["QGKMeans", "RaBitQKMeans"]) @@ -153,13 +200,12 @@ def test_documented_search_result_types(kind): index.build(data, 40) result = index.search(data[:2], 3, ef=20) else: - centroids = data[:2].copy() - labels = np.arange(80, dtype=np.uint32) % 2 if kind == "ivf": - index = rabitqlib.IvfIndex(64, 80, 2, 4) + index = rabitqlib.IvfIndex(64, 2, nbits=4) else: - index = rabitqlib.HnswIndex(64, 80) - index.build(data, centroids, labels) + index = rabitqlib.HnswIndex(64, nbits=4, num_clusters=2) + index.set_centroids(data[:2]) + index.add(data, cluster_ids=np.arange(80, dtype=np.uint32) % 2) if isinstance(index, rabitqlib.IvfIndex): result = index.search(data[:2], 3, nprobe=2) else: diff --git a/tests/python/test_examples.py b/tests/python/test_examples.py index e62c8237..df1e2795 100644 --- a/tests/python/test_examples.py +++ b/tests/python/test_examples.py @@ -282,6 +282,32 @@ def test_separate_clustering_and_indexing( check_saved_index_and_query(tmp_path, kind, metric, data) +@pytest.mark.parametrize("metric", ["l2", "ip"]) +@pytest.mark.parametrize("kind", ["ivf", "hnsw"]) +def test_indexing_trains_without_a_clustering_file( + tmp_path, example_data, metric, kind, monkeypatch +): + # The samples teach the 0.8 API, so they must not hit the compatibility layer. + monkeypatch.setenv("PYTHONWARNINGS", "error::DeprecationWarning") + data_path, data = example_data + output = run_example( + f"{kind}_rabitq_indexing.py", + data_path, + tmp_path / "test.index", + "--num-clusters", + 4, + "--metric", + metric, + "--total-bits", + 5, + "--num-threads", + 2, + ) + if kind == "ivf": + assert "Trained 4 centroids" in output + check_saved_index_and_query(tmp_path, kind, metric, data) + + def check_exact_cluster_labels(data, centroids, labels, metric): original = data.astype(np.float64) centers = centroids.astype(np.float64) diff --git a/tests/python/test_hnsw.py b/tests/python/test_hnsw.py index a0ec3e00..6d8c1cd4 100644 --- a/tests/python/test_hnsw.py +++ b/tests/python/test_hnsw.py @@ -190,7 +190,7 @@ def test_wrong_query_dim_raises(built_hnsw): def test_search_before_build_raises(query_data): idx = HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4) with pytest.raises( - RuntimeError, match="^HnswIndex must be built or loaded before search$" + RuntimeError, match="^HnswIndex must be trained or loaded before search$" ): idx.search(query_data[:1], k=1) @@ -284,10 +284,12 @@ def test_add_without_cluster_ids_preserves_self_retrieval(partial_hnsw, base_dat assert np.count_nonzero(found[:, 0] == np.arange(400, N_VECTORS)) >= 0.95 * 100 -def test_add_beyond_capacity_raises(partial_hnsw, base_data): - with pytest.raises(ValueError, match="capacity"): - partial_hnsw.add(base_data) - assert partial_hnsw.num_points == 400 +def test_add_beyond_capacity_grows_it(partial_hnsw, base_data): + capacity = partial_hnsw.max_elements + partial_hnsw.add(base_data) + assert partial_hnsw.num_points == 400 + N_VECTORS + # Geometric growth: at least double, so repeated small adds stay cheap. + assert partial_hnsw.max_elements == max(400 + N_VECTORS, 2 * capacity) def test_add_rejects_bad_cluster_ids(partial_hnsw, base_data): @@ -310,7 +312,7 @@ def test_add_rejects_wrong_dimension(partial_hnsw): def test_add_before_build_raises(base_data): idx = HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4) - with pytest.raises(RuntimeError, match="built or loaded"): + with pytest.raises(RuntimeError, match="trained or loaded"): idx.add(base_data[:4]) @@ -320,9 +322,6 @@ def test_resize_grows_capacity_and_allows_more(base_data, clusters): idx.build(base_data[:400], centroids, cluster_ids[:400]) assert idx.max_elements == 400 - with pytest.raises(ValueError, match="capacity"): - idx.add(base_data[400:]) - idx.resize(N_VECTORS) assert idx.max_elements == N_VECTORS assert idx.num_points == 400 @@ -495,7 +494,7 @@ def test_remove_rejects_invalid_ids_without_removing_anything(): assert idx.remove([0, 1]) == 2 unbuilt = HnswIndex(64, 4, M=8, ef_construction=50, nbits=4) - with pytest.raises(RuntimeError, match="built or loaded"): + with pytest.raises(RuntimeError, match="trained or loaded"): unbuilt.remove([0]) diff --git a/tests/python/test_index_repr.py b/tests/python/test_index_repr.py index a2c79452..16c23005 100644 --- a/tests/python/test_index_repr.py +++ b/tests/python/test_index_repr.py @@ -8,12 +8,12 @@ "ivf": ( "IvfIndex(dim=64, metric='ip', nbits={bits}, num_points={live}, " "num_stored_points={stored}, max_elements={capacity}, " - "num_clusters={clusters}, initializer='flat', is_built={built})" + "num_clusters={clusters}, initializer='flat', is_trained={built})" ), "hnsw": ( "HnswIndex(dim=64, metric='ip', nbits={bits}, num_points={live}, " "num_stored_points={stored}, max_elements={capacity}, " - "num_clusters={clusters}, is_built={built})" + "num_clusters={clusters}, is_trained={built})" ), "symqg": ( "SymqgIndex(dim=64, metric='ip', quantization_bits={bits}, " @@ -55,7 +55,7 @@ def test_repr_lifecycle(kind, bits, vectors, tmp_path): live=0, stored=0, capacity=capacity, - clusters=0 if kind == "hnsw" else 1, + clusters=16 if kind == "hnsw" else 1, built=False, ) assert repr(index) == str(index) == before diff --git a/tests/python/test_input_boundaries.py b/tests/python/test_input_boundaries.py index 1fbb4dc4..1d8caab3 100644 --- a/tests/python/test_input_boundaries.py +++ b/tests/python/test_input_boundaries.py @@ -227,14 +227,15 @@ def __fspath__(self): np.testing.assert_array_equal(restored.search(queries, 3, 32)[0], expected[0]) -@pytest.mark.parametrize("index_kind,rows", [("ivf", 79), ("hnsw", 81)]) +@pytest.mark.parametrize("index_kind,rows", [("ivf", 79), ("hnsw", 0)]) def test_build_row_count_diagnostics(index_kind, rows): if index_kind == "ivf": index = IvfIndex(64, 80, 1, 4) expected = "number of data rows must match index max_elements; expected 80, got shape (79, 64)" else: + # HNSW capacity grows, so only an empty build is rejected. index = HnswIndex(64, 80) - expected = "number of data rows must be between 1 and max_elements; expected 1 to 80, got shape (81, 64)" + expected = "number of data rows must be at least 1; got shape (0, 64)" with pytest.raises(ValueError) as error: index.build( np.zeros((rows, 64)), np.zeros((1, 64)), np.zeros(rows, dtype=np.uint32) diff --git a/tests/python/test_ivf.py b/tests/python/test_ivf.py index a83f0c75..59764f96 100644 --- a/tests/python/test_ivf.py +++ b/tests/python/test_ivf.py @@ -246,7 +246,7 @@ def test_wrong_query_dim_raises(built_ivf): def test_search_before_build_raises(query_data): idx = IvfIndex(DIM, N_VECTORS, N_CLUSTERS, nbits=4) with pytest.raises( - RuntimeError, match="^IvfIndex must be built or loaded before search$" + RuntimeError, match="^IvfIndex must be trained or loaded before search$" ): idx.search(query_data[:1], k=1, nprobe=1) @@ -635,7 +635,7 @@ def test_reserve_rejects_invalid_arguments(): assert idx.max_elements == len(old) unbuilt = IvfIndex(65, 4, 1, 4) - with pytest.raises(RuntimeError, match="built or loaded"): + with pytest.raises(RuntimeError, match="trained or loaded"): unbuilt.reserve(0, 1) @@ -663,7 +663,7 @@ def test_add_rejects_invalid_arguments(): assert idx.max_elements == len(old) unbuilt = IvfIndex(65, 4, 1, 4) - with pytest.raises(RuntimeError, match="built or loaded"): + with pytest.raises(RuntimeError, match="trained or loaded"): unbuilt.add(new[:2]) @@ -743,7 +743,7 @@ def test_remove_rejects_invalid_ids_without_removing_anything(): np.testing.assert_array_equal(before[1], after[1]) unbuilt = IvfIndex(65, 4, 1, 4) - with pytest.raises(RuntimeError, match="built or loaded"): + with pytest.raises(RuntimeError, match="trained or loaded"): unbuilt.remove([0]) diff --git a/tests/python/test_train_api.py b/tests/python/test_train_api.py new file mode 100644 index 00000000..ecf5f38f --- /dev/null +++ b/tests/python/test_train_api.py @@ -0,0 +1,395 @@ +"""The 0.8 train / set_centroids / add / search workflow of IvfIndex and HnswIndex.""" + +import inspect +import warnings + +import numpy as np +import pytest +from conftest import DIM, N_VECTORS, brute_force_knn, recall_at_k +from rabitqlib import HnswIndex, IvfIndex, RaBitQKMeans + +KINDS = ["ivf", "hnsw"] + + +def new_index(kind, metric="l2", **kwargs): + if kind == "ivf": + kwargs.setdefault("seed", 7) + return IvfIndex( + DIM, kwargs.pop("num_clusters", 4), nbits=4, metric=metric, **kwargs + ) + return HnswIndex( + DIM, + nbits=4, + num_clusters=kwargs.pop("num_clusters", 4), + metric=metric, + **kwargs, + ) + + +def search(index, queries, k): + if isinstance(index, IvfIndex): + return index.search(queries, k, index.num_clusters) + return index.search(queries, k, 200) + + +def saved_bytes(index, path): + index.save(path) + return path.read_bytes() + + +@pytest.fixture(autouse=True) +def no_deprecations(): + # The 0.8 API must never go through the compatibility layer. + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + yield + + +def test_constructor_options_are_keyword_only(): + ivf = inspect.signature(IvfIndex.__init__).parameters + hnsw = inspect.signature(HnswIndex.__init__).parameters + keyword = inspect.Parameter.KEYWORD_ONLY + assert [p for p in ivf if ivf[p].kind is not keyword] == [ + "self", + "dim", + "num_clusters", + ] + assert [p for p in hnsw if hnsw[p].kind is not keyword] == ["self", "dim"] + assert ivf["nbits"].default is inspect.Parameter.empty + assert hnsw["nbits"].default is inspect.Parameter.empty + assert ivf["seed"].default is None + assert hnsw["seed"].default == 100 + assert hnsw["num_clusters"].default == 16 + for cls in (IvfIndex, HnswIndex): + add = inspect.signature(cls.add).parameters + assert [p for p in add if add[p].kind is not keyword] == ["self", "data"] + assert add["num_threads"].default == 0 + # train is native; pybind11 documents its signature on the first line. + assert ", x: object, *, num_threads:" in cls.train.__doc__.splitlines()[0] + assert cls.train.__doc__.splitlines()[0].endswith("= None) -> None") + + +@pytest.mark.parametrize( + "make,message", + [ + ( + lambda: IvfIndex(DIM, 4), + "IvfIndex() missing a required keyword-only argument: 'nbits'", + ), + ( + lambda: HnswIndex(DIM), + "HnswIndex() missing a required keyword-only argument: 'nbits'", + ), + ( + lambda: IvfIndex(DIM, 4, nbits=4, max_degree=8), + "IvfIndex() got an unexpected keyword argument 'max_degree'", + ), + ], +) +def test_constructor_argument_errors(make, message): + with pytest.raises(TypeError) as error: + make() + assert str(error.value) == message + + +def test_positional_nbits_is_read_as_the_old_form_and_explained(): + with pytest.raises( + TypeError, match=r"missing a required argument: 'nbits' \(read as the 0\.7 form" + ): + IvfIndex(DIM, 4, 4) + + +@pytest.mark.parametrize("cls", [IvfIndex, HnswIndex]) +@pytest.mark.parametrize( + "seed,error,message", + [ + (-1, ValueError, "seed must be an integer in"), + (True, TypeError, "seed must be an integer in"), + (1.5, TypeError, "seed must be an integer in"), + ], +) +def test_seed_validation(cls, seed, error, message): + with pytest.raises(error, match=message): + cls(DIM, 4, nbits=4, seed=seed) if cls is IvfIndex else cls( + DIM, nbits=4, seed=seed + ) + + +def test_ivf_seed_range_and_random_default(): + with pytest.raises(ValueError, match=r"\[0, 4294967295\]"): + IvfIndex(DIM, 4, nbits=4, seed=2**32) + drawn = [IvfIndex(DIM, 4, nbits=4, seed=None).seed for _ in range(4)] + assert all(isinstance(seed, int) and 0 <= seed < 2**32 for seed in drawn) + assert len(set(drawn)) > 1 + assert IvfIndex(DIM, 4, nbits=4, seed=np.uint32(9)).seed == 9 + + +@pytest.mark.parametrize("metric", ["l2", "ip"]) +def test_cp_defaults_follow_metric_and_seed(metric): + ivf = IvfIndex(DIM, 4, nbits=4, metric=metric, seed=11) + hnsw = HnswIndex(DIM, nbits=4, metric=metric, seed=2**40 + 5) + for index, seed in ((ivf, 11), (hnsw, 5)): + assert index.cp.niter == 25 + assert index.cp.max_points_per_centroid == 256 + assert index.cp.spherical is (metric == "ip") + assert index.cp.seed == seed + assert index.cp.num_threads == 0 + assert hnsw.seed == 2**40 + 5 + # cp is the index's own object: changes persist, and assignment copies. + ivf.cp.niter = 3 + assert ivf.cp.niter == 3 + other = RaBitQKMeans(DIM, 2, niter=7).cp + ivf.cp = other + other.niter = 8 + assert ivf.cp.niter == 7 + + +@pytest.mark.parametrize("kind", KINDS) +@pytest.mark.parametrize("metric", ["l2", "ip"]) +def test_train_add_search(kind, metric, base_data, query_data): + index = new_index(kind, metric) + assert not index.is_trained and len(index) == 0 + index.train(base_data) + assert index.is_trained and len(index) == 0 + ids = index.add(base_data) + np.testing.assert_array_equal(ids, np.arange(N_VECTORS, dtype=np.uint32)) + assert len(index) == index.num_points == N_VECTORS + if metric == "l2": + exact, _ = brute_force_knn(base_data, query_data, 10) + found, _ = search(index, query_data, 10) + assert recall_at_k(found, exact, 10) > 0.8 + + +@pytest.mark.parametrize("kind", KINDS) +def test_train_uses_cp_as_rabitqkmeans_does(kind, base_data, tmp_path): + # set_centroids with the centroids RaBitQKMeans trains from cp must give the + # same index as train, so train really clusters with cp. + trained = new_index(kind) + trained.cp.niter = 3 + trained.cp.num_threads = 1 + trained.train(base_data) + trained.add(base_data, num_threads=1) + + model = RaBitQKMeans( + DIM, + 4, + niter=3, + max_points_per_centroid=256, + seed=trained.cp.seed, + num_threads=1, + ) + model.train(base_data) + given = new_index(kind) + given.set_centroids(model.centroids) + given.add(base_data, num_threads=1) + assert saved_bytes(trained, tmp_path / "a.index") == saved_bytes( + given, tmp_path / "b.index" + ) + + +@pytest.mark.parametrize("kind", KINDS) +def test_train_num_threads_overrides_cp_only_for_the_call(kind, base_data): + index = new_index(kind) + index.cp.num_threads = 3 + index.train(base_data, num_threads=1) + assert index.cp.num_threads == 3 + + +@pytest.mark.parametrize("kind", KINDS) +def test_retraining_is_allowed_only_while_empty(kind, base_data, tmp_path): + index = new_index(kind) + index.train(base_data) + index.set_centroids(base_data[:4]) + index.train(base_data) + index.add(base_data[:50]) + name = type(index).__name__ + message = f"^{name} cannot be trained once it stores points$" + before = search(index, base_data[:5], 3) + with pytest.raises(RuntimeError, match=message): + index.train(base_data) + with pytest.raises(RuntimeError, match=message): + index.set_centroids(base_data[:4]) + after = search(index, base_data[:5], 3) + np.testing.assert_array_equal(before[0], after[0]) + + index.save(tmp_path / "trained.index") + loaded = type(index).load(tmp_path / "trained.index") + assert loaded.is_trained + assert loaded.seed is None + with pytest.raises(RuntimeError, match=message): + loaded.set_centroids(base_data[:4]) + + +@pytest.mark.parametrize("kind", KINDS) +def test_set_centroids_requires_num_clusters_by_dim(kind, base_data): + index = new_index(kind) + for bad in (base_data[:3], base_data[:4, :32], base_data[0]): + with pytest.raises(ValueError) as error: + index.set_centroids(bad) + if bad.ndim == 2: + assert str(error.value).startswith( + "centroids must have shape (num_clusters, dim); expected (4, 64), got shape" + ) + assert not index.is_trained + + +@pytest.mark.parametrize("kind", KINDS) +def test_train_rejects_too_few_rows_and_bad_values(kind, base_data): + index = new_index(kind) + with pytest.raises( + ValueError, + match=r"^x must have at least num_clusters rows; expected at least 4, got shape \(3, 64\)$", + ): + index.train(base_data[:3]) + bad = base_data.copy() + bad[3, 3] = np.nan + with pytest.raises(ValueError, match="finite"): + index.train(bad) + assert not index.is_trained + + +@pytest.mark.parametrize("kind", KINDS) +def test_operations_before_training_raise_and_change_nothing(kind, base_data, tmp_path): + index = new_index(kind) + name = type(index).__name__ + operations = { + "add": lambda: index.add(base_data), + "search": lambda: search(index, base_data[:1], 1), + "remove": lambda: index.remove([0]), + "save": lambda: index.save(tmp_path / "untrained.index"), + } + if kind == "ivf": + operations["reserve"] = lambda: index.reserve(0, 10) + for operation, call in operations.items(): + with pytest.raises( + RuntimeError, match=f"^{name} must be trained or loaded before {operation}$" + ): + call() + assert not (tmp_path / "untrained.index").exists() + assert not index.is_trained and len(index) == 0 + index.train(base_data) + index.add(base_data) + assert len(index) == N_VECTORS + + +@pytest.mark.parametrize("kind", KINDS) +def test_trained_index_without_points_cannot_be_saved(kind, base_data, tmp_path): + index = new_index(kind) + index.train(base_data) + path = tmp_path / "empty.index" + with pytest.raises( + RuntimeError, + match=f"^{type(index).__name__} has no points to save; add points first$", + ): + index.save(path) + assert not path.exists() + index.add(base_data) + index.save(path) + assert type(index).load(path).num_points == N_VECTORS + + +@pytest.mark.parametrize("kind", KINDS) +def test_search_k_is_bounded_by_stored_points(kind, base_data): + # Not by HNSW capacity, which grows geometrically. + index = new_index(kind) + index.train(base_data) + bound = "num_stored_points" if kind == "hnsw" else "max_elements" + with pytest.raises(ValueError, match=f"^k must be between 1 and {bound}$"): + search(index, base_data[:1], 1) + index.add(base_data[:10]) + index.add(base_data[10:11]) + assert search(index, base_data[:1], 11)[0].shape == (1, 11) + with pytest.raises(ValueError, match=f"^k must be between 1 and {bound}$"): + search(index, base_data[:1], 12) + + +@pytest.mark.parametrize("kind", KINDS) +def test_len_counts_live_points(kind, base_data): + index = new_index(kind) + index.train(base_data) + index.add(base_data) + assert index.remove([0, 1, 1]) == 2 + assert len(index) == index.num_points == N_VECTORS - 2 + assert index.num_stored_points == N_VECTORS + + +def test_same_seed_reproduces_an_ivf_index(base_data, tmp_path): + first = IvfIndex(DIM, 4, nbits=4) + second = IvfIndex(DIM, 4, nbits=4, seed=first.seed) + for index in (first, second): + index.cp.num_threads = 1 + index.train(base_data) + index.add(base_data, num_threads=1) + assert saved_bytes(first, tmp_path / "a.index") == saved_bytes( + second, tmp_path / "b.index" + ) + + +def test_ivf_point_count_index_cannot_be_trained(base_data): + with pytest.warns(DeprecationWarning): + index = IvfIndex(DIM, N_VECTORS, 4, 4) + message = r"^IvfIndex created with max_elements is filled by build; create IvfIndex\(dim, num_clusters, nbits=\.\.\.\) to train or set centroids$" + with pytest.raises(RuntimeError, match=message): + index.train(base_data) + with pytest.raises(RuntimeError, match=message): + index.set_centroids(base_data[:4]) + assert index.max_elements == N_VECTORS + + +def test_hnsw_capacity_grows_and_resize_preallocates(base_data): + index = HnswIndex(DIM, nbits=4, num_clusters=4) + assert index.max_elements == 1 + index.train(base_data) + index.add(base_data[:100]) + assert index.max_elements == 100 + index.add(base_data[100:110]) + assert index.max_elements == 200 # doubled + preallocated = HnswIndex(DIM, nbits=4, num_clusters=4) + preallocated.resize(1000) + preallocated.train(base_data) + preallocated.add(base_data) + assert preallocated.max_elements == 1000 + + +def test_hnsw_failed_add_keeps_the_capacity(base_data): + index = HnswIndex(DIM, nbits=4, num_clusters=4) + index.train(base_data) + index.add(base_data[:100]) + bad = base_data.copy() + bad[150, 0] = np.nan + with pytest.raises(ValueError, match="must contain only finite values"): + index.add(bad) + assert index.max_elements == 100 + assert len(index) == 100 + + +def test_hnsw_num_clusters_is_validated(): + with pytest.raises(ValueError, match=r"^num_clusters must be in \[1, "): + HnswIndex(DIM, nbits=4, num_clusters=0) + + +def test_hnsw_first_add_matches_deprecated_build(base_data, clusters, tmp_path): + # The first add into an empty index links like the 0.7 build, so the two + # paths give the same bytes when single-threaded. + centroids, cluster_ids = clusters + added = HnswIndex( + DIM, nbits=4, num_clusters=len(centroids), M=8, ef_construction=50 + ) + added.set_centroids(centroids) + added.add(base_data, cluster_ids=cluster_ids, num_threads=1) + with pytest.warns(DeprecationWarning): + built = HnswIndex(DIM, N_VECTORS, M=8, ef_construction=50, nbits=4) + built.build(base_data, centroids, cluster_ids) + assert saved_bytes(added, tmp_path / "a.index") == saved_bytes( + built, tmp_path / "b.index" + ) + + +@pytest.mark.parametrize("kind", KINDS) +def test_add_without_cluster_ids_routes_to_trained_centroids(kind, base_data): + index = new_index(kind) + index.train(base_data) + index.add(base_data[:250]) + index.add(base_data[250:]) + found, _ = search(index, base_data[:20], 1) + assert np.mean(found[:, 0] == np.arange(20)) >= 0.9 diff --git a/tests/python/test_txtai_compat.py b/tests/python/test_txtai_compat.py new file mode 100644 index 00000000..a005cb6b --- /dev/null +++ b/tests/python/test_txtai_compat.py @@ -0,0 +1,95 @@ +"""txtai's rabitqlib backend, call for call, must keep working in 0.8. + +Mirrors txtai.ann.dense.rabitq.RabitQ (index, append, delete, search, count, +save, load). txtai trains centroids with scikit-learn; any centroids exercise the +same calls, so these come from the data. The only warnings allowed are the +compatibility layer's DeprecationWarnings. +""" + +import math +import warnings + +import numpy as np +import pytest +from rabitqlib import HnswIndex, IvfIndex + +SENTINEL = 0xFFFFFFFF + + +@pytest.fixture +def embeddings(): + vectors = np.random.default_rng(3).normal(size=(400, 64)).astype(np.float32) + return vectors / np.linalg.norm(vectors, axis=1, keepdims=True) + + +def index_like_txtai(mode, vectors): + rows = vectors.shape[0] + clusters = max(1, min(round(4 * math.sqrt(rows)), rows)) + centroids = np.ascontiguousarray(vectors[:clusters], dtype=np.float32) + clusterids = np.ascontiguousarray( + np.argmax(vectors @ centroids.T, axis=1), dtype=np.uint32 + ) + dim = vectors.shape[1] + if mode == "ivf": + backend = IvfIndex( + dim=dim, max_elements=rows, num_clusters=clusters, nbits=1, metric="ip" + ) + else: + backend = HnswIndex( + dim=dim, + max_elements=rows, + M=16, + ef_construction=200, + nbits=1, + metric="ip", + random_seed=100, + ) + backend.build(vectors, centroids, clusterids) + return backend + + +def search_like_txtai(backend, mode, queries, limit, deletes): + count = backend.max_elements - deletes + k = min(limit, count) + if mode == "ivf": + nprobe = max(1, round(backend.num_clusters / 16)) + ids, distances = backend.search(queries, k, nprobe) + else: + ids, distances = backend.search(queries, k, None or 0) + return [ + [ + (int(uid), float(1.0 - dist)) + for uid, dist in zip(uids, dists) + if uid != SENTINEL + ] + for uids, dists in zip(ids.tolist(), distances.tolist()) + ] + + +@pytest.mark.parametrize("mode", ["ivf", "hnsw"]) +def test_txtai_calls(mode, embeddings, tmp_path): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + backend = index_like_txtai(mode, embeddings[:300]) + deletes = 0 + results = search_like_txtai(backend, mode, embeddings[:5], 3, deletes) + assert [hits[0][0] for hits in results] == list(range(5)) + + if mode == "ivf": + # append and delete exist for ivf only. + backend.add(np.ascontiguousarray(embeddings[300:], dtype=np.float32)) + assert backend.max_elements == 400 + ids = [uid for uid in [0, 1, 399, 1000] if 0 <= uid < backend.max_elements] + deletes += backend.remove(np.array(ids, dtype=np.int64)) + assert deletes == 3 + results = search_like_txtai(backend, mode, embeddings[300:305], 3, deletes) + assert results[0][0][0] == 300 + + path = str(tmp_path / f"{mode}.index") + backend.save(path) + loaded = IvfIndex.load(path) if mode == "ivf" else HnswIndex.load(path) + assert search_like_txtai( + loaded, mode, embeddings[:5], 3, deletes + ) == search_like_txtai(backend, mode, embeddings[:5], 3, deletes) + categories = {w.category for w in caught} + assert categories == {DeprecationWarning}, [str(w.message) for w in caught] diff --git a/tests/python/typing/public_api.py b/tests/python/typing/public_api.py index 4888b273..91559127 100644 --- a/tests/python/typing/public_api.py +++ b/tests/python/typing/public_api.py @@ -23,17 +23,22 @@ def check(x: npt.NDArray[np.float32], labels: npt.NDArray[np.uint32]) -> None: assert_type(rabitqlib.__all__, list[str]) - ivf = IvfIndex(64, 64, 2, 4) - hnsw = HnswIndex(64, 128) + ivf = IvfIndex(64, 2, nbits=4, seed=7) + hnsw = HnswIndex(64, nbits=4, num_clusters=2) graph = SymqgIndex(64, 32) for index in (ivf, hnsw, graph): assert_type(index.__repr__(), str) for index in (ivf, hnsw): - index.build(x, x[:2], labels) - assert_type(index.add(x, labels), npt.NDArray[np.uint32]) + assert_type(index.cp, RaBitQKMeansParameters) + index.cp.niter = 10 + index.train(x, num_threads=1) + index.set_centroids(x[:2]) + assert_type(index.is_trained, bool) + assert_type(index.seed, int | None) + assert_type(index.add(x, cluster_ids=labels), npt.NDArray[np.uint32]) assert_type(index.remove([0, 1]), int) assert_type(index.dim, int) - assert_type(index.is_built, bool) + assert_type(len(index), int) index.save("index.bin") ids, distances = ivf.search(x, k=5, nprobe=2) assert_type(ids, npt.NDArray[np.uint32]) @@ -74,6 +79,8 @@ def check(x: npt.NDArray[np.float32], labels: npt.NDArray[np.uint32]) -> None: qg.cp.niter = "two" # type: ignore[assignment] RaBitQKMeans(64, 2, graph_degree=32) # type: ignore[call-arg] hnsw.search(x, 5, nprobe=2) # type: ignore[call-arg] + IvfIndex(64, 2, 4) # type: ignore[call-overload] + hnsw.train(x, 1) # type: ignore[call-arg] def stats_alias(stats: RaBitQKMeansIterationStats) -> QGKMeansIterationStats: @@ -98,14 +105,16 @@ def check_numpy_scalars() -> None: x = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) labels = np.zeros(80, dtype=np.uint32) centroids = x.mean(axis=0, keepdims=True) - ivf = IvfIndex(np.int64(64), np.int64(80), np.int64(1), np.int64(4)) - hnsw = HnswIndex(np.int64(64), np.int64(80), M=np.int64(16)) + ivf = IvfIndex(np.int64(64), np.int64(1), nbits=np.int64(4), seed=np.uint32(7)) + hnsw = HnswIndex( + np.int64(64), nbits=np.int64(4), num_clusters=np.int64(1), M=np.int64(16) + ) graph = SymqgIndex(np.int64(64), np.int64(32), quantization_bits=np.int64(4)) for index in (ivf, hnsw): - index.build( + index.set_centroids(centroids) + index.add( x, - centroids, - labels, + cluster_ids=labels, num_threads=np.int64(1), fast_quantization=np.bool_(False), ) @@ -172,8 +181,8 @@ def check_input_contracts() -> None: "innerproduct", ) for metric in metrics: - ivf = IvfIndex(64, 80, 1, 4, metric=metric) - HnswIndex(64, 80, metric=metric) + ivf = IvfIndex(64, 1, nbits=4, metric=metric) + HnswIndex(64, nbits=4, metric=metric) SymqgIndex(64, 32, metric=metric) initializers: tuple[Literal["auto", "flat", "flat_rabitq", "hnsw"], ...] = ( "auto", @@ -182,21 +191,23 @@ def check_input_contracts() -> None: "hnsw", ) for initializer in initializers: - IvfIndex(64, 80, 1, 4, initializer=initializer) + IvfIndex(64, 1, nbits=4, initializer=initializer) for ids in (signed, unsigned, [0] * 80, (0,) * 80): - ivf = IvfIndex(64, 80, 1, 4) - ivf.build(x, centroids, ids) - ivf.add(x[:1], [np.int64(0)]) - ivf.add(x[:1], (np.uint32(0),)) + ivf = IvfIndex(64, 1, nbits=4) + ivf.set_centroids(centroids) + ivf.add(x, cluster_ids=ids) + ivf.add(x[:1], cluster_ids=[np.int64(0)]) + ivf.add(x[:1], cluster_ids=(np.uint32(0),)) ivf.remove(np.array([0], dtype=np.int32)) ivf.remove(np.array([1], dtype=np.uint64)) ivf.remove([2]) ivf.remove((3,)) ivf.remove([]) ivf.remove(()) - hnsw = HnswIndex(64, 82) - hnsw.build(x, centroids, ids) - hnsw.add(x[:1], [np.uint32(0)]) + hnsw = HnswIndex(64, nbits=4, num_clusters=1) + hnsw.set_centroids(centroids) + hnsw.add(x, cluster_ids=ids, num_threads=1) + hnsw.add(x[:1], cluster_ids=[np.uint32(0)]) hnsw.remove([0]) hnsw.remove(np.array([1], dtype=np.int64)) initializations: tuple[Literal["random", "pipnn"], ...] = ("random", "pipnn") @@ -204,25 +215,25 @@ def check_input_contracts() -> None: graph = SymqgIndex(64, 32) graph.build(x, 40, init=init) # Returned option values must be usable as inputs without casts. - IvfIndex(64, 80, 1, 4, metric=ivf.metric, initializer=ivf.initializer) + IvfIndex(64, 1, nbits=4, metric=ivf.metric, initializer=ivf.initializer) def reject_invalid_inputs(x: npt.NDArray[np.float32]) -> None: float_ids: npt.NDArray[np.float64] = np.zeros(80, dtype=np.float64) bool_ids: npt.NDArray[np.bool_] = np.array([True], dtype=np.bool_) string_ids: npt.NDArray[np.str_] = np.array(["1"], dtype=np.str_) - ivf = IvfIndex(64, 80, 1, 4) - hnsw = HnswIndex(64, 80) + ivf = IvfIndex(64, 1, nbits=4) + hnsw = HnswIndex(64, nbits=4) graph = SymqgIndex(64, 32) - IvfIndex(64, 80, 1, 4, metric="cosine") # type: ignore[arg-type] - HnswIndex(64, 80, metric="L2") # type: ignore[arg-type] + IvfIndex(64, 1, nbits=4, metric="cosine") # type: ignore[call-overload] + HnswIndex(64, nbits=4, metric="L2") # type: ignore[call-overload] SymqgIndex(64, 32, metric="inner_product") # type: ignore[arg-type] - IvfIndex(64, 80, 1, 4, initializer="flat_rabit") # type: ignore[arg-type] + IvfIndex(64, 1, nbits=4, initializer="flat_rabit") # type: ignore[call-overload] graph.build(x, 40, init="pippn") # type: ignore[arg-type] - ivf.build(x, x[:1], [0.0] * 80) # type: ignore[list-item] - hnsw.build(x, x[:1], float_ids) # type: ignore[arg-type] - ivf.add(x, float_ids) # type: ignore[arg-type] - hnsw.add(x, (0.0,)) # type: ignore[arg-type] + ivf.add(x, cluster_ids=[0.0] * 80) # type: ignore[list-item] + hnsw.add(x, cluster_ids=float_ids) # type: ignore[arg-type] + ivf.train(x, num_threads="1") # type: ignore[arg-type] + hnsw.add(x, cluster_ids=(0.0,)) # type: ignore[arg-type] ivf.remove([1.5]) # type: ignore[list-item] hnsw.remove(bool_ids) # type: ignore[arg-type] ivf.remove(string_ids) # type: ignore[arg-type] diff --git a/tests/test_mypy.py b/tests/test_mypy.py index 224de9a0..5b3ee1fe 100644 --- a/tests/test_mypy.py +++ b/tests/test_mypy.py @@ -56,3 +56,50 @@ def test_static_public_usage(tmp_path): ) } assert error_lines == expected_lines, result.stdout + + +# Each 0.7 call form the compatibility layer accepts is marked deprecated in the +# stubs, so editors flag it while it still type-checks. +DEPRECATED_USAGE = """\ +import numpy as np +from rabitqlib import HnswIndex, IvfIndex + +x = np.zeros((80, 64), dtype=np.float32) +labels = np.zeros(80, dtype=np.uint32) +ivf = IvfIndex(dim=64, max_elements=80, num_clusters=1, nbits=4, metric="ip") +hnsw = HnswIndex(dim=64, max_elements=80, M=16, ef_construction=200, nbits=4) +seeded = HnswIndex(64, nbits=4, random_seed=5) +ivf.build(x, x[:1], labels) +hnsw.build(x, x[:1], labels) +ivf.add(x, labels) +flag = ivf.is_built +""" + + +def test_deprecated_forms_are_flagged(tmp_path): + usage = tmp_path / "deprecated_usage.py" + usage.write_text(DEPRECATED_USAGE) + result = subprocess.run( + [ + sys.executable, + "-m", + "mypy", + "--strict", + "--enable-error-code", + "deprecated", + "--no-incremental", + "--cache-dir", + str(tmp_path / "cache"), + str(usage), + ], + text=True, + capture_output=True, + check=False, + ) + flagged = { + int(number) + for number in re.findall( + rf"{re.escape(str(usage))}:(\d+): error: .*\[deprecated\]", result.stdout + ) + } + assert flagged == set(range(6, 13)), result.stdout diff --git a/tests/unit/rabitqlib/index/ivf_test.cpp b/tests/unit/rabitqlib/index/ivf_test.cpp index 93939d5d..c617ea42 100644 --- a/tests/unit/rabitqlib/index/ivf_test.cpp +++ b/tests/unit/rabitqlib/index/ivf_test.cpp @@ -1712,7 +1712,7 @@ TEST(IvfTrainTest, SetCentroidsAndAddMatchConstructByteForByte) { EXPECT_TRUE(saved_bytes(explicit_ids, path) == expected) << "explicit IDs"; IVF routed(kDynDim, kDynNumClusters, bits, metric, kSeed); - routed.set_centroids(centroids.data(), 1); + routed.set_centroids(centroids.data()); // Rotates on every thread. routed.add(points.data(), kCount, nullptr, false, 1); EXPECT_TRUE(saved_bytes(routed, path) == expected) << "routed"; } From 61f0e0493970ccadb41c723834ef7e2418acbf06 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sun, 1 Nov 2026 18:15:26 +0800 Subject: [PATCH 2/2] test(python): accept Python 3.11 wording for missing keyword-only arguments The compatibility layer passes inspect.Signature.bind's message through. Python 3.11 reports a missing keyword-only argument as "missing a required argument", while 3.12+ says "missing a required keyword-only argument", so the wheel jobs failed on 3.11. Match both. --- tests/python/test_train_api.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/python/test_train_api.py b/tests/python/test_train_api.py index ecf5f38f..0f8a20e3 100644 --- a/tests/python/test_train_api.py +++ b/tests/python/test_train_api.py @@ -72,24 +72,24 @@ def test_constructor_options_are_keyword_only(): @pytest.mark.parametrize( "make,message", [ + # Python 3.11 omits "keyword-only" from inspect's message; 3.12+ includes it. ( lambda: IvfIndex(DIM, 4), - "IvfIndex() missing a required keyword-only argument: 'nbits'", + r"^IvfIndex\(\) missing a required (keyword-only )?argument: 'nbits'$", ), ( lambda: HnswIndex(DIM), - "HnswIndex() missing a required keyword-only argument: 'nbits'", + r"^HnswIndex\(\) missing a required (keyword-only )?argument: 'nbits'$", ), ( lambda: IvfIndex(DIM, 4, nbits=4, max_degree=8), - "IvfIndex() got an unexpected keyword argument 'max_degree'", + r"^IvfIndex\(\) got an unexpected keyword argument 'max_degree'$", ), ], ) def test_constructor_argument_errors(make, message): - with pytest.raises(TypeError) as error: + with pytest.raises(TypeError, match=message): make() - assert str(error.value) == message def test_positional_nbits_is_read_as_the_old_form_and_explained():