diff --git a/docs/docs/index/hnsw.md b/docs/docs/index/hnsw.md index 8294f60c..72cc529e 100644 --- a/docs/docs/index/hnsw.md +++ b/docs/docs/index/hnsw.md @@ -68,7 +68,7 @@ see [C++ integration](../cpp_integration.md). | `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. | +| `k`, `ef` | Positive result count and search window. Python requires `k <= num_stored_points`, so a trained index without points rejects every `k`; omitted, `None` or 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 diff --git a/docs/docs/index/qg.md b/docs/docs/index/qg.md index 4a89ab1e..c574eee8 100644 --- a/docs/docs/index/qg.md +++ b/docs/docs/index/qg.md @@ -23,15 +23,22 @@ step. The degree-32 example requires at least 33 data points. from rabitqlib import SymqgIndex # data and queries are float32 arrays with shape (count, dim). -index = SymqgIndex(dim=data.shape[1], max_degree=32, metric="l2", quantization_bits=0) -index.build(data, ef_construction=200, num_threads=1, init="pipnn") +index = SymqgIndex(data.shape[1], max_degree=32, ef_construction=200, init="pipnn") +index.add(data) # builds the whole graph index.save("qg_example.index") loaded = SymqgIndex.load("qg_example.index") -ids, distances = loaded.search_batch(queries, k=10, ef=100, num_threads=1) +ids, distances = loaded.search(queries, k=10, ef=100) ``` -`search()` accepts the same arguments and is equivalent to `search_batch()`. +SymphonyQG follows the same `train` → `add` → `search` workflow as the other +indexes, but has nothing to train: `train(x)` only checks `x` as `add` would, +and `is_trained` is always `True`. `add` builds the graph from every vector at once +and can be called only once, as in Faiss's NSG index; a second call, or `add` on +a loaded index, raises `RuntimeError`. To index more vectors, create a new index +and add all of them. Point IDs are the row numbers of the added data. Files do +not store the construction settings, so a loaded index reports +`ef_construction`, `init` and `seed` as `None`. ### C++ @@ -64,15 +71,16 @@ qg.save("qg_example.index"); | Parameter | Contract | | --- | --- | | `dim` | Original input dimension; Python supports 64–65,536 and pads internally. | -| `max_degree` | Required positive multiple of 32, strictly smaller than the point count. | -| `quantization_bits` | 0 for owned raw vectors (default), or 4/8 for quantized storage. | -| `metric` | Python `"l2"` or `"ip"`; defaults to `"l2"`. | -| `ef_construction` | Required positive construction search window. | -| `init` | Python `"pipnn"` (default) or `"random"`. | -| `k`, `ef` | Python requires `1 <= k <= num_points` and `ef >= k`. | -| `num_threads` | Python defaults to 1; 0 selects available logical CPUs. | - -Both Python search names return `(ids, distances)` arrays. Distances are squared +| `nbits` / C++ `quantization_bits` | Python keyword: 32 (default) for owned raw vectors, or 4/8 for quantized storage. C++ uses 0 for raw vectors, as Python's `quantization_bits` did before 0.8. | +| `max_degree` | Python keyword, default 32: a positive multiple of 32, strictly smaller than the point count. It is a hard degree bound, unlike HNSW-RaBitQ's `M`. | +| `ef_construction` | Python keyword, default 200: positive construction search window used by `add`. | +| `init` | Python keyword: `"pipnn"` (default) or `"random"`. | +| `metric` | Python keyword: `"l2"` (default) or `"ip"`. | +| `seed` | Python keyword: rotation and graph-construction seed in `[0, 2**32 - 1]`. The default `None` draws one at random; `index.seed` reports it, and passing it back with `num_threads=1` rebuilds the same graph from the same data. | +| `k`, `ef` | Python requires `1 <= k <= num_points`. Omitted or zero `ef` uses `max(k, 10)`, and smaller values are raised to `k`, as for HNSW-RaBitQ. | +| `num_threads` | Python `add` defaults to 0 (every available logical CPU) and `search` to 1. | + +Python search returns `(ids, distances)` arrays. Distances are squared L2 or `1 - dot`, ordered from smallest to largest; quantized storage returns estimates. SymphonyQG has no incremental add/remove API. See shared [input, count and concurrency contracts](usage.md). @@ -117,7 +125,7 @@ as the random-initialization seed. ### Input lifetime and rebuilds The caller can release input vectors after the `QGBuilder` constructor returns. -Python retains the caller's array during `build`. Complete construction before +Python retains the caller's array during `add`. Complete construction before querying or saving. `builder.reset(data)` reinitializes vectors of the same shape and reuses workspace. @@ -135,7 +143,7 @@ The graph seed controls the rotation. An FHT/Kac seed gives the same rotation on every platform; a MatrixRotator seed reproduces the matrix only with the same standard library and build, so rely on the saved index rather than the seed. The builder seed controls random initialization and fallback edges. PiPNN uses -fixed internal seeds. +fixed internal seeds. Python's `seed` sets both the graph seed and the builder seed. ## Querying @@ -172,10 +180,11 @@ do no work and may use null buffers. Inputs and outputs must not overlap. Concurrent searches need separate outputs; do not modify the index or call `set_ef()` while searches are running. -Python `search_batch(queries, k, ef, num_threads=1)` and `search()` accept -the same two-dimensional query array. Both return `(ids, distances)` arrays -of shape `(num_queries, k)`, including empty batches. Both APIs default to one worker; `num_threads=0` -selects the available logical CPU count, capped by the number of queries. +Python `search(queries, k, ef=None, num_threads=1)` accepts a two-dimensional +query array and returns `(ids, distances)` arrays of shape `(num_queries, k)`, +including empty batches. It defaults to one worker; `num_threads=0` selects the +available logical CPU count, capped by the number of queries. The 0.7 name +`search_batch` still works in 0.8 with a `DeprecationWarning`. Queries use the original dimension, including when the internal rotation pads it. Python searches use per-call `ef` values. For access rules, see the diff --git a/docs/docs/index/usage.md b/docs/docs/index/usage.md index b0516f6d..1bd2b06a 100644 --- a/docs/docs/index/usage.md +++ b/docs/docs/index/usage.md @@ -17,8 +17,7 @@ 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. `train`, `set_centroids`, `add` and `search` (and SymphonyQG `build`) -raise `ValueError` (C++: +float32. `train`, `set_centroids`, `add` and `search` 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. @@ -45,8 +44,9 @@ IVF-RaBitQ, HNSW-RaBitQ, and SymphonyQG expose these read-only Python properties `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`). +all three indexes, `len(index)` equals `num_points`. `is_trained` reports +whether IVF-RaBitQ or HNSW-RaBitQ centroids are set (by `train`, `set_centroids`, +or `load`), and is always `True` for SymphonyQG. 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, @@ -55,10 +55,10 @@ and `num_points` for a live count. Counts survive save/load. 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. +updates: `add` builds the graph once, and all three counts are equal afterwards. `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 +plus the training state for IVF/HNSW. 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. @@ -119,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 | -| `train`, `set_centroids`, `build`, `add`, `remove`, `reserve`, `resize`, and setting `cp`, where available | Require exclusive access | +| `train`, `set_centroids`, `build`, `add`, `remove`, `reserve`, `resize`, and setting `cp`, where available | Require exclusive access, except SymphonyQG's `train`, which only checks its input and may run with reads | A conflicting call immediately raises `RuntimeError("Index is busy: conflicting operation in progress")`. @@ -139,8 +139,8 @@ 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. -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. +`train` and `add` on every index default to `0`, as does clustering; `search` +and the deprecated `build` 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 diff --git a/docs/docs/python_migration.md b/docs/docs/python_migration.md index 6f4d80dd..da1bc83d 100644 --- a/docs/docs/python_migration.md +++ b/docs/docs/python_migration.md @@ -1,12 +1,20 @@ # 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 +In 0.8, `IvfIndex`, `HnswIndex` and `SymqgIndex` follow a `train` → `add` → +`search` workflow, as Faiss indexes do. IVF-RaBitQ and HNSW-RaBitQ learn their +own centroids, so you no longer run k-means and pass the centroids and +assignments to `build`; SymphonyQG has nothing to train and builds its graph in +`add`. 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. +and will be removed in **0.9**. -## Old and new calls +Unless an item below 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. + +## IVF-RaBitQ and HNSW-RaBitQ + +### Old and new calls | 0.7 | 0.8 | | --- | --- | @@ -29,11 +37,7 @@ 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. +### Behavior changes - **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 @@ -66,6 +70,33 @@ keeps its 0.7 behavior. Calls that are valid in both versions, such as `set_centroids` raise `RuntimeError` for it. Create the index in the 0.8 form to train it. +## SymphonyQG + +| 0.7 | 0.8 | +| --- | --- | +| `SymqgIndex(dim, max_degree, metric, quantization_bits)` | `SymqgIndex(dim, max_degree=..., metric=..., nbits=...)`; every argument after `dim` is keyword-only | +| `quantization_bits=0` (raw vectors), `4`, `8` | `nbits=32` (raw vectors, the default), `4`, `8` | +| `index.build(x, ef_construction, num_threads, init)` | `SymqgIndex(dim, ef_construction=..., init=..., seed=...)`
`index.add(x, num_threads=...)` | +| `index.search_batch(q, k, ef)` | `index.search(q, k, ef)`; `ef` is now optional | +| `index.quantization_bits` | `index.nbits` (32 instead of 0 for raw vectors) | +| `index.is_built` | `len(index) > 0` | + +`train(x)` exists for a uniform workflow but only checks `x` as `add` would, +and `is_trained` is always `True`. `add` builds the graph from all vectors and can run only once; +a second `add`, or `add` on a loaded index, raises `RuntimeError`, as Faiss's +NSG index does. The deprecated `build` keeps its 0.7 behavior: it replaces any +existing graph and, unless you passed `seed`, draws a new random seed each time. + +`search` now treats `ef` like HNSW-RaBitQ: an omitted or zero `ef` uses +`max(k, 10)`, and a smaller one is raised to `k`. In 0.7, both raised +`ValueError`. `add` uses every CPU by default, while the deprecated `build` +keeps one thread. The seed now also drives graph construction: without `seed`, +each index draws one random seed and reports it as `index.seed`; 0.7 drew +separate random seeds for the rotation and the graph. Before vectors are added, +`search` and `save` raise `RuntimeError("SymqgIndex has no vectors; ...")`. + +## Files and C++ + Saved files, the `max_elements` property, and `search` results for the same index are unchanged. Files written by 0.7 load as trained indexes. diff --git a/python_bindings/__init__.py b/python_bindings/__init__.py index f2d96f4d..017c7252 100644 --- a/python_bindings/__init__.py +++ b/python_bindings/__init__.py @@ -247,12 +247,131 @@ def is_built(self): cls.is_built = property(is_built, doc=is_built.__doc__) +def _symqg_signature( + self, + dim, + *, + nbits=32, + max_degree=32, + ef_construction=200, + init="pipnn", + metric="l2", + seed=None, +): ... + + +def _symqg_legacy_signature( + self, dim, max_degree, metric="l2", quantization_bits=0 +): ... + + +def _symqg_nbits(quantization_bits): + # 0.7 stored raw vectors as 0 bits; 0.8 calls that 32. + if isinstance(quantization_bits, bool) or quantization_bits not in (0, 4, 8): + raise ValueError("quantization_bits must be 0, 4, or 8") + return 32 if quantization_bits == 0 else quantization_bits + + +def _patch_symqg(cls): + native_init = cls.__init__ + signature = inspect.signature(_symqg_signature) + legacy_signature = inspect.signature(_symqg_legacy_signature) + + def __init__(self, *args, **kwargs): + if len(args) > 1: + arguments = _bind( + legacy_signature, + "SymqgIndex", + (self, *args), + kwargs, + " (read as the 0.7 form SymqgIndex(dim, max_degree, metric, " + "quantization_bits); the 0.8 form is SymqgIndex(dim, max_degree=...))", + ) + _deprecated( + "SymqgIndex(dim, max_degree, metric, quantization_bits) is deprecated; " + "use SymqgIndex(dim, max_degree=..., metric=..., nbits=...)" + ) + nbits = _symqg_nbits(arguments["quantization_bits"]) + native_init( + self, + arguments["dim"], + nbits=nbits, + max_degree=arguments["max_degree"], + metric=arguments["metric"], + ) + return + quantization_bits = kwargs.pop("quantization_bits", _MISSING) + if quantization_bits is not _MISSING: + if "nbits" in kwargs: + raise TypeError( + "SymqgIndex() got both nbits and its old form quantization_bits" + ) + _deprecated( + "SymqgIndex(quantization_bits=...) is deprecated; use nbits=, with 32 " + "instead of 0 for raw vectors" + ) + kwargs["nbits"] = _symqg_nbits(quantization_bits) + native_init(self, **_bind(signature, "SymqgIndex", (self, *args), kwargs)) + + def build(self, data, ef_construction, num_threads=1, init="pipnn"): + """Build the graph from input vectors, replacing any existing graph. + + Deprecated since 0.8 and removed in 0.9: pass ef_construction and init to + the constructor, then call add() once. The defaults match 0.7, including + a single thread. + + Parameters + ---------- + data : array_like of float, shape (n, dim) + Vectors in the original, unpadded coordinates, with n > max_degree. + Converted to contiguous float32; a copy may be made. + ef_construction : int + Positive construction search breadth fitting uint32. + num_threads : int, optional + Worker count. Default is 1; 0 selects the available CPU count. + init : {'pipnn', 'random'}, optional + Graph initialization method. Default is 'pipnn'. + """ + _deprecated( + "SymqgIndex.build() is deprecated; pass ef_construction and init to " + "SymqgIndex() and call add()" + ) + self._build(data, ef_construction, num_threads, init) + + def search_batch(self, queries, k, ef, num_threads=1): + """Deprecated since 0.8: the same as search(); use search().""" + _deprecated("SymqgIndex.search_batch() is deprecated; use search()") + return self.search(queries, k, ef, num_threads) + + def quantization_bits(self): + """Deprecated since 0.8: storage bits with 0 for raw vectors; use nbits.""" + _deprecated( + "SymqgIndex.quantization_bits is deprecated; use nbits, which is 32 " + "instead of 0 for raw vectors" + ) + return self._quantization_bits + + def is_built(self): + """Deprecated since 0.8: whether the index holds a graph; use len().""" + _deprecated("SymqgIndex.is_built is deprecated; use len()") + return self.num_stored_points != 0 + + __init__.__doc__ = cls.__doc__ + __init__.__signature__ = signature + cls.__init__ = __init__ + cls.build = build + cls.search_batch = search_batch + cls.quantization_bits = property(quantization_bits, doc=quantization_bits.__doc__) + 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 +_patch_symqg(SymqgIndex) # --- end of the 0.8 compatibility layer ------------------------------------- diff --git a/python_bindings/__init__.pyi b/python_bindings/__init__.pyi index 7c5a677e..ca545545 100644 --- a/python_bindings/__init__.pyi +++ b/python_bindings/__init__.pyi @@ -28,6 +28,7 @@ _Path: TypeAlias = str | bytes | bytearray | PathLike[str] | PathLike[bytes] _Metric: TypeAlias = Literal["l2", "ip", "innerproduct"] _Initializer: TypeAlias = Literal["auto", "flat", "flat_rabitq", "hnsw"] +_QGInit: TypeAlias = Literal["pipnn", "random"] # Rank, ID ranges and Python bool (an int subtype) are validated at runtime. _IntegerIds: TypeAlias = npt.NDArray[np.integer[Any]] | Sequence[int | np.integer[Any]] @@ -781,7 +782,7 @@ class HnswIndex: self, queries: npt.ArrayLike, k: SupportsIndex, - ef: SupportsIndex = 0, + ef: SupportsIndex | None = 0, num_threads: SupportsIndex = 1, ) -> tuple[npt.NDArray[np.uint32], npt.NDArray[np.float32]]: """Search a batch of query vectors. @@ -794,7 +795,7 @@ class HnswIndex: k : int 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 + Search breadth. None or 0 (default) selects max(k, 10); values above the stored point count are clamped. num_threads : int, optional Worker count. Default is 1; 0 selects the available CPU count. @@ -861,29 +862,60 @@ class SymqgIndex: ---------- dim : int Original vector dimension, from 64 through 65536. - max_degree : int - Graph degree bound; a positive multiple of 32. + nbits : {4, 8, 32}, optional + Keyword-only. Storage bits: 32 (default) keeps raw float32 vectors, and 4 + or 8 store centroid-relative RaBitQ codes. + max_degree : int, optional + Keyword-only. Graph degree bound; a positive multiple of 32, smaller than + the number of vectors added. Default is 32. + ef_construction : int, optional + Keyword-only. Positive construction search breadth fitting uint32. + Default is 200. + init : {'pipnn', 'random'}, optional + Keyword-only. Graph initialization method. Default is 'pipnn'. metric : {'l2', 'ip', 'innerproduct'}, optional - Squared L2 or inner-product distance. Default is 'l2'. Vectors are not - automatically normalized for cosine similarity. - quantization_bits : {0, 4, 8}, optional - Centroid-relative RaBitQ storage at 4 or 8 bits, or raw float32 at 0 - (default). + Keyword-only. Squared L2 or inner-product distance. Default is 'l2'. + Vectors are not automatically normalized for cosine similarity. + seed : int, optional + Keyword-only. Rotation and graph-construction seed in [0, 2**32 - 1]. + 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. - Data, centroids and queries must be finite float32 values within the index - coordinate range; other values raise ValueError. + There is nothing to train: train only checks its input, and is_trained is + always True. add builds the graph from all vectors at once and may be called + only once; to add more vectors, create a new index. Load an index or add + vectors before searching or saving. Concurrent reads are supported; an + overlapping mutation raises RuntimeError. + Data and queries must be finite float32 values within the index coordinate + range; other values raise ValueError. + The 0.7 forms SymqgIndex(dim, max_degree, metric, quantization_bits), + quantization_bits=, build() and search_batch() still work in 0.8, raise + DeprecationWarning, and will be removed in 0.9. + + Examples + -------- + >>> import numpy as np + >>> from rabitqlib import SymqgIndex + >>> x = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) + >>> index = SymqgIndex(64, ef_construction=40, seed=7) + >>> 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_degree: SupportsIndex, + *, + nbits: SupportsIndex = 32, + max_degree: SupportsIndex = 32, + ef_construction: SupportsIndex = 200, + init: _QGInit = "pipnn", metric: _Metric = "l2", - quantization_bits: SupportsIndex = 0, + seed: SupportsIndex | None = None, ) -> None: """An approximate nearest-neighbor index using SymphonyQG. @@ -891,25 +923,70 @@ class SymqgIndex: ---------- dim : int Original vector dimension, from 64 through 65536. - max_degree : int - Graph degree bound; a positive multiple of 32. + nbits : {4, 8, 32}, optional + Keyword-only. Storage bits: 32 (default) keeps raw float32 vectors, and 4 + or 8 store centroid-relative RaBitQ codes. + max_degree : int, optional + Keyword-only. Graph degree bound; a positive multiple of 32, smaller than + the number of vectors added. Default is 32. + ef_construction : int, optional + Keyword-only. Positive construction search breadth fitting uint32. + Default is 200. + init : {'pipnn', 'random'}, optional + Keyword-only. Graph initialization method. Default is 'pipnn'. metric : {'l2', 'ip', 'innerproduct'}, optional - Squared L2 or inner-product distance. Default is 'l2'. Vectors are not - automatically normalized for cosine similarity. - quantization_bits : {0, 4, 8}, optional - Centroid-relative RaBitQ storage at 4 or 8 bits, or raw float32 at 0 - (default). + Keyword-only. Squared L2 or inner-product distance. Default is 'l2'. + Vectors are not automatically normalized for cosine similarity. + seed : int, optional + Keyword-only. Rotation and graph-construction seed in [0, 2**32 - 1]. + 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. - Data, centroids and queries must be finite float32 values within the index - coordinate range; other values raise ValueError. + There is nothing to train: train only checks its input, and is_trained is + always True. add builds the graph from all vectors at once and may be called + only once; to add more vectors, create a new index. Load an index or add + vectors before searching or saving. Concurrent reads are supported; an + overlapping mutation raises RuntimeError. + Data and queries must be finite float32 values within the index coordinate + range; other values raise ValueError. + The 0.7 forms SymqgIndex(dim, max_degree, metric, quantization_bits), + quantization_bits=, build() and search_batch() still work in 0.8, raise + DeprecationWarning, and will be removed in 0.9. """ - + @overload + @deprecated( + "SymqgIndex(dim, max_degree, metric, quantization_bits) is deprecated since " + "0.8; use SymqgIndex(dim, max_degree=..., metric=..., nbits=...)" + ) + def __init__( + self, + dim: SupportsIndex, + max_degree: SupportsIndex, + metric: _Metric = "l2", + quantization_bits: SupportsIndex = 0, + ) -> None: ... + @overload + @deprecated( + "SymqgIndex(quantization_bits=...) is deprecated since 0.8; use nbits=, " + "with 32 instead of 0 for raw vectors" + ) + def __init__( + self, + dim: SupportsIndex, + *, + quantization_bits: SupportsIndex, + max_degree: SupportsIndex = 32, + ef_construction: SupportsIndex = 200, + init: _QGInit = "pipnn", + metric: _Metric = "l2", + seed: SupportsIndex | None = None, + ) -> None: ... def __repr__(self) -> str: - """Return a one-line summary of configuration, counts and build state.""" + """Return a one-line summary of configuration and counts.""" + + def __len__(self) -> int: + """Return the number of live, searchable points.""" @property def dim(self) -> int: @@ -928,8 +1005,13 @@ class SymqgIndex: """Number of stored points, including removed points.""" @property + def is_trained(self) -> bool: + """Always True: SymphonyQG needs no training.""" + + @property + @deprecated("SymqgIndex.is_built is deprecated since 0.8; use len()") def is_built(self) -> bool: - """Whether the index has been built or loaded.""" + """Deprecated since 0.8: whether the index holds a graph; use len().""" @property def metric(self) -> Literal["l2", "ip"]: @@ -940,41 +1022,105 @@ class SymqgIndex: """Graph degree bound.""" @property + def ef_construction(self) -> int | None: + """Construction search breadth of the current graph or the next add, or None for an index loaded from disk.""" + + @property + def init(self) -> _QGInit | None: + """Graph initialization method of the current graph or the next add, or None for an index loaded from disk.""" + + @property + def nbits(self) -> int: + """Storage bits: 32 for raw float32, or 4 or 8 for RaBitQ.""" + + @property + @deprecated( + "SymqgIndex.quantization_bits is deprecated since 0.8; use nbits, which is " + "32 instead of 0 for raw vectors" + ) def quantization_bits(self) -> int: - """Storage bits: 0 for raw float32, or 4 or 8 for RaBitQ.""" + """Deprecated since 0.8: storage bits with 0 for raw vectors; use nbits.""" + @property + def seed(self) -> int | None: + """Rotation and graph-construction seed, or None for an index loaded from disk.""" + + def train( + self, x: npt.ArrayLike, *, num_threads: SupportsIndex | None = None + ) -> None: + """Check training vectors; SymphonyQG has nothing to train. + + Parameters + ---------- + x : array_like of float, shape (n, dim) + Vectors in the original, unpadded coordinates, checked like the data of + add. + num_threads : int, optional + Keyword-only. Accepted for interface consistency with the other indexes. + Default is None. + + Notes + ----- + Provided so that SymphonyQG follows the same train/add/search workflow as + the other indexes. It does not change the index, so it may run alongside + searches. + """ + + def add(self, data: npt.ArrayLike, *, num_threads: SupportsIndex = 0) -> None: + """Build the graph from all vectors. + + Parameters + ---------- + data : array_like of float, shape (n, dim) + Vectors in the original, unpadded coordinates, with n > max_degree. + Converted to contiguous float32; a copy may be made. + num_threads : int, optional + Keyword-only. Worker count. Default is 0, which selects the available + CPU count. + + Notes + ----- + Uses the constructor's ef_construction, init and seed. The graph cannot be + extended: a second call raises RuntimeError, so add every vector in one + call, or create a new index to rebuild. Point IDs are the row numbers of + data. Input arrays need not be retained after this synchronous call. + """ + + @deprecated( + "SymqgIndex.build() is deprecated since 0.8; pass ef_construction and init " + "to SymqgIndex() and call add()" + ) def build( self, data: npt.ArrayLike, ef_construction: SupportsIndex, num_threads: SupportsIndex = 1, - init: Literal["pipnn", "random"] = "pipnn", + init: _QGInit = "pipnn", ) -> None: - """Build the graph from input vectors. + """Build the graph from input vectors, replacing any existing graph. + + Deprecated since 0.8 and removed in 0.9: pass ef_construction and init to + the constructor, then call add() once. 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, with n > max_degree. + Converted to contiguous float32; a copy may be made. ef_construction : int Positive construction search breadth fitting uint32. num_threads : int, optional Worker count. Default is 1; 0 selects the available CPU count. init : {'pipnn', 'random'}, optional Graph initialization method. Default is 'pipnn'. - - Notes - ----- - Requires n > max_degree. Input arrays need not be retained after this - synchronous call. """ def search( self, queries: npt.ArrayLike, k: SupportsIndex, - ef: SupportsIndex, + ef: SupportsIndex | None = None, num_threads: SupportsIndex = 1, ) -> tuple[npt.NDArray[np.uint32], npt.NDArray[np.float32]]: """Search a batch of query vectors. @@ -986,8 +1132,9 @@ class SymqgIndex: float32; a copy may be made. k : int Number of neighbors requested, between 1 and num_points. - ef : int - Search breadth. Must be positive and at least k. + ef : int, optional + Search breadth. None or 0 (default) selects max(k, 10); smaller values + are raised to k. num_threads : int, optional Worker count. Default is 1; 0 selects the available CPU count. @@ -1005,6 +1152,7 @@ class SymqgIndex: checked at runtime, not by static type checkers. """ + @deprecated("SymqgIndex.search_batch() is deprecated since 0.8; use search()") def search_batch( self, queries: npt.ArrayLike, @@ -1012,36 +1160,10 @@ class SymqgIndex: ef: SupportsIndex, num_threads: SupportsIndex = 1, ) -> tuple[npt.NDArray[np.uint32], npt.NDArray[np.float32]]: - """Search a batch of query vectors. - - Parameters - ---------- - 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 num_points. - ef : int - Search breadth. Must be positive and at least k. - num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. - - Returns - ------- - ids : ndarray of uint32, shape (n_queries, k) - Point IDs of the selected neighbors. - distances : ndarray of float32, shape (n_queries, k) - Lower values are better; squared L2 estimates for l2 and inner-product - distance estimates for ip. These are distances, not similarity scores. - - Notes - ----- - Returns (ids, distances), in that order. Array shapes and value ranges are - checked at runtime, not by static type checkers. - """ + """Deprecated since 0.8: the same as search(); use search().""" def save(self, path: _Path) -> None: - """Save a built index to disk. + """Save an index with vectors to disk. Parameters ---------- @@ -1067,7 +1189,7 @@ class SymqgIndex: Returns ------- index : SymqgIndex - A built index ready for search. + An index ready for search; add raises because it already has a graph. """ class QGKMeans: diff --git a/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index e30fd2bd..ae8182bd 100644 --- a/python_bindings/hnsw_bindings.cpp +++ b/python_bindings/hnsw_bindings.cpp @@ -238,16 +238,17 @@ class HnswIndex { return index_->num_stored_points(); } - py::tuple search(py::handle queries, size_t k, size_t ef, size_t num_threads) { + py::tuple search( + py::handle queries, size_t k, std::optional ef, size_t num_threads + ) { const auto operation = access_.read(); auto query_array = ensure_2d_array(queries, "queries"); check_trained("search"); ensure_dimension( query_array, index_->dimension(), "query dimension does not match index dim" ); - if (ef == 0) { - ef = std::max(k, 10); - } + // None and 0 select max(k, 10), as SymqgIndex does. + const size_t ef_search = ef.value_or(0) == 0 ? std::max(k, 10) : *ef; // 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"); @@ -262,7 +263,7 @@ class HnswIndex { std::vector>> results; { py::gil_scoped_release release; - results = index_->search(queries_data, num_queries, k, ef, num_threads); + results = index_->search(queries_data, num_queries, k, ef_search, num_threads); } for (py::ssize_t i = 0; i < static_cast(results.size()); ++i) { @@ -589,7 +590,7 @@ queries : array_like of float, shape (n_queries, dim) k : int 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 + Search breadth. None or 0 (default) selects max(k, 10); values above the stored point count are clamped. num_threads : int, optional Worker count. Default is 1; 0 selects the available CPU count. diff --git a/python_bindings/symqg_bindings.cpp b/python_bindings/symqg_bindings.cpp index e144b5c3..a91e7262 100644 --- a/python_bindings/symqg_bindings.cpp +++ b/python_bindings/symqg_bindings.cpp @@ -4,11 +4,15 @@ #include #include #include +#include // IWYU pragma: keep; registers std::optional casters +#include #include #include #include #include +#include +#include #include #include #include @@ -21,6 +25,7 @@ #include "rabitqlib/index/symqg/qg.hpp" #include "rabitqlib/index/symqg/qg_builder.hpp" #include "rabitqlib/index/types.hpp" +#include "rabitqlib/utils/input_checks.hpp" namespace py = pybind11; @@ -29,76 +34,86 @@ namespace rabitqlib::python_bindings { class SymqgIndex { public: SymqgIndex( - size_t dim, size_t max_degree, const std::string& metric, size_t quantization_bits + size_t dim, + size_t nbits, + size_t max_degree, + size_t ef_construction, + const std::string& init, + const std::string& metric, + const py::object& seed ) : dim_(dim) , max_degree_(max_degree) , metric_type_(metric_from_string(metric)) - , quantization_bits_(quantization_bits) { - if (max_degree == 0 || max_degree % fastscan::kBatchSize != 0) { - throw std::invalid_argument("max_degree must be a positive multiple of 32"); + , ef_construction_(ef_construction) + , init_(init) { + // Reject now what add would reject only after converting all the data. + if (dim < 64 || dim > rotator_impl::FhtKacRotator::kMaxDim) { + throw std::invalid_argument("Unsupported dimension for FhtKacRotator"); } - if (quantization_bits != 0 && quantization_bits != 4 && quantization_bits != 8) { - throw std::invalid_argument("quantization_bits must be 0, 4, or 8"); + check_max_degree(max_degree); + if (nbits != 4 && nbits != 8 && nbits != 32) { + throw std::invalid_argument("nbits must be 4, 8, or 32"); } + quantization_bits_ = nbits == 32 ? 0 : nbits; + check_construction(ef_construction, init); + seed_given_ = !seed.is_none(); + seed_ = seed_given_ ? static_cast( + checked_seed(seed, std::numeric_limits::max()) + ) + : random_seed(); + } + + // SymphonyQG learns nothing before add: train only checks x as add would, so + // the index is unchanged and a shared guard suffices. + void train(py::handle x, std::optional /*num_threads*/) { + const auto operation = access_.read(); + auto x_array = ensure_2d_array(x, "x"); + ensure_dimension(x_array, dim_, "x dimension does not match index dim"); + const auto* x_ptr = x_array.data(); + const size_t count = static_cast(x_array.size()); + py::gil_scoped_release release; + check_coordinates(x_ptr, count, dim_, "x"); } + void add(py::handle data, size_t num_threads) { + const auto operation = access_.write(); + if (built()) { + throw std::runtime_error( + "SymqgIndex does not support incremental addition; create a new index " + "and add all vectors at once" + ); + } + // Only a loaded index lacks these settings, and it is already built. + build_graph(data, *ef_construction_, num_threads, *init_, *seed_); + } + + // Deprecated: the 0.7 build, which replaces any existing graph. Without a + // user seed, each build draws a new one, as 0.7 did. void build( py::handle data, size_t ef_construction, size_t num_threads, const std::string& init ) { const auto operation = access_.write(); - if (init != "random" && init != "pipnn") { - throw std::invalid_argument("init must be 'random' or 'pipnn'"); - } - if (ef_construction == 0 || - ef_construction > std::numeric_limits::max()) { - throw std::invalid_argument("ef_construction must be positive and fit uint32"); - } - auto data_array = ensure_2d_array(data, "data"); - ensure_dimension(data_array, dim_, "data dimension does not match index dim"); - - const size_t num_points = static_cast(data_array.shape(0)); - auto index = std::make_unique>( - num_points, - dim_, - max_degree_, - metric_type_, - RotatorType::FhtKacRotator, - quantization_bits_ - ); - - const auto initialization = init == "pipnn" ? symqg::QGInitialization::PiPNN - : symqg::QGInitialization::Random; - const auto* data_ptr = data_array.data(); - { - py::gil_scoped_release release; - symqg::QGBuilder builder( - *index, - static_cast(ef_construction), - data_ptr, - num_threads, - initialization - ); - builder.build(); - } - index_ = std::move(index); + check_construction(ef_construction, init); + const uint32_t seed = seed_given_ ? *seed_ : random_seed(); + build_graph(data, ef_construction, num_threads, init, seed); + ef_construction_ = ef_construction; + init_ = init; + seed_ = seed; } - py::tuple search(py::handle queries, size_t k, size_t ef, size_t num_threads) { + py::tuple search( + py::handle queries, size_t k, std::optional ef, size_t num_threads + ) { const auto operation = access_.read(); auto query_array = ensure_2d_array(queries, "queries"); - if (!built()) { - throw std::runtime_error("SymqgIndex must be built or loaded before search"); - } + check_built("search"); if (k == 0 || k > index_->num_points()) { throw std::invalid_argument("k must be between 1 and num_points"); } - if (ef == 0) { - throw std::invalid_argument("ef must be positive"); - } - if (ef < k) { - throw std::invalid_argument("ef must be at least k"); - } + // As in HnswIndex: no or zero ef selects max(k, 10), and ef is at least k. + const size_t ef_search = + std::max(ef.value_or(0) == 0 ? std::max(k, 10) : *ef, k); ensure_dimension(query_array, dim_, "query dimension does not match index dim"); const size_t num_queries = static_cast(query_array.shape(0)); @@ -117,7 +132,7 @@ class SymqgIndex { static_cast(k), ids_data, distances_data, - ef, + ef_search, num_threads ); } @@ -128,9 +143,7 @@ class SymqgIndex { void save(py::handle path) const { const auto filename = path_string(path); const auto operation = access_.read(); - if (!built()) { - throw std::runtime_error("SymqgIndex must be built or loaded before save"); - } + check_built("save"); py::gil_scoped_release release; index_->save(filename.c_str()); } @@ -147,17 +160,18 @@ class SymqgIndex { wrapper.max_degree_ = wrapper.index_->degree_bound(); wrapper.metric_type_ = wrapper.index_->metric_type(); wrapper.quantization_bits_ = wrapper.index_->quantization_bits(); + // Files store neither the seed nor the construction settings. + wrapper.ef_construction_.reset(); + wrapper.init_.reset(); return wrapper; } [[nodiscard]] std::string repr() const { const auto operation = access_.read(); return "SymqgIndex(dim=" + std::to_string(dim_) + ", metric='" + - metric_to_string(metric_type_) + - "', quantization_bits=" + std::to_string(quantization_bits_) + + metric_to_string(metric_type_) + "', nbits=" + std::to_string(nbits()) + ", num_points=" + std::to_string(built() ? index_->num_points() : 0) + - ", max_degree=" + std::to_string(max_degree_) + - ", is_built=" + (built() ? "True" : "False") + ")"; + ", max_degree=" + std::to_string(max_degree_) + ")"; } [[nodiscard]] size_t dim() const { @@ -168,6 +182,14 @@ class SymqgIndex { const auto operation = access_.read(); return max_degree_; } + [[nodiscard]] std::optional ef_construction() const { + const auto operation = access_.read(); + return ef_construction_; + } + [[nodiscard]] std::optional init() const { + const auto operation = access_.read(); + return init_; + } [[nodiscard]] size_t num_points() const { const auto operation = access_.read(); return built() ? index_->num_points() : 0; @@ -180,29 +202,113 @@ class SymqgIndex { const auto operation = access_.read(); return built() ? index_->max_elements() : 0; } - [[nodiscard]] bool is_built() const { - const auto operation = access_.read(); - return built(); - } [[nodiscard]] std::string metric() const { const auto operation = access_.read(); return metric_to_string(metric_type_); } + [[nodiscard]] size_t nbits_property() const { + const auto operation = access_.read(); + return nbits(); + } [[nodiscard]] size_t quantization_bits() const { const auto operation = access_.read(); return quantization_bits_; } + [[nodiscard]] std::optional seed() const { + const auto operation = access_.read(); + return seed_; + } private: + static constexpr size_t kDefaultEfConstruction = 200; + SymqgIndex() = default; - // A failed build keeps the previous index; a failed load discards its wrapper. + static uint32_t random_seed() { return std::random_device{}(); } + + static void check_max_degree(size_t max_degree) { + if (max_degree == 0 || max_degree % fastscan::kBatchSize != 0) { + throw std::invalid_argument("max_degree must be a positive multiple of 32"); + } + } + + static void check_construction(size_t ef_construction, const std::string& init) { + if (init != "random" && init != "pipnn") { + throw std::invalid_argument("init must be 'random' or 'pipnn'"); + } + if (ef_construction == 0 || + ef_construction > std::numeric_limits::max()) { + throw std::invalid_argument("ef_construction must be positive and fit uint32"); + } + } + + void check_built(const char* operation) const { + if (!built()) { + throw std::runtime_error( + std::string( + "SymqgIndex has no vectors; add them or load an index before " + ) + + operation + ); + } + } + + // Build a new graph from data; a failure keeps the previous one. Requires the + // write guard. + void build_graph( + py::handle data, + size_t ef_construction, + size_t num_threads, + const std::string& init, + uint32_t seed + ) { + auto data_array = ensure_2d_array(data, "data"); + ensure_dimension(data_array, dim_, "data dimension does not match index dim"); + + const size_t num_points = static_cast(data_array.shape(0)); + auto index = std::make_unique>( + num_points, + dim_, + max_degree_, + metric_type_, + RotatorType::FhtKacRotator, + quantization_bits_, + seed + ); + + const auto initialization = init == "pipnn" ? symqg::QGInitialization::PiPNN + : symqg::QGInitialization::Random; + const auto* data_ptr = data_array.data(); + { + py::gil_scoped_release release; + symqg::QGBuilder builder( + *index, + static_cast(ef_construction), + data_ptr, + num_threads, + initialization, + seed + ); + builder.build(); + } + index_ = std::move(index); + } + + [[nodiscard]] size_t nbits() const { + return quantization_bits_ == 0 ? 32 : quantization_bits_; + } + [[nodiscard]] bool built() const { return index_ != nullptr; } size_t dim_ = 0; size_t max_degree_ = 0; MetricType metric_type_ = METRIC_L2; - size_t quantization_bits_ = 0; + size_t quantization_bits_ = 0; // 0 stores raw vectors. + // Settings of the current graph, or of the next add; unknown after load. + std::optional ef_construction_ = kDefaultEfConstruction; + std::optional init_ = "pipnn"; + std::optional seed_; // Unknown for loaded indexes. + bool seed_given_ = false; // The user chose the seed. mutable IndexAccess access_; std::unique_ptr> index_; }; @@ -220,21 +326,47 @@ Parameters ---------- dim : int Original vector dimension, from 64 through 65536. -max_degree : int - Graph degree bound; a positive multiple of 32. +nbits : {4, 8, 32}, optional + Keyword-only. Storage bits: 32 (default) keeps raw float32 vectors, and 4 + or 8 store centroid-relative RaBitQ codes. +max_degree : int, optional + Keyword-only. Graph degree bound; a positive multiple of 32, smaller than + the number of vectors added. Default is 32. +ef_construction : int, optional + Keyword-only. Positive construction search breadth fitting uint32. + Default is 200. +init : {'pipnn', 'random'}, optional + Keyword-only. Graph initialization method. Default is 'pipnn'. metric : {'l2', 'ip', 'innerproduct'}, optional - Squared L2 or inner-product distance. Default is 'l2'. Vectors are not - automatically normalized for cosine similarity. -quantization_bits : {0, 4, 8}, optional - Centroid-relative RaBitQ storage at 4 or 8 bits, or raw float32 at 0 - (default). + Keyword-only. Squared L2 or inner-product distance. Default is 'l2'. + Vectors are not automatically normalized for cosine similarity. +seed : int, optional + Keyword-only. Rotation and graph-construction seed in [0, 2**32 - 1]. + 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. -Data, centroids and queries must be finite float32 values within the index -coordinate range; other values raise ValueError.)doc"; +There is nothing to train: train only checks its input, and is_trained is +always True. add builds the graph from all vectors at once and may be called +only once; to add more vectors, create a new index. Load an index or add +vectors before searching or saving. Concurrent reads are supported; an +overlapping mutation raises RuntimeError. +Data and queries must be finite float32 values within the index coordinate +range; other values raise ValueError. +The 0.7 forms SymqgIndex(dim, max_degree, metric, quantization_bits), +quantization_bits=, build() and search_batch() still work in 0.8, raise +DeprecationWarning, and will be removed in 0.9. + +Examples +-------- +>>> import numpy as np +>>> from rabitqlib import SymqgIndex +>>> x = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) +>>> index = SymqgIndex(64, ef_construction=40, seed=7) +>>> 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 kSearchDoc[] = R"doc(Search a batch of query vectors. @@ -245,8 +377,9 @@ queries : array_like of float, shape (n_queries, dim) float32; a copy may be made. k : int Number of neighbors requested, between 1 and num_points. -ef : int - Search breadth. Must be positive and at least k. +ef : int, optional + Search breadth. None or 0 (default) selects max(k, 10); smaller values + are raised to k. num_threads : int, optional Worker count. Default is 1; 0 selects the available CPU count. @@ -265,54 +398,85 @@ checked at runtime, not by static type checkers.)doc"; py::class_(m, "SymqgIndex", kIndexDoc) .def( - py::init(), + py::init< + size_t, + size_t, + size_t, + size_t, + const std::string&, + const std::string&, + const py::object&>(), py::arg("dim"), - py::arg("max_degree"), + py::kw_only(), + py::arg("nbits") = 32, + py::arg("max_degree") = 32, + py::arg("ef_construction") = 200, + py::arg("init") = "pipnn", py::arg("metric") = "l2", - py::arg("quantization_bits") = 0, + py::arg("seed") = py::none(), kIndexDoc ) .def( - "build", - &SymqgIndex::build, + "train", + &SymqgIndex::train, + py::arg("x"), + py::kw_only(), + py::arg("num_threads") = py::none(), + R"doc(Check training vectors; SymphonyQG has nothing to train. + +Parameters +---------- +x : array_like of float, shape (n, dim) + Vectors in the original, unpadded coordinates, checked like the data of + add. +num_threads : int, optional + Keyword-only. Accepted for interface consistency with the other indexes. + Default is None. + +Notes +----- +Provided so that SymphonyQG follows the same train/add/search workflow as +the other indexes. It does not change the index, so it may run alongside +searches.)doc" + ) + .def( + "add", + &SymqgIndex::add, py::arg("data"), - py::arg("ef_construction"), - py::arg("num_threads") = 1, - py::arg("init") = "pipnn", - R"doc(Build the graph from input vectors. + py::kw_only(), + py::arg("num_threads") = 0, + R"doc(Build the graph from all vectors. Parameters ---------- data : array_like of float, shape (n, dim) - Vectors in the original, unpadded coordinates. Converted to contiguous - float32; a copy may be made. -ef_construction : int - Positive construction search breadth fitting uint32. + Vectors in the original, unpadded coordinates, with n > max_degree. + Converted to contiguous float32; a copy may be made. num_threads : int, optional - Worker count. Default is 1; 0 selects the available CPU count. -init : {'pipnn', 'random'}, optional - Graph initialization method. Default is 'pipnn'. + Keyword-only. Worker count. Default is 0, which selects the available + CPU count. Notes ----- -Requires n > max_degree. Input arrays need not be retained after this -synchronous call.)doc" +Uses the constructor's ef_construction, init and seed. The graph cannot be +extended: a second call raises RuntimeError, so add every vector in one +call, or create a new index to rebuild. Point IDs are the row numbers of +data. Input arrays need not be retained after this synchronous call.)doc" ) .def( - "search", - &SymqgIndex::search, - py::arg("queries"), - py::arg("k"), - py::arg("ef"), - py::arg("num_threads") = 1, - kSearchDoc + "_build", + &SymqgIndex::build, + py::arg("data"), + py::arg("ef_construction"), + py::arg("num_threads"), + py::arg("init") ) .def( - "search_batch", + "search", &SymqgIndex::search, py::arg("queries"), py::arg("k"), - py::arg("ef"), + py::arg("ef") = py::none(), py::arg("num_threads") = 1, kSearchDoc ) @@ -320,7 +484,7 @@ synchronous call.)doc" "save", &SymqgIndex::save, py::arg("path"), - R"doc(Save a built index to disk. + R"doc(Save an index with vectors to disk. Parameters ---------- @@ -347,12 +511,17 @@ path : str, bytes, bytearray, or os.PathLike Returns ------- index : SymqgIndex - A built index ready for search.)doc" + An index ready for search; add raises because it already has a graph.)doc" ) .def( "__repr__", &SymqgIndex::repr, - "Return a one-line summary of configuration, counts and build state." + "Return a one-line summary of configuration and counts." + ) + .def( + "__len__", + &SymqgIndex::num_points, + "Return the number of live, searchable points." ) .def_property_readonly( "dim", &SymqgIndex::dim, R"doc(Original vector dimension.)doc" @@ -360,6 +529,16 @@ index : SymqgIndex .def_property_readonly( "max_degree", &SymqgIndex::max_degree, R"doc(Graph degree bound.)doc" ) + .def_property_readonly( + "ef_construction", + &SymqgIndex::ef_construction, + R"doc(Construction search breadth of the current graph or the next add, or None for an index loaded from disk.)doc" + ) + .def_property_readonly( + "init", + &SymqgIndex::init, + R"doc(Graph initialization method of the current graph or the next add, or None for an index loaded from disk.)doc" + ) .def_property_readonly( "num_points", &SymqgIndex::num_points, @@ -374,16 +553,22 @@ index : SymqgIndex "max_elements", &SymqgIndex::max_elements, R"doc(Current storage capacity.)doc" ) .def_property_readonly( - "is_built", - &SymqgIndex::is_built, - R"doc(Whether the index has been built or loaded.)doc" + "is_trained", + [](const SymqgIndex&) { return true; }, + R"doc(Always True: SymphonyQG needs no training.)doc" ) .def_property_readonly( "metric", &SymqgIndex::metric, R"doc(Canonical metric name: 'l2' or 'ip'.)doc" ) .def_property_readonly( - "quantization_bits", - &SymqgIndex::quantization_bits, - R"doc(Storage bits: 0 for raw float32, or 4 or 8 for RaBitQ.)doc" + "nbits", + &SymqgIndex::nbits_property, + R"doc(Storage bits: 32 for raw float32, or 4 or 8 for RaBitQ.)doc" + ) + .def_property_readonly("_quantization_bits", &SymqgIndex::quantization_bits) + .def_property_readonly( + "seed", + &SymqgIndex::seed, + R"doc(Rotation and graph-construction seed, or None for an index loaded from disk.)doc" ); } diff --git a/sample/python/symqg_indexing.py b/sample/python/symqg_indexing.py index c09ce0cd..a0410b61 100644 --- a/sample/python/symqg_indexing.py +++ b/sample/python/symqg_indexing.py @@ -11,7 +11,7 @@ MAX_DEGREE = 32 # degree bound for SymphonyQG EF_CONSTRUCTION = 200 # ef for indexing METRIC = "l2" # "l2" or "ip" -QUANTIZATION_BITS = 0 # 0 for raw SymphonyQG, or 4/8 for quantized SymphonyQG +NBITS = 32 # 32 for raw SymphonyQG, or 4/8 for quantized SymphonyQG NUM_THREADS = 16 # number of threads for build INIT = "pipnn" # "pipnn" or "random" initialization of SymphonyQG # ────────────────────────────────────────────── @@ -22,28 +22,26 @@ def main(args=None) -> None: data = read_fvecs(args.data_file) print(f"Data shape: {data.shape}") - # 3. Build SymphonyQG index + # 2. Build the SymphonyQG graph; there is nothing to train. n, dim = data.shape print( f"\nBuilding SymphonyQG index: n={n}, dim={dim}, MaxDegree={args.max_degree}, " f"ef={args.ef_construction}, metric={args.metric}, " - f"quantization_bits={args.quantization_bits}, init={args.init}" + f"nbits={args.nbits}, init={args.init}" ) idx = SymqgIndex( - dim=dim, + dim, + nbits=args.nbits, max_degree=args.max_degree, + ef_construction=args.ef_construction, + init=args.init, metric=args.metric, - quantization_bits=args.quantization_bits, ) t0 = time() - idx.build( - data, - ef_construction=args.ef_construction, - num_threads=args.num_threads, - init=args.init, - ) + # add builds the whole graph at once; it cannot be extended later. + idx.add(data, num_threads=args.num_threads) print(f"Indexing time: {time() - t0:.2f}s") idx.save(args.index_file) @@ -79,12 +77,12 @@ def main(args=None) -> None: help="Distance metric (l2 or ip)", ) parser.add_argument( - "--quantization-bits", - dest="quantization_bits", + "--nbits", + dest="nbits", type=int, - choices=[0, 4, 8], - default=QUANTIZATION_BITS, - help="Vector quantization bits: 0 for raw SymphonyQG, or 4/8 for quantized SymphonyQG", + choices=[4, 8, 32], + default=NBITS, + help="Storage bits: 32 for raw SymphonyQG, or 4/8 for quantized SymphonyQG", ) parser.add_argument( "--init", diff --git a/sample/python/symqg_querying.py b/sample/python/symqg_querying.py index b837da7f..035282a4 100644 --- a/sample/python/symqg_querying.py +++ b/sample/python/symqg_querying.py @@ -37,7 +37,7 @@ def main(args=None) -> None: for i_probe, ef in enumerate(EFS): for r in range(args.test_rounds): t0 = perf_counter() - ids, _ = idx.search_batch( + ids, _ = idx.search( queries, k=args.topk, ef=ef, num_threads=args.num_threads ) elapsed = perf_counter() - t0 # seconds diff --git a/tests/python/test_compat_08.py b/tests/python/test_compat_08.py index 5cdcda14..288a5f35 100644 --- a/tests/python/test_compat_08.py +++ b/tests/python/test_compat_08.py @@ -6,7 +6,7 @@ import numpy as np import pytest from conftest import DIM, N_VECTORS -from rabitqlib import HnswIndex, IvfIndex +from rabitqlib import HnswIndex, IvfIndex, SymqgIndex def single_deprecation(call): @@ -212,3 +212,106 @@ def test_hnsw_build_on_a_trained_index_replaces_centroids( 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() + + +@pytest.mark.parametrize( + "make,bits", + [ + (lambda: SymqgIndex(DIM, 32), 0), + (lambda: SymqgIndex(DIM, 32, "ip", 4), 4), + (lambda: SymqgIndex(DIM, 32, quantization_bits=8), 8), + ], +) +def test_symqg_positional_form(make, bits): + index, message = single_deprecation(make) + assert message.startswith( + "SymqgIndex(dim, max_degree, metric, quantization_bits) is deprecated" + ) + assert index.nbits == (32 if bits == 0 else bits) + assert index.max_degree == 32 + assert (index.ef_construction, index.init) == (200, "pipnn") + + +@pytest.mark.parametrize("bits,nbits", [(0, 32), (4, 4), (8, 8)]) +def test_symqg_quantization_bits_keyword(bits, nbits, base_data, tmp_path): + index, message = single_deprecation( + lambda: SymqgIndex(dim=DIM, max_degree=32, quantization_bits=bits, seed=4) + ) + assert message.startswith("SymqgIndex(quantization_bits=...) is deprecated") + reference = SymqgIndex(DIM, nbits=nbits, seed=4) + for graph in (index, reference): + graph.add(base_data, num_threads=1) + index.save(tmp_path / "a.index") + reference.save(tmp_path / "b.index") + assert (tmp_path / "a.index").read_bytes() == (tmp_path / "b.index").read_bytes() + value, message = single_deprecation(lambda: index.quantization_bits) + assert value == bits + assert message.startswith("SymqgIndex.quantization_bits is deprecated") + + +@pytest.mark.parametrize( + "make,error,message", + [ + ( + lambda: SymqgIndex(DIM, nbits=4, quantization_bits=4), + TypeError, + "^SymqgIndex\\(\\) got both nbits and its old form quantization_bits$", + ), + ( + lambda: SymqgIndex(DIM, quantization_bits=2), + ValueError, + "^quantization_bits must be 0, 4, or 8$", + ), + ( + lambda: SymqgIndex(DIM, 32, nbits=4), + TypeError, + "got an unexpected keyword argument 'nbits' \\(read as the 0.7 form", + ), + ( + lambda: SymqgIndex(DIM, 32, "l2", 3), + ValueError, + "^quantization_bits must be 0, 4, or 8$", + ), + ], +) +def test_symqg_mixed_or_invalid_old_forms(make, error, message): + with pytest.raises(error, match=message): + quietly(make) + + +def test_symqg_build_search_batch_and_is_built(base_data, query_data): + index = quietly(lambda: SymqgIndex(DIM, 32)) + built, _ = single_deprecation(lambda: index.is_built) + assert built is False + _, message = single_deprecation(lambda: index.build(base_data, 50)) + assert message.startswith("SymqgIndex.build() is deprecated") + assert len(index) == N_VECTORS + built, _ = single_deprecation(lambda: index.is_built) + assert built is True + batch, message = single_deprecation(lambda: index.search_batch(query_data, 5, 40)) + assert message.startswith("SymqgIndex.search_batch() is deprecated; use search()") + expected = index.search(query_data, 5, 40) + np.testing.assert_array_equal(batch[0], expected[0]) + np.testing.assert_array_equal(batch[1], expected[1]) + # As in 0.7, build replaces the graph, even after add or another build. + quietly(lambda: index.build(base_data[:100], 50, 1, "random")) + assert len(index) == 100 + added = SymqgIndex(DIM) + added.add(base_data[:80]) + quietly(lambda: added.build(base_data, 50)) + assert len(added) == N_VECTORS + + +def test_symqg_build_draws_a_new_seed_unless_one_was_given(base_data): + unseeded = quietly(lambda: SymqgIndex(DIM, 32)) + seeds = set() + for init in ("pipnn", "random"): + quietly(lambda init=init: unseeded.build(base_data, 50, 1, init)) + seeds.add(unseeded.seed) + # The settings of the current graph. + assert (unseeded.ef_construction, unseeded.init) == (50, init) + assert len(seeds) == 2 + seeded = SymqgIndex(DIM, seed=3) + for _ in range(2): + quietly(lambda: seeded.build(base_data, 50)) + assert seeded.seed == 3 diff --git a/tests/python/test_editor_support.py b/tests/python/test_editor_support.py index eaccc5d1..eb14d8c4 100644 --- a/tests/python/test_editor_support.py +++ b/tests/python/test_editor_support.py @@ -196,8 +196,8 @@ def test_documented_search_result_types(kind): data = np.random.default_rng(42).normal(size=(80, 64)).astype(np.float32) if kind == "symqg": - index = rabitqlib.SymqgIndex(64, 32) - index.build(data, 40) + index = rabitqlib.SymqgIndex(64, ef_construction=40) + index.add(data) result = index.search(data[:2], 3, ef=20) else: if kind == "ivf": diff --git a/tests/python/test_examples.py b/tests/python/test_examples.py index df1e2795..27c19020 100644 --- a/tests/python/test_examples.py +++ b/tests/python/test_examples.py @@ -375,7 +375,11 @@ def test_clustering_default_requires_graph_compatible_cluster_count(monkeypatch) @pytest.mark.parametrize("metric", ["l2", "ip"]) -def test_symqg_examples_do_not_import_faiss(tmp_path, example_data, metric): +def test_symqg_examples_do_not_import_faiss( + tmp_path, example_data, metric, 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 run_example( "symqg_indexing.py", diff --git a/tests/python/test_index_repr.py b/tests/python/test_index_repr.py index 16c23005..667c5217 100644 --- a/tests/python/test_index_repr.py +++ b/tests/python/test_index_repr.py @@ -16,8 +16,8 @@ "num_clusters={clusters}, is_trained={built})" ), "symqg": ( - "SymqgIndex(dim=64, metric='ip', quantization_bits={bits}, " - "num_points={live}, max_degree=32, is_built={built})" + "SymqgIndex(dim=64, metric='ip', nbits={nbits}, " + "num_points={live}, max_degree=32)" ), } @@ -52,6 +52,7 @@ def test_repr_lifecycle(kind, bits, vectors, tmp_path): capacity = 96 if kind == "hnsw" else 80 before = SUMMARIES[kind].format( bits=bits, + nbits=32 if bits == 0 else bits, live=0, stored=0, capacity=capacity, @@ -62,7 +63,13 @@ def test_repr_lifecycle(kind, bits, vectors, tmp_path): build_index(index, vectors) after = SUMMARIES[kind].format( - bits=bits, live=80, stored=80, capacity=capacity, clusters=1, built=True + bits=bits, + nbits=32 if bits == 0 else bits, + live=80, + stored=80, + capacity=capacity, + clusters=1, + built=True, ) reference = index.search(vectors[:2], 3, 32) assert repr(index) == str(index) == after diff --git a/tests/python/test_symqg.py b/tests/python/test_symqg.py index a558de70..37e12468 100644 --- a/tests/python/test_symqg.py +++ b/tests/python/test_symqg.py @@ -1,5 +1,6 @@ """Tests for SymqgIndex: construction, search, properties, error handling, save/load.""" +import inspect import os import re import struct @@ -102,12 +103,8 @@ def test_search_batch_validates_inputs(built_symqg, query_data): for k in [0, N_VECTORS + 1]: with pytest.raises(ValueError, match="k must be between 1 and num_points"): built_symqg.search_batch(query_data, k=k, ef=_EF) - with pytest.raises(ValueError, match="ef must be positive"): - built_symqg.search_batch(query_data, k=1, ef=0) - with pytest.raises(ValueError, match="ef must be at least k"): - built_symqg.search_batch(query_data, k=5, ef=4) unbuilt = SymqgIndex(DIM, max_degree=_MAX_DEGREE) - with pytest.raises(RuntimeError, match="must be built or loaded before search"): + with pytest.raises(RuntimeError, match="add them or load an index before search"): unbuilt.search_batch(query_data, k=1, ef=_EF) @@ -164,7 +161,8 @@ def test_wrong_query_dim_raises(built_symqg): def test_save_before_build_raises(tmp_path): idx = SymqgIndex(DIM, max_degree=_MAX_DEGREE) with pytest.raises( - RuntimeError, match="^SymqgIndex must be built or loaded before save$" + RuntimeError, + match="^SymqgIndex has no vectors; add them or load an index before save$", ): idx.save(str(tmp_path / "unbuilt.index")) @@ -173,14 +171,24 @@ def test_search_before_build_raises(): idx = SymqgIndex(DIM, max_degree=_MAX_DEGREE) queries = np.zeros((5, DIM), dtype=np.float32) with pytest.raises( - RuntimeError, match="^SymqgIndex must be built or loaded before search$" + RuntimeError, + match="^SymqgIndex has no vectors; add them or load an index before search$", ): idx.search(queries, k=1, ef=_EF) -def test_search_rejects_ef_smaller_than_k(built_symqg, query_data): - with pytest.raises(ValueError, match="ef must be at least k"): - built_symqg.search(query_data, k=10, ef=9) +def test_search_raises_ef_to_k_and_defaults_it(built_symqg, query_data): + # As in HnswIndex: no or zero ef selects max(k, 10), and ef is at least k. + at_k = built_symqg.search(query_data, k=10, ef=10) + for ef in (9, 1): + raised = built_symqg.search(query_data, k=10, ef=ef) + np.testing.assert_array_equal(raised[0], at_k[0]) + default = built_symqg.search(query_data, k=3, ef=10) + for ef in (None, 0): + np.testing.assert_array_equal( + built_symqg.search(query_data, 3, ef)[0], default[0] + ) + np.testing.assert_array_equal(built_symqg.search(query_data, 3)[0], default[0]) @pytest.mark.parametrize("bits", [0, 4, 8]) @@ -396,5 +404,6 @@ def test_native_pipnn_rejects_invalid_build_arguments(base_data): def test_default_initialization_is_pipnn(): - # pybind11 generates this signature from the actual Python argument defaults. - assert "init: str = 'pipnn'" in SymqgIndex.build.__doc__ + assert inspect.signature(SymqgIndex.__init__).parameters["init"].default == "pipnn" + assert SymqgIndex(DIM).init == "pipnn" + assert inspect.signature(SymqgIndex.build).parameters["init"].default == "pipnn" diff --git a/tests/python/test_symqg_api.py b/tests/python/test_symqg_api.py new file mode 100644 index 00000000..8c98b0bc --- /dev/null +++ b/tests/python/test_symqg_api.py @@ -0,0 +1,188 @@ +"""The 0.8 train / add / search workflow of SymqgIndex.""" + +import inspect +import warnings + +import numpy as np +import pytest +from conftest import DIM, N_VECTORS +from rabitqlib import SymqgIndex + + +@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 saved_bytes(index, path): + index.save(path) + return path.read_bytes() + + +def test_constructor_options_are_keyword_only_with_defaults(): + parameters = inspect.signature(SymqgIndex.__init__).parameters + keyword = inspect.Parameter.KEYWORD_ONLY + assert [p for p in parameters if parameters[p].kind is not keyword] == [ + "self", + "dim", + ] + defaults = {name: p.default for name, p in parameters.items() if name != "self"} + assert defaults == { + "dim": inspect.Parameter.empty, + "nbits": 32, + "max_degree": 32, + "ef_construction": 200, + "init": "pipnn", + "metric": "l2", + "seed": None, + } + index = SymqgIndex(DIM) + assert (index.nbits, index.max_degree, index.ef_construction, index.init) == ( + 32, + 32, + 200, + "pipnn", + ) + assert index.is_trained + assert len(index) == 0 + + +@pytest.mark.parametrize( + "make,error,message", + [ + ( + lambda: SymqgIndex(10), + ValueError, + "^Unsupported dimension for FhtKacRotator$", + ), + ( + lambda: SymqgIndex(65537), + ValueError, + "^Unsupported dimension for FhtKacRotator$", + ), + (lambda: SymqgIndex(DIM, nbits=0), ValueError, "^nbits must be 4, 8, or 32$"), + (lambda: SymqgIndex(DIM, nbits=16), ValueError, "^nbits must be 4, 8, or 32$"), + ( + lambda: SymqgIndex(DIM, max_degree=48), + ValueError, + "^max_degree must be a positive multiple of 32$", + ), + ( + lambda: SymqgIndex(DIM, ef_construction=0), + ValueError, + "^ef_construction must be positive and fit uint32$", + ), + ( + lambda: SymqgIndex(DIM, init="pippn"), + ValueError, + "^init must be 'random' or 'pipnn'$", + ), + (lambda: SymqgIndex(DIM, seed=-1), ValueError, "seed must be an integer in"), + (lambda: SymqgIndex(DIM, seed=2**32), ValueError, r"\[0, 4294967295\]"), + ], +) +def test_constructor_argument_errors(make, error, message): + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + with pytest.raises(error, match=message): + make() + + +def test_seed_is_drawn_once_and_reproduces_the_index(base_data, tmp_path): + drawn = [SymqgIndex(DIM).seed for _ in range(4)] + assert all(isinstance(seed, int) and 0 <= seed < 2**32 for seed in drawn) + assert len(set(drawn)) > 1 + first = SymqgIndex(DIM, ef_construction=50) + second = SymqgIndex(DIM, ef_construction=50, seed=first.seed) + for index in (first, second): + index.add(base_data, num_threads=1) + assert saved_bytes(first, tmp_path / "a.index") == saved_bytes( + second, tmp_path / "b.index" + ) + + +def test_train_only_checks_its_input(base_data): + index = SymqgIndex(DIM) + index.train(base_data) + index.train(base_data, num_threads=1) + index.train(base_data, num_threads=None) + assert len(index) == 0 + bad = base_data.copy() + bad[3, 3] = np.nan + with pytest.raises(ValueError, match="^x must contain only finite values$"): + index.train(bad) + with pytest.raises(ValueError, match="x dimension does not match index dim"): + index.train(base_data[:, :-1]) + with pytest.raises(ValueError, match="x must be a 2D NumPy array"): + index.train(base_data[0]) + + +@pytest.mark.parametrize("nbits", [4, 8, 32]) +def test_add_builds_once_and_search_finds_the_vectors(nbits, base_data, tmp_path): + index = SymqgIndex(DIM, nbits=nbits, ef_construction=60, seed=5) + assert index.add(base_data) is None + assert len(index) == index.num_points == N_VECTORS + ids, distances = index.search(base_data[:20], 1) + assert np.mean(ids[:, 0] == np.arange(20)) >= 0.9 + message = "^SymqgIndex does not support incremental addition; create a new index and add all vectors at once$" + with pytest.raises(RuntimeError, match=message): + index.add(base_data[:40]) + assert len(index) == N_VECTORS + index.save(tmp_path / "graph.index") + loaded = SymqgIndex.load(tmp_path / "graph.index") + assert loaded.nbits == nbits + # Files store neither the seed nor the construction settings. + assert (loaded.seed, loaded.ef_construction, loaded.init) == (None, None, None) + assert loaded.is_trained + with pytest.raises(RuntimeError, match=message): + loaded.add(base_data) + + +def test_add_uses_the_constructor_settings(base_data, tmp_path): + for init in ("pipnn", "random"): + added = SymqgIndex(DIM, ef_construction=60, init=init, seed=9) + added.add(base_data, num_threads=1) + with pytest.warns(DeprecationWarning): + built = SymqgIndex(DIM, seed=9) + built.build(base_data, 60, 1, init) + assert saved_bytes(added, tmp_path / "a.index") == saved_bytes( + built, tmp_path / "b.index" + ) + + +def test_search_ef_defaults_to_max_of_k_and_10_and_is_raised_to_k(base_data): + index = SymqgIndex(DIM, ef_construction=60, seed=5) + index.add(base_data) + queries = base_data[:10] + for k, implied_ef in ((3, 10), (20, 20)): + expected = index.search(queries, k, implied_ef) + for ef in (None, 0): + actual = index.search(queries, k, ef) + np.testing.assert_array_equal(actual[0], expected[0]) + np.testing.assert_array_equal(actual[1], expected[1]) + np.testing.assert_array_equal( + index.search(queries, 5, 2)[0], index.search(queries, 5, 5)[0] + ) + + +def test_operations_before_add_raise(tmp_path, base_data): + index = SymqgIndex(DIM) + for operation, call in ( + ("search", lambda: index.search(base_data[:1], 1)), + ("save", lambda: index.save(tmp_path / "empty.index")), + ): + with pytest.raises( + RuntimeError, + match=f"^SymqgIndex has no vectors; add them or load an index before {operation}$", + ): + call() + assert not (tmp_path / "empty.index").exists() + + +def test_type_errors_list_only_the_public_constructor(): + with pytest.raises(TypeError) as error: + SymqgIndex(DIM, nbits="four") + assert "_quantization_bits" not in str(error.value) diff --git a/tests/python/test_train_api.py b/tests/python/test_train_api.py index 0f8a20e3..d5d92f50 100644 --- a/tests/python/test_train_api.py +++ b/tests/python/test_train_api.py @@ -351,6 +351,19 @@ def test_hnsw_capacity_grows_and_resize_preallocates(base_data): assert preallocated.max_elements == 1000 +def test_hnsw_search_ef_none_means_the_default(base_data): + index = HnswIndex(DIM, nbits=4, num_clusters=4) + index.train(base_data) + index.add(base_data, num_threads=1) + expected = index.search(base_data[:10], 5, 0) + for found in ( + index.search(base_data[:10], 5, None), + index.search(base_data[:10], 5), + ): + np.testing.assert_array_equal(found[0], expected[0]) + np.testing.assert_array_equal(found[1], expected[1]) + + def test_hnsw_failed_add_keeps_the_capacity(base_data): index = HnswIndex(DIM, nbits=4, num_clusters=4) index.train(base_data) diff --git a/tests/python/test_txtai_compat.py b/tests/python/test_txtai_compat.py index a005cb6b..91257889 100644 --- a/tests/python/test_txtai_compat.py +++ b/tests/python/test_txtai_compat.py @@ -1,9 +1,9 @@ """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. +Mirrors txtai.ann.dense.rabitq.RabitQ at neuml/txtai 4d5e05f (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 @@ -48,8 +48,21 @@ def index_like_txtai(mode, vectors): return backend -def search_like_txtai(backend, mode, queries, limit, deletes): - count = backend.max_elements - deletes +def append_like_txtai(backend, vectors): + backend.add(np.ascontiguousarray(vectors, dtype=np.float32)) + return backend.num_stored_points # txtai's next id offset + + +def delete_like_txtai(backend, ids): + ids = [uid for uid in ids if 0 <= uid < backend.num_stored_points] + if ids: + backend.remove(np.array(ids, dtype=np.int64)) + + +def search_like_txtai(backend, mode, queries, limit): + count = backend.num_points + if not count: + return [[] for _ in queries] k = min(limit, count) if mode == "ivf": nprobe = max(1, round(backend.num_clusters / 16)) @@ -71,25 +84,22 @@ 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) + results = search_like_txtai(backend, mode, embeddings[:5], 3) 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 append_like_txtai(backend, embeddings[300:]) == 400 + delete_like_txtai(backend, [0, 1, 399, 1000]) + assert backend.num_points == 397 + results = search_like_txtai(backend, mode, embeddings[300:305], 3) 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) + assert search_like_txtai(loaded, mode, embeddings[:5], 3) == ( + search_like_txtai(backend, mode, embeddings[:5], 3) + ) 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 91559127..81b5b561 100644 --- a/tests/python/typing/public_api.py +++ b/tests/python/typing/public_api.py @@ -25,7 +25,7 @@ def check(x: npt.NDArray[np.float32], labels: npt.NDArray[np.uint32]) -> None: assert_type(rabitqlib.__all__, list[str]) ivf = IvfIndex(64, 2, nbits=4, seed=7) hnsw = HnswIndex(64, nbits=4, num_clusters=2) - graph = SymqgIndex(64, 32) + graph = SymqgIndex(64, ef_construction=64) for index in (ivf, hnsw, graph): assert_type(index.__repr__(), str) for index in (ivf, hnsw): @@ -46,9 +46,13 @@ def check(x: npt.NDArray[np.float32], labels: npt.NDArray[np.uint32]) -> None: assert_type( hnsw.search(x, 5), tuple[npt.NDArray[np.uint32], npt.NDArray[np.float32]] ) - graph.build(x, 64) + graph.train(x) + graph.add(x, num_threads=1) + assert_type(graph.is_trained, bool) + assert_type(graph.nbits, int) + assert_type(len(graph), int) assert_type( - graph.search_batch(x, 5, 32), + graph.search(x, 5), tuple[npt.NDArray[np.uint32], npt.NDArray[np.float32]], ) assert_type(IvfIndex.load("index.bin"), IvfIndex) @@ -109,7 +113,13 @@ def check_numpy_scalars() -> None: 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)) + graph = SymqgIndex( + np.int64(64), + nbits=np.int64(4), + max_degree=np.int64(32), + ef_construction=np.int64(40), + seed=np.uint32(5), + ) for index in (ivf, hnsw): index.set_centroids(centroids) index.add( @@ -128,7 +138,7 @@ def check_numpy_scalars() -> None: assert_type(distances, npt.NDArray[np.float32]) hnsw.resize(np.int64(100)) hnsw.search(x[:2], np.int64(3), ef=np.int64(10)) - graph.build(x, np.int64(40), num_threads=np.int64(1)) + graph.add(x, num_threads=np.int64(1)) graph.search(x[:2], np.int64(3), np.int64(10)) flat = RaBitQKMeans( np.int64(64), @@ -183,7 +193,7 @@ def check_input_contracts() -> None: for metric in metrics: ivf = IvfIndex(64, 1, nbits=4, metric=metric) HnswIndex(64, nbits=4, metric=metric) - SymqgIndex(64, 32, metric=metric) + SymqgIndex(64, metric=metric) initializers: tuple[Literal["auto", "flat", "flat_rabitq", "hnsw"], ...] = ( "auto", "flat", @@ -212,8 +222,8 @@ def check_input_contracts() -> None: hnsw.remove(np.array([1], dtype=np.int64)) initializations: tuple[Literal["random", "pipnn"], ...] = ("random", "pipnn") for init in initializations: - graph = SymqgIndex(64, 32) - graph.build(x, 40, init=init) + graph = SymqgIndex(64, ef_construction=40, init=init) + graph.add(x) # Returned option values must be usable as inputs without casts. IvfIndex(64, 1, nbits=4, metric=ivf.metric, initializer=ivf.initializer) @@ -224,12 +234,13 @@ def reject_invalid_inputs(x: npt.NDArray[np.float32]) -> None: string_ids: npt.NDArray[np.str_] = np.array(["1"], dtype=np.str_) ivf = IvfIndex(64, 1, nbits=4) hnsw = HnswIndex(64, nbits=4) - graph = SymqgIndex(64, 32) + graph = SymqgIndex(64) 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] + SymqgIndex(64, metric="inner_product") # type: ignore[call-overload] IvfIndex(64, 1, nbits=4, initializer="flat_rabit") # type: ignore[call-overload] - graph.build(x, 40, init="pippn") # type: ignore[arg-type] + SymqgIndex(64, init="pippn") # type: ignore[call-overload] + graph.add(x, 1) # type: ignore[call-arg] 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] diff --git a/tests/test_mypy.py b/tests/test_mypy.py index 5b3ee1fe..430f2322 100644 --- a/tests/test_mypy.py +++ b/tests/test_mypy.py @@ -62,7 +62,7 @@ def test_static_public_usage(tmp_path): # stubs, so editors flag it while it still type-checks. DEPRECATED_USAGE = """\ import numpy as np -from rabitqlib import HnswIndex, IvfIndex +from rabitqlib import HnswIndex, IvfIndex, SymqgIndex x = np.zeros((80, 64), dtype=np.float32) labels = np.zeros(80, dtype=np.uint32) @@ -73,6 +73,12 @@ def test_static_public_usage(tmp_path): hnsw.build(x, x[:1], labels) ivf.add(x, labels) flag = ivf.is_built +graph = SymqgIndex(64, 32, "l2", 4) +alias = SymqgIndex(64, quantization_bits=0) +graph.build(x, 40) +graph.search_batch(x, 5, 32) +bits = graph.quantization_bits +built = graph.is_built """ @@ -102,4 +108,4 @@ def test_deprecated_forms_are_flagged(tmp_path): rf"{re.escape(str(usage))}:(\d+): error: .*\[deprecated\]", result.stdout ) } - assert flagged == set(range(6, 13)), result.stdout + assert flagged == set(range(6, 19)), result.stdout