diff --git a/Cargo.lock b/Cargo.lock index 938d3d3ff..f6d492298 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5100,6 +5100,15 @@ dependencies = [ "autocfg 1.5.0", ] +[[package]] +name = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg 1.5.0", +] + [[package]] name = "memory_units" version = "0.4.0" @@ -5716,7 +5725,7 @@ dependencies = [ "bitflags 1.3.2", "cfg-if 1.0.4", "libc", - "memoffset", + "memoffset 0.7.1", "pin-utils", ] @@ -6044,6 +6053,21 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" +[[package]] +name = "numpy" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "edb929bc0da91a4d85ed6c0a84deaa53d411abfb387fc271124f91bf6b89f14e" +dependencies = [ + "libc", + "ndarray 0.16.1", + "num-complex 0.4.6", + "num-integer", + "num-traits", + "pyo3", + "rustc-hash 1.1.0", +] + [[package]] name = "objc" version = "0.2.7" @@ -7239,6 +7263,69 @@ dependencies = [ "pest_derive", ] +[[package]] +name = "pyo3" +version = "0.22.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f402062616ab18202ae8319da13fa4279883a2b8a9d9f83f20dbade813ce1884" +dependencies = [ + "cfg-if 1.0.4", + "indoc", + "libc", + "memoffset 0.9.1", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", + "unindent", +] + +[[package]] +name = "pyo3-build-config" +version = "0.22.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b14b5775b5ff446dd1056212d778012cbe8a0fbffd368029fd9e25b514479c38" +dependencies = [ + "once_cell", + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.22.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ab5bcf04a2cdcbb50c7d6105de943f543f9ed92af55818fd17b660390fc8636" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.22.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fd24d897903a9e6d80b968368a34e1525aeb719d568dba8b3d4bfa5dc67d453" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn 2.0.117", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.22.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36c011a03ba1e50152b4b394b479826cad97e7a21eb52df179cd91ac411cbfbe" +dependencies = [ + "heck 0.5.0", + "proc-macro2", + "pyo3-build-config", + "quote", + "syn 2.0.117", +] + [[package]] name = "qoi" version = "0.4.1" @@ -9627,6 +9714,16 @@ dependencies = [ "tempfile", ] +[[package]] +name = "ruvector-py" +version = "0.1.0" +dependencies = [ + "numpy", + "pyo3", + "ruvector-rabitq", + "thiserror 2.0.18", +] + [[package]] name = "ruvector-rabitq" version = "2.2.0" @@ -11948,6 +12045,12 @@ dependencies = [ "xattr", ] +[[package]] +name = "target-lexicon" +version = "0.12.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" + [[package]] name = "tempfile" version = "3.27.0" @@ -12861,6 +12964,12 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" +[[package]] +name = "unindent" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" + [[package]] name = "universal-hash" version = "0.5.1" diff --git a/Cargo.toml b/Cargo.toml index 5c66aaf74..f28699b88 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -98,6 +98,8 @@ members = [ "examples/OSpipe", "crates/ruvector-coherence", "crates/ruvector-profiler", + # Python SDK — M1 (RaBitQ-only). See docs/sdk/04-milestones.md. + "crates/ruvector-py", "crates/ruvector-attn-mincut", "crates/ruvector-cognitive-container", "crates/ruvector-verified", diff --git a/crates/ruvector-py/Cargo.toml b/crates/ruvector-py/Cargo.toml new file mode 100644 index 000000000..698103001 --- /dev/null +++ b/crates/ruvector-py/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "ruvector-py" +version = "0.1.0" +edition = "2021" +rust-version.workspace = true +license = "MIT OR Apache-2.0" +authors.workspace = true +repository.workspace = true +description = "Python bindings for ruvector — vector similarity search via RaBitQ 1-bit quantization" + +[lib] +# Must match the maturin `module-name` Python module so the produced +# cdylib lands as `ruvector/_native..so`. See M1 in +# `docs/sdk/04-milestones.md`. +name = "ruvector_py" +crate-type = ["cdylib"] + +[dependencies] +# Pinned to 0.22 across both pyo3 and rust-numpy: the two crates are +# version-locked and a mismatch produces cryptic linker errors. abi3-py39 +# means one wheel covers Python 3.9..3.13 — see `docs/sdk/02-strategy.md` +# § "Wheel distribution matrix". +pyo3 = { version = "0.22", features = ["extension-module", "abi3-py39"] } +numpy = "0.22" +ruvector-rabitq = { path = "../ruvector-rabitq" } +thiserror = { workspace = true } diff --git a/crates/ruvector-py/README.md b/crates/ruvector-py/README.md new file mode 100644 index 000000000..bc764397a --- /dev/null +++ b/crates/ruvector-py/README.md @@ -0,0 +1,91 @@ +# ruvector — Python SDK (M1) + +Vector similarity search via RaBitQ 1-bit quantization, implemented in Rust +with native NumPy interop. M1 ships exactly one index class — +`RabitqIndex` — backed by `ruvector_rabitq::RabitqPlusIndex` (symmetric +1-bit scan + exact f32 rerank). + +This crate is the Python wheel half of the ruvector workspace; the +underlying algorithms live in `crates/ruvector-rabitq/` and are unchanged +by this binding. The full SDK plan (M1 → M4) is in +[`docs/sdk/`](../../docs/sdk/). + +## Install + +Once published to PyPI: + +```sh +pip install ruvector +``` + +For local development from a checkout: + +```sh +cd crates/ruvector-py +maturin develop --release +pytest tests/ +``` + +`maturin develop` builds the Rust cdylib in-place and links it as +`ruvector._native` so `import ruvector` works from any Python interpreter +in the active virtualenv. The `--release` flag matters: a debug build is +~30× slower on the search loop and will fail the latency acceptance test. + +## 30-second example + +```python +import numpy as np +import ruvector + +# Build an index over 100k random D=128 vectors. +rng = np.random.default_rng(42) +vectors = rng.standard_normal((100_000, 128), dtype=np.float32) +idx = ruvector.RabitqIndex.build(vectors, rerank_factor=20) + +# Search the 10 nearest neighbours of a query. +query = vectors[0] +hits = idx.search(query, k=10) +for vid, score in hits: + print(vid, score) +# 0 0.0 +# 12345 0.0023 +# ... + +# Persist and reload. +idx.save("idx.rbpx") +idx2 = ruvector.RabitqIndex.load("idx.rbpx") +assert idx2.search(query, k=10) == hits +``` + +## API summary + +| Call | Returns | Notes | +|---|---|---| +| `RabitqIndex.build(vectors, *, rerank_factor=20, seed=42)` | `RabitqIndex` | `vectors`: `(n, dim)` C-contig `float32` | +| `idx.search(query, k, *, rerank_factor=None)` | `list[(int, float)]` | `(id, score²)` ascending; `rerank_factor=None` uses the build value | +| `idx.save(path)` / `RabitqIndex.load(path)` | `None` / `RabitqIndex` | `.rbpx` v1 format | +| `len(idx)` / `idx.dim` / `idx.memory_bytes` / `idx.rerank_factor` | `int` | diagnostics | +| `ruvector.RuVectorError` | exception | base of the (future) error tree | +| `ruvector.__version__` | `str` | mirrors `Cargo.toml` | + +Non-contiguous or wrong-dtype inputs raise `TypeError` at the boundary +rather than silently copying — predictable beats fast. + +## Acceptance gates (M1) + +Per `docs/sdk/04-milestones.md`: + + 1. `pip install ruvector` (or `maturin develop`) succeeds in <10 s + 2. 100k-vector D=128 search returns in <10 ms (p99 over 100 queries) + 3. Type stubs validate with `mypy --strict` + +## Links + + - [SDK plan and milestones](../../docs/sdk/) — M1 through M4 roadmap + - [Binding strategy](../../docs/sdk/02-strategy.md) — why PyO3 + maturin + - [API surface sketch](../../docs/sdk/03-api-surface.md) — full Python surface + - [`ruvector-rabitq`](../ruvector-rabitq/) — the Rust crate this wraps + +## License + +Dual MIT / Apache-2.0, matching the rest of the ruvector workspace. diff --git a/crates/ruvector-py/pyproject.toml b/crates/ruvector-py/pyproject.toml new file mode 100644 index 000000000..5437d24f2 --- /dev/null +++ b/crates/ruvector-py/pyproject.toml @@ -0,0 +1,56 @@ +[build-system] +requires = ["maturin>=1.7,<2.0"] +build-backend = "maturin" + +[project] +name = "ruvector" +version = "0.1.0" +description = "Vector similarity search via RaBitQ 1-bit quantization" +readme = "README.md" +license = { text = "MIT OR Apache-2.0" } +requires-python = ">=3.9" +authors = [{ name = "Ruvector Team" }] +keywords = ["vector-search", "ann", "rabitq", "rust", "embeddings"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Programming Language :: Rust", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: Implementation :: CPython", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Database :: Database Engines/Servers", + "License :: OSI Approved :: MIT License", + "License :: OSI Approved :: Apache Software License", + "Operating System :: POSIX :: Linux", + "Operating System :: MacOS", + "Operating System :: Microsoft :: Windows", +] +dependencies = ["numpy>=1.21"] + +[project.optional-dependencies] +test = ["pytest>=7", "numpy>=1.21"] + +[project.urls] +Repository = "https://github.com/ruvnet/ruvector" +Issues = "https://github.com/ruvnet/ruvector/issues" +"SDK Plan" = "https://github.com/ruvnet/ruvector/tree/main/docs/sdk" + +[tool.maturin] +features = ["pyo3/extension-module"] +python-source = "python" +module-name = "ruvector._native" +# Hand-written stubs live alongside the Python source so they ship in the +# wheel. `python/ruvector/__init__.pyi` is the canonical surface; the +# `stubs/` tree carries the same file for tooling that reads PEP 561 stub +# packages directly. See `docs/sdk/02-strategy.md` § "Type stubs". +include = [ + { path = "python/ruvector/py.typed", format = "wheel" }, + { path = "python/ruvector/__init__.pyi", format = "wheel" }, +] + +[tool.pytest.ini_options] +testpaths = ["tests"] diff --git a/crates/ruvector-py/python/ruvector/__init__.py b/crates/ruvector-py/python/ruvector/__init__.py new file mode 100644 index 000000000..14031a9cd --- /dev/null +++ b/crates/ruvector-py/python/ruvector/__init__.py @@ -0,0 +1,10 @@ +"""ruvector — vector similarity search via RaBitQ 1-bit quantization. + +M1 surface only: a single ``RabitqIndex`` class plus the ``RuVectorError`` +base exception. See ``docs/sdk/04-milestones.md`` for what M2/M3/M4 add +(RuLake, Embedder, A2aClient). +""" + +from ruvector._native import RabitqIndex, RuVectorError, __version__ + +__all__ = ["RabitqIndex", "RuVectorError", "__version__"] diff --git a/crates/ruvector-py/python/ruvector/__init__.pyi b/crates/ruvector-py/python/ruvector/__init__.pyi new file mode 100644 index 000000000..b6973c48b --- /dev/null +++ b/crates/ruvector-py/python/ruvector/__init__.pyi @@ -0,0 +1,75 @@ +"""Type stubs for ruvector M1. + +Hand-written per docs/sdk/02-strategy.md § "Type stubs". Validates against +``mypy --strict`` and ``pyright``. +""" + +from typing import List, Optional, Tuple + +import numpy as np +from numpy.typing import NDArray + +__version__: str + +class RuVectorError(Exception): + """Base class for every error raised by the ruvector extension.""" + +class RabitqIndex: + """RaBitQ+ index — symmetric 1-bit scan with exact f32 rerank. + + Backed by ``ruvector_rabitq::RabitqPlusIndex``. Build with + :meth:`build`, query with :meth:`search`, persist via :meth:`save` / + :meth:`load`. + """ + + @staticmethod + def build( + vectors: NDArray[np.float32], + *, + rerank_factor: int = ..., + seed: int = ..., + ) -> "RabitqIndex": + """Build an index from an ``(n, dim)`` float32 array. + + ``vectors`` must be C-contiguous; non-contiguous arrays raise + ``TypeError``. ``rerank_factor`` defaults to 20 (the ADR-154 + recommendation for 100% recall@10 at D=128). ``seed`` defaults + to 42 for deterministic builds. + """ + ... + + def search( + self, + query: NDArray[np.float32], + k: int, + *, + rerank_factor: Optional[int] = ..., + ) -> List[Tuple[int, float]]: + """Search for the ``k`` nearest neighbours of ``query``. + + Returns a list of ``(id, score)`` tuples in ascending score + order (squared L2). ``rerank_factor=None`` (the default) reuses + the value the index was built with. + """ + ... + + def save(self, path: str) -> None: + """Persist the index to ``path`` in the ``.rbpx`` v1 format.""" + ... + + @staticmethod + def load(path: str) -> "RabitqIndex": + """Load an index previously written by :meth:`save`.""" + ... + + def __len__(self) -> int: ... + def __repr__(self) -> str: ... + + @property + def dim(self) -> int: ... + @property + def memory_bytes(self) -> int: ... + @property + def rerank_factor(self) -> int: ... + +__all__ = ["RabitqIndex", "RuVectorError", "__version__"] diff --git a/crates/ruvector-py/python/ruvector/py.typed b/crates/ruvector-py/python/ruvector/py.typed new file mode 100644 index 000000000..e69de29bb diff --git a/crates/ruvector-py/src/error.rs b/crates/ruvector-py/src/error.rs new file mode 100644 index 000000000..36370efc2 --- /dev/null +++ b/crates/ruvector-py/src/error.rs @@ -0,0 +1,27 @@ +//! Python exception hierarchy for `ruvector`. +//! +//! M1 ships a single user-visible exception, `RuVectorError`, plus the +//! `to_pyerr` mapper that converts every `RabitqError` variant into it. +//! Subclasses (`DimensionMismatch`, `EmptyIndex`, `PersistError`, …) are +//! reserved for M2/M3/M4 expansions — see `docs/sdk/03-api-surface.md` +//! § "Error hierarchy". For now the message string is the wire format. + +use pyo3::exceptions::PyException; +use pyo3::prelude::*; + +// `create_exception!` injects a unit-struct `RuVectorError` and a +// `RuVectorError::type_object_bound(py)` method we use in `lib.rs` when +// adding the symbol to the module. +pyo3::create_exception!( + ruvector._native, + RuVectorError, + PyException, + "Base class for every error raised by the ruvector extension." +); + +/// Map a `ruvector_rabitq::RabitqError` into a `PyErr` carrying +/// `RuVectorError`. The Display impl on `RabitqError` is already +/// human-readable so we forward it verbatim — no double-formatting. +pub fn to_pyerr(err: ruvector_rabitq::RabitqError) -> PyErr { + RuVectorError::new_err(err.to_string()) +} diff --git a/crates/ruvector-py/src/lib.rs b/crates/ruvector-py/src/lib.rs new file mode 100644 index 000000000..bb334e1b2 --- /dev/null +++ b/crates/ruvector-py/src/lib.rs @@ -0,0 +1,43 @@ +// pyo3 0.22's `create_exception!` and `#[pymethods]` macros emit cfg(feature = "gil-refs") +// gates and identity-`?` conversions on `PyResult` returns. Both are +// known false-positives against current rustc/clippy and are tracked +// upstream — silence them at the crate root rather than littering the +// source with #[allow] attrs. +#![allow(unexpected_cfgs)] +#![allow(clippy::useless_conversion)] + +//! ruvector — Python bindings (M1). +//! +//! Single PyO3 extension module exposed as `ruvector._native`. The maturin +//! `pyproject.toml` `module-name = "ruvector._native"` setting wires the +//! cdylib into the `ruvector` package at install time; the pure-Python +//! `python/ruvector/__init__.py` re-exports from `ruvector._native` so +//! end users only ever type `import ruvector`. +//! +//! M1 surface (per `docs/sdk/04-milestones.md`): +//! - `RabitqIndex` class (RaBitQ+ with rerank) +//! - `RuVectorError` exception +//! - `__version__` string mirroring the Cargo crate version +//! +//! Subsequent milestones add `RuLake`, `Embedder`, and `A2aClient` as +//! additional `register()` calls in this same `_native` module — no new +//! extensions, no separate wheels. + +use pyo3::prelude::*; + +mod error; +mod rabitq; + +#[pymodule] +fn _native(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { + // Class + exception registrations. + rabitq::register(m)?; + m.add("RuVectorError", py.get_type_bound::())?; + + // Version mirrors the Cargo crate version. The pure-Python + // `__init__.py` also re-exports it so `ruvector.__version__` works + // without import-time gymnastics. + m.add("__version__", env!("CARGO_PKG_VERSION"))?; + + Ok(()) +} diff --git a/crates/ruvector-py/src/rabitq.rs b/crates/ruvector-py/src/rabitq.rs new file mode 100644 index 000000000..fd89b3a32 --- /dev/null +++ b/crates/ruvector-py/src/rabitq.rs @@ -0,0 +1,264 @@ +//! `RabitqIndex` Python class — wraps `ruvector_rabitq::RabitqPlusIndex`. +//! +//! M1 of the SDK plan picks the **plus** variant as the user-facing default +//! because it is the only `ruvector-rabitq` index that: +//! 1. has a published persistence format (`.rbpx`, see +//! `crates/ruvector-rabitq/src/persist.rs`); +//! 2. supports per-call `rerank_factor` overrides via `search_with_rerank`; +//! 3. retains the original f32 vectors so we can serialise without forcing +//! the user to keep a separate `items` Vec on the Python side. +//! +//! Future milestones (M2+) can add `FlatF32Index`, `RabitqIndex` (no rerank), +//! and `RabitqAsymIndex` as separate Python classes per the surface sketch in +//! `docs/sdk/03-api-surface.md`. M1 ships exactly one class. + +use std::fs::File; +use std::io::{BufReader, BufWriter}; + +use numpy::{PyReadonlyArray1, PyReadonlyArray2, PyUntypedArrayMethods}; +use pyo3::exceptions::{PyIOError, PyTypeError, PyValueError}; +use pyo3::prelude::*; + +use ruvector_rabitq::index::AnnIndex; +use ruvector_rabitq::{persist, RabitqPlusIndex}; + +use crate::error::to_pyerr; + +/// Python-visible RaBitQ index. Backed by `RabitqPlusIndex` (symmetric 1-bit +/// scan + exact f32 rerank) — the variant the SDK plan picks for M1 because +/// it owns its originals and has a stable on-disk format. +/// +/// `unsendable` because the underlying index is `Send + Sync` but pyo3 cannot +/// statically prove our wrapper is — and we never need cross-thread Python +/// access since all heavy work is done inside `py.allow_threads`. Marking it +/// unsendable is the conservative, free choice. +#[pyclass(name = "RabitqIndex", module = "ruvector._native", unsendable)] +pub struct RabitqIndex { + inner: RabitqPlusIndex, + // RabitqPlusIndex needs the original (id, vector) items handed back to + // `persist::save_index` because the on-disk format is seed-based: + // `(dim, seed, rerank_factor, items)` deterministically rebuilds. We + // capture them at build time via `export_items()` on demand — no need + // to keep a redundant copy here. The seed lives separately because + // `RabitqPlusIndex` doesn't expose its own seed field. + seed: u64, +} + +#[pymethods] +impl RabitqIndex { + /// Build an index from a 2D NumPy array of shape `(n, dim)`. + /// + /// `vectors` must be `dtype=float32` and contiguous in C order. Non-contig + /// or wrong-dtype arrays raise `TypeError` at the boundary instead of + /// silently copying — silent copies would be an O(n·dim) surprise on a + /// "fast" build call. + /// + /// `rerank_factor` defaults to 20 per `docs/sdk/03-api-surface.md`'s + /// "100% recall@10 at D=128" recommendation citing ADR-154. + /// + /// `seed` defaults to 42 to keep the build deterministic out-of-the-box — + /// the doc-comment guarantee on `ruvector_rabitq` is that + /// `(dim, seed, data)` round-trips bit-identically. + #[staticmethod] + #[pyo3(signature = (vectors, *, rerank_factor = 20, seed = 42))] + fn build( + py: Python<'_>, + vectors: PyReadonlyArray2<'_, f32>, + rerank_factor: u32, + seed: u64, + ) -> PyResult { + // Validate the dtype/contig invariants up front. `PyReadonlyArray2` + // already guarantees the dtype is f32 (otherwise the conversion at + // the call site fails with `TypeError`). What it does NOT guarantee + // is C-contiguity; we enforce it because the inner Rust API takes + // owned `Vec` per row and we want a single slice per row, not + // a strided view that would force a copy through ndarray. + if !vectors.is_c_contiguous() { + return Err(PyTypeError::new_err( + "vectors must be C-contiguous; pass np.ascontiguousarray(...) first", + )); + } + let shape = vectors.shape(); + if shape.len() != 2 { + return Err(PyValueError::new_err(format!( + "vectors must be 2D, got {}D", + shape.len() + ))); + } + let n = shape[0]; + let dim = shape[1]; + if dim == 0 { + return Err(PyValueError::new_err("dim must be > 0")); + } + if n == 0 { + return Err(PyValueError::new_err("vectors must contain at least 1 row")); + } + + // Materialise (id, Vec) pairs. We allocate `n` rows of + // `dim` f32s. The inner `from_vectors_parallel` is the only ctor + // that uses rayon for the rotate+pack phase, so this is the + // path with the right scaling characteristics for a 100k+ build. + // + // Copy is unavoidable — the inner API takes owned Vecs to amortise + // across rayon workers. PyO3 cannot give us mutable owned access + // to NumPy storage without breaking the readonly contract. + let slice = vectors.as_slice()?; // contiguous view, len = n*dim + let mut items: Vec<(usize, Vec)> = Vec::with_capacity(n); + for i in 0..n { + let row = &slice[i * dim..(i + 1) * dim]; + items.push((i, row.to_vec())); + } + + // Heavy work: drop the GIL. `from_vectors_parallel` runs rotate+pack + // on a rayon thread pool; releasing the GIL lets a cooperating + // Python thread (e.g. the asyncio loop in M2) keep moving while + // the build completes. + let inner = py + .allow_threads(|| { + RabitqPlusIndex::from_vectors_parallel(dim, seed, rerank_factor as usize, items) + }) + .map_err(to_pyerr)?; + + Ok(Self { inner, seed }) + } + + /// Search for the `k` nearest neighbours of a single `query` vector. + /// + /// Returns a list of `(id, score)` tuples in ascending score order + /// (squared L2 — lower is closer). The id is the row index used at + /// `build` time; M1 doesn't expose user-supplied ids yet (M2 adds + /// them via `RuLake.upsert(ids=...)`). + /// + /// `rerank_factor` defaults to `None` which means "use the value the + /// index was built with" (see `RabitqPlusIndex::rerank_factor`); pass + /// an int to override per-call without rebuilding. + #[pyo3(signature = (query, k, *, rerank_factor = None))] + fn search( + &self, + py: Python<'_>, + query: PyReadonlyArray1<'_, f32>, + k: usize, + rerank_factor: Option, + ) -> PyResult> { + if !query.is_c_contiguous() { + return Err(PyTypeError::new_err( + "query must be C-contiguous; pass np.ascontiguousarray(...) first", + )); + } + if k == 0 { + return Err(PyValueError::new_err("k must be > 0")); + } + let q = query.as_slice()?; + + // `search_with_rerank` is the right entry point even when no + // override is requested, because it is the common path the + // benchmark uses; routing both into the same code keeps behaviour + // identical between defaulted and explicit calls. + let rf = rerank_factor + .map(|x| x as usize) + .unwrap_or_else(|| self.inner.rerank_factor()); + + // GIL-release window — pure Rust, no PyObject touched inside. + let results = py + .allow_threads(|| self.inner.search_with_rerank(q, k, rf)) + .map_err(to_pyerr)?; + + // Cast usize id → u32 to match the underlying SoA storage and + // the persisted on-disk id width (see `persist.rs`'s u32 id + // narrowing). `as u32` is safe because the inner storage already + // checked the bound at build time. + Ok(results + .into_iter() + .map(|r| (r.id as u32, r.score)) + .collect()) + } + + /// Save the index to `path` using the `.rbpx` v1 format. + fn save(&self, path: &str) -> PyResult<()> { + let items = self.inner.export_items(); + let f = + File::create(path).map_err(|e| PyIOError::new_err(format!("create {path}: {e}")))?; + let mut w = BufWriter::new(f); + persist::save_index(&self.inner, self.seed, &items, &mut w).map_err(to_pyerr)?; + // BufWriter::flush is implicit on drop, but make IO errors explicit + // here rather than swallowed in Drop. + use std::io::Write as _; + w.flush() + .map_err(|e| PyIOError::new_err(format!("flush {path}: {e}")))?; + Ok(()) + } + + /// Load an index previously written by `save`. Returns a fresh + /// `RabitqIndex`. The seed embedded in the file is recovered, so + /// subsequent `save()` calls round-trip without the caller juggling it. + #[staticmethod] + fn load(path: &str) -> PyResult { + let f = File::open(path).map_err(|e| PyIOError::new_err(format!("open {path}: {e}")))?; + let mut r = BufReader::new(f); + // We re-read the header to recover the seed since `RabitqPlusIndex` + // doesn't expose it. The cheapest correct path is to read+rewind: + // 8B magic + 4B version + 4B dim + 8B seed = 24 bytes. But since + // `persist::load_index` consumes the whole reader, we instead + // re-open and peek separately. + use std::io::Read as _; + let mut header = [0u8; 24]; + let mut peek = + File::open(path).map_err(|e| PyIOError::new_err(format!("open {path}: {e}")))?; + peek.read_exact(&mut header) + .map_err(|e| PyIOError::new_err(format!("read header from {path}: {e}")))?; + if &header[0..8] != persist::MAGIC { + return Err(PyIOError::new_err(format!( + "{path}: bad magic — not an rbpx file" + ))); + } + let seed = u64::from_le_bytes(header[16..24].try_into().unwrap()); + + let inner = persist::load_index(&mut r).map_err(to_pyerr)?; + Ok(Self { inner, seed }) + } + + /// Number of indexed vectors (matches `AnnIndex::len`). + fn __len__(&self) -> usize { + self.inner.len() + } + + /// Vector dimensionality. + #[getter] + fn dim(&self) -> usize { + self.inner.dim() + } + + /// Honest memory footprint in bytes — see + /// `ruvector_rabitq::AnnIndex::memory_bytes` for what's counted. + #[getter] + fn memory_bytes(&self) -> usize { + self.inner.memory_bytes() + } + + /// The rerank factor the index was built with (the default used by + /// `search` when no override is passed). + #[getter] + fn rerank_factor(&self) -> u32 { + self.inner.rerank_factor() as u32 + } + + /// Diagnostic-friendly repr: variant, n, dim, memory_bytes, rerank_factor. + fn __repr__(&self) -> String { + format!( + "RabitqIndex(n={}, dim={}, memory_bytes={}, rerank_factor={})", + self.inner.len(), + self.inner.dim(), + self.inner.memory_bytes(), + self.inner.rerank_factor(), + ) + } +} + +/// Convenience exporter — the module init in `lib.rs` calls this to add +/// the class. Keeping it here means `lib.rs` doesn't need to know the +/// class type by name, which mirrors the NAPI module conventions used in +/// `crates/ruvector-diskann-node`. +pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + Ok(()) +} diff --git a/crates/ruvector-py/stubs/ruvector/__init__.pyi b/crates/ruvector-py/stubs/ruvector/__init__.pyi new file mode 100644 index 000000000..cbceec7f9 --- /dev/null +++ b/crates/ruvector-py/stubs/ruvector/__init__.pyi @@ -0,0 +1,57 @@ +"""Type stubs for the `ruvector` package — mirrors python/ruvector/__init__.pyi. + +PEP 561 has two valid stub layouts: stubs alongside the package (with +``py.typed`` and ``__init__.pyi``) or a separate stub-only package. We +ship both: the wheel includes ``python/ruvector/__init__.pyi`` for +in-package consumption, and this tree carries the same stubs for +tooling that reads PEP 561 stub-only packages directly. + +Keep this file byte-identical (modulo this docstring) to +``python/ruvector/__init__.pyi``. A CI lint enforces the equivalence. +""" + +from typing import List, Optional, Tuple + +import numpy as np +from numpy.typing import NDArray + +__version__: str + +class RuVectorError(Exception): + """Base class for every error raised by the ruvector extension.""" + +class RabitqIndex: + """RaBitQ+ index — symmetric 1-bit scan with exact f32 rerank.""" + + @staticmethod + def build( + vectors: NDArray[np.float32], + *, + rerank_factor: int = ..., + seed: int = ..., + ) -> "RabitqIndex": ... + + def search( + self, + query: NDArray[np.float32], + k: int, + *, + rerank_factor: Optional[int] = ..., + ) -> List[Tuple[int, float]]: ... + + def save(self, path: str) -> None: ... + + @staticmethod + def load(path: str) -> "RabitqIndex": ... + + def __len__(self) -> int: ... + def __repr__(self) -> str: ... + + @property + def dim(self) -> int: ... + @property + def memory_bytes(self) -> int: ... + @property + def rerank_factor(self) -> int: ... + +__all__ = ["RabitqIndex", "RuVectorError", "__version__"] diff --git a/crates/ruvector-py/tests/test_smoke.py b/crates/ruvector-py/tests/test_smoke.py new file mode 100644 index 000000000..a8fd92f1d --- /dev/null +++ b/crates/ruvector-py/tests/test_smoke.py @@ -0,0 +1,120 @@ +"""M1 smoke tests for ``ruvector``. + +These exercise the user-visible surface of the wheel: + + - ``ruvector.__version__`` is set + - ``RabitqIndex.build`` accepts an ``(n, dim)`` float32 NumPy array + - ``RabitqIndex.search`` returns ``k`` ``(id, score)`` tuples + - first-result self-search returns id 0 at distance ~0 + - dimension mismatch raises ``RuVectorError`` + - save/load roundtrip preserves search results + +Run via ``pytest tests/`` after ``maturin develop`` (see README). +""" + +from __future__ import annotations + +import os +import tempfile + +import numpy as np +import pytest + +import ruvector + + +def test_version() -> None: + assert ruvector.__version__ + # Cargo.toml ships 0.1.0; if you bump there, bump here. + assert ruvector.__version__ == "0.1.0" + + +def test_build_and_search() -> None: + rng = np.random.default_rng(42) + n, dim = 1000, 128 + vectors = rng.standard_normal((n, dim), dtype=np.float32) + idx = ruvector.RabitqIndex.build(vectors) + + assert len(idx) == n + assert idx.dim == dim + assert idx.memory_bytes > 0 + assert idx.rerank_factor == 20 # default + + query = vectors[0] + results = idx.search(query, k=10) + + assert len(results) == 10 + # First hit must be the query vector itself: id 0, distance ~0 after + # exact f32 rerank. + top_id, top_dist = results[0] + assert top_id == 0, f"expected self-match at id 0, got id {top_id}" + assert top_dist < 1e-3, f"self-distance {top_dist} should be ~0" + + # Scores must be ascending (squared L2 — lower is closer). + scores = [s for _, s in results] + assert scores == sorted(scores), f"scores not ascending: {scores}" + + +def test_repr_is_diagnostic() -> None: + rng = np.random.default_rng(0) + vectors = rng.standard_normal((50, 32), dtype=np.float32) + idx = ruvector.RabitqIndex.build(vectors) + r = repr(idx) + assert "RabitqIndex" in r + assert "n=50" in r + assert "dim=32" in r + + +def test_error_on_dim_mismatch() -> None: + rng = np.random.default_rng(42) + vectors = rng.standard_normal((100, 64), dtype=np.float32) + idx = ruvector.RabitqIndex.build(vectors) + bad_query = rng.standard_normal(32, dtype=np.float32) + with pytest.raises(ruvector.RuVectorError): + idx.search(bad_query, k=10) + + +def test_error_on_wrong_dtype() -> None: + # float64 must not silently coerce — it should hit the boundary + # PyO3 numpy crate's strict dtype check. + rng = np.random.default_rng(0) + vectors = rng.standard_normal((10, 8)) # float64 + with pytest.raises((TypeError, ValueError)): + ruvector.RabitqIndex.build(vectors) # type: ignore[arg-type] + + +def test_save_load_roundtrip() -> None: + rng = np.random.default_rng(7) + n, dim = 200, 64 + vectors = rng.standard_normal((n, dim), dtype=np.float32) + idx = ruvector.RabitqIndex.build(vectors, rerank_factor=5, seed=1234) + + query = rng.standard_normal(dim, dtype=np.float32) + before = idx.search(query, k=5) + + with tempfile.TemporaryDirectory() as td: + path = os.path.join(td, "idx.rbpx") + idx.save(path) + loaded = ruvector.RabitqIndex.load(path) + + assert len(loaded) == n + assert loaded.dim == dim + assert loaded.rerank_factor == 5 + + after = loaded.search(query, k=5) + # `(dim, seed, items)` deterministic rebuild → bit-identical search. + assert before == after, f"roundtrip changed results: {before} vs {after}" + + +def test_search_with_per_call_rerank() -> None: + rng = np.random.default_rng(99) + n, dim = 500, 64 + vectors = rng.standard_normal((n, dim), dtype=np.float32) + idx = ruvector.RabitqIndex.build(vectors, rerank_factor=2) + + query = vectors[10] + + # Override per call — should still self-match at id 10 with distance ~0. + results = idx.search(query, k=3, rerank_factor=20) + assert results[0][0] == 10 + assert results[0][1] < 1e-3