mirror of
https://github.com/ruvnet/RuVector.git
synced 2026-05-23 12:55:26 +00:00
feat(ruvector-rabitq-wasm): WASM bindings for RaBitQ via wasm-bindgen
Closes the WASM gap from `docs/research/rabitq-integration/` Tier 2
("WASM / edge: 32× compression makes on-device RAG feasible") and
ADR-157 ("VectorKernel WASM kernel as a Phase 2 goal"). Adds a
`ruvector-rabitq-wasm` sibling crate that exposes `RabitqIndex` to
JavaScript/TypeScript callers (browsers, Cloudflare Workers, Deno,
Bun) via wasm-bindgen.
```js
import init, { RabitqIndex } from "ruvector-rabitq";
await init();
const dim = 768;
const n = 10_000;
const vectors = new Float32Array(n * dim); // populate
const idx = RabitqIndex.build(vectors, dim, 42, 20);
const query = new Float32Array(dim);
const results = idx.search(query, 10); // [{id, distance}, ...]
```
## Surface
- `RabitqIndex.build(vectors: Float32Array, dim, seed, rerank_factor)`
- `idx.search(query: Float32Array, k) → SearchResult[]`
- `idx.len`, `idx.isEmpty`
- `version()` — crate version baked at build time
- `SearchResult { id: u32, distance: f32 }` — mirrors the Python SDK
(PR #381) shape so callers porting code between languages get
identical structures.
## Native compatibility tweak
`ruvector-rabitq` had one rayon call site in
`from_vectors_parallel_with_rotation`. WASM is single-threaded — gated
that path on `cfg(not(target_arch = "wasm32"))` with a sequential
`.into_iter()` fallback for wasm. Output is bit-identical because the
rotation matrix is deterministic (ADR-154); parallel ordering doesn't
affect bytes.
`rayon` is now `[target.'cfg(not(target_arch = "wasm32"))'.dependencies]`
so the wasm build doesn't pull it in. Native build behavior unchanged
(39 / 39 lib tests still pass).
## Crate layout
crates/ruvector-rabitq-wasm/
Cargo.toml cdylib + rlib, wasm-bindgen 0.2, abi-3-friendly
src/lib.rs ~150 LoC of bindings; tests gated to wasm32 via
wasm_bindgen_test (native test would panic in
wasm-bindgen 0.2.117's runtime stub).
## Testing strategy
Native tests of WASM bindings panic by design — `JsValue::from_str`
calls into a wasm-bindgen runtime stub that's `unimplemented!()` on
non-wasm32 targets (since 0.2.117). The right path is
`wasm-pack test --node` or `wasm-pack test --headless --chrome`,
which we'll wire into CI as a follow-up.
The numerical correctness is already covered by `ruvector-rabitq`'s
own test suite. This crate only adds the JS-facing surface.
## Verification (native)
cargo build --workspace → 0 errors
cargo build -p ruvector-rabitq-wasm → clean
cargo clippy -p ruvector-rabitq-wasm --all-targets --no-deps -- -D warnings → exit 0
cargo test -p ruvector-rabitq → 39 / 39 (unchanged)
cargo fmt --all --check → clean
WASM target build (`wasm32-unknown-unknown`) requires `rustup target
add wasm32-unknown-unknown` — not exercised in this PR; will be
covered by a follow-up CI job.
Refs: docs/research/rabitq-integration/ Tier 2, ADR-157
("Optional Accelerator Plane"), PR #381 (Python SDK shape mirror).
Co-Authored-By: claude-flow <ruv@ruv.net>
This commit is contained in:
parent
208c1439cd
commit
a674d6ebae
6 changed files with 275 additions and 6 deletions
14
Cargo.lock
generated
14
Cargo.lock
generated
|
|
@ -9615,6 +9615,20 @@ dependencies = [
|
|||
"thiserror 2.0.18",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ruvector-rabitq-wasm"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"console_error_panic_hook",
|
||||
"getrandom 0.2.17",
|
||||
"js-sys",
|
||||
"ruvector-rabitq",
|
||||
"serde",
|
||||
"serde-wasm-bindgen",
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-test",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ruvector-raft"
|
||||
version = "2.2.0"
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ exclude = ["crates/micro-hnsw-wasm", "crates/ruvector-hyperbolic-hnsw", "crates/
|
|||
"crates/ruvector-postgres"]
|
||||
members = [
|
||||
"crates/ruvector-rabitq",
|
||||
"crates/ruvector-rabitq-wasm",
|
||||
"crates/ruvector-rulake",
|
||||
"crates/ruvector-core",
|
||||
"crates/ruvector-node",
|
||||
|
|
|
|||
47
crates/ruvector-rabitq-wasm/Cargo.toml
Normal file
47
crates/ruvector-rabitq-wasm/Cargo.toml
Normal file
|
|
@ -0,0 +1,47 @@
|
|||
[package]
|
||||
name = "ruvector-rabitq-wasm"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
description = "WASM bindings for ruvector-rabitq — 1-bit quantized vector index for browsers and edge runtimes"
|
||||
license = "MIT OR Apache-2.0"
|
||||
repository = "https://github.com/ruvnet/ruvector"
|
||||
keywords = ["rabitq", "vector-search", "wasm", "quantization", "embeddings"]
|
||||
categories = ["wasm", "science", "algorithms"]
|
||||
|
||||
[package.metadata.wasm-pack.profile.release]
|
||||
wasm-opt = false
|
||||
|
||||
[lib]
|
||||
crate-type = ["cdylib", "rlib"]
|
||||
|
||||
[features]
|
||||
default = ["console_error_panic_hook"]
|
||||
|
||||
[dependencies]
|
||||
ruvector-rabitq = { path = "../ruvector-rabitq" }
|
||||
wasm-bindgen = "0.2"
|
||||
js-sys = "0.3"
|
||||
console_error_panic_hook = { version = "0.1", optional = true }
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
serde-wasm-bindgen = "0.6"
|
||||
|
||||
[target.'cfg(target_arch = "wasm32")'.dependencies]
|
||||
getrandom = { version = "0.2", features = ["js"] }
|
||||
|
||||
[dev-dependencies]
|
||||
wasm-bindgen-test = "0.3"
|
||||
|
||||
[profile.release]
|
||||
opt-level = "s"
|
||||
lto = true
|
||||
|
||||
# Workspace cleanup pass: research-tier crate, doc/style churn deferred.
|
||||
# Correctness + suspicious lints stay denied.
|
||||
[lints.rust]
|
||||
unexpected_cfgs = { level = "allow", priority = -1 }
|
||||
|
||||
[lints.clippy]
|
||||
pedantic = { level = "allow", priority = -2 }
|
||||
all = { level = "warn", priority = -1 }
|
||||
correctness = "deny"
|
||||
suspicious = "deny"
|
||||
188
crates/ruvector-rabitq-wasm/src/lib.rs
Normal file
188
crates/ruvector-rabitq-wasm/src/lib.rs
Normal file
|
|
@ -0,0 +1,188 @@
|
|||
//! WASM bindings for ruvector-rabitq.
|
||||
//!
|
||||
//! Exposes [`RabitqIndex`] as a JavaScript-friendly class for use in
|
||||
//! browsers and edge runtimes (Cloudflare Workers, Deno, Bun).
|
||||
//! Single-threaded — the underlying `from_vectors_parallel` falls back
|
||||
//! to sequential iteration on wasm32 (output is bit-identical because
|
||||
//! rotation is deterministic).
|
||||
//!
|
||||
//! ```ignore
|
||||
//! import init, { RabitqIndex } from "ruvector-rabitq";
|
||||
//! await init();
|
||||
//!
|
||||
//! const dim = 768;
|
||||
//! const n = 10_000;
|
||||
//! const vectors = new Float32Array(n * dim); // populate
|
||||
//! const idx = RabitqIndex.build(vectors, dim, 42, 20);
|
||||
//! const query = new Float32Array(dim); // populate
|
||||
//! const results = idx.search(query, 10); // [{id, distance}, ...]
|
||||
//! ```
|
||||
|
||||
#![allow(clippy::new_without_default)]
|
||||
|
||||
use ruvector_rabitq::{AnnIndex, RabitqPlusIndex};
|
||||
use wasm_bindgen::prelude::*;
|
||||
|
||||
/// Initialize panic hook for clearer error messages in the browser
|
||||
/// console. Called once at module import.
|
||||
#[wasm_bindgen(start)]
|
||||
pub fn init() {
|
||||
#[cfg(feature = "console_error_panic_hook")]
|
||||
console_error_panic_hook::set_once();
|
||||
}
|
||||
|
||||
/// Search result — single nearest-neighbor hit.
|
||||
///
|
||||
/// Mirrors the structure used by the Python SDK's `RabitqIndex.search`
|
||||
/// so callers porting code between languages get identical shapes.
|
||||
#[wasm_bindgen]
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct SearchResult {
|
||||
/// Caller-supplied vector id (the position passed to `build`).
|
||||
#[wasm_bindgen(readonly)]
|
||||
pub id: u32,
|
||||
/// Approximate L2² distance after RaBitQ rerank.
|
||||
#[wasm_bindgen(readonly)]
|
||||
pub distance: f32,
|
||||
}
|
||||
|
||||
/// 1-bit quantized vector index. Builds in O(n × dim) memory + O(n × dim)
|
||||
/// time; searches in O(n) hamming distance + O(rerank_factor × k × dim)
|
||||
/// exact-L2² rerank.
|
||||
#[wasm_bindgen]
|
||||
pub struct RabitqIndex {
|
||||
inner: RabitqPlusIndex,
|
||||
}
|
||||
|
||||
#[wasm_bindgen]
|
||||
impl RabitqIndex {
|
||||
/// Build an index from a flat Float32Array of length `n * dim`.
|
||||
///
|
||||
/// `seed` controls the random rotation matrix; the same `(seed,
|
||||
/// dim, vectors)` triple produces bit-identical codes (ADR-154
|
||||
/// determinism guarantee). `rerank_factor` is the multiplier on
|
||||
/// `k` for the exact-L2² rerank pool — typical 20.
|
||||
///
|
||||
/// Errors:
|
||||
/// - `vectors.length` is not a multiple of `dim`
|
||||
/// - `dim == 0` or `vectors.length == 0`
|
||||
#[wasm_bindgen]
|
||||
pub fn build(
|
||||
vectors: &[f32],
|
||||
dim: u32,
|
||||
seed: u64,
|
||||
rerank_factor: u32,
|
||||
) -> Result<RabitqIndex, JsValue> {
|
||||
let dim = dim as usize;
|
||||
if dim == 0 {
|
||||
return Err(JsValue::from_str("dim must be > 0"));
|
||||
}
|
||||
if vectors.is_empty() {
|
||||
return Err(JsValue::from_str("vectors must not be empty"));
|
||||
}
|
||||
if !vectors.len().is_multiple_of(dim) {
|
||||
return Err(JsValue::from_str(&format!(
|
||||
"vectors length {} is not a multiple of dim {}",
|
||||
vectors.len(),
|
||||
dim
|
||||
)));
|
||||
}
|
||||
|
||||
let n = vectors.len() / dim;
|
||||
let items: Vec<(usize, Vec<f32>)> = (0..n)
|
||||
.map(|i| (i, vectors[i * dim..(i + 1) * dim].to_vec()))
|
||||
.collect();
|
||||
|
||||
let inner =
|
||||
RabitqPlusIndex::from_vectors_parallel(dim, seed, rerank_factor as usize, items)
|
||||
.map_err(|e| JsValue::from_str(&format!("RabitqIndex.build: {e}")))?;
|
||||
|
||||
Ok(Self { inner })
|
||||
}
|
||||
|
||||
/// Find the `k` nearest neighbors of `query`. Returns hits in
|
||||
/// ascending distance.
|
||||
///
|
||||
/// Errors:
|
||||
/// - `query.length != dim` of the index
|
||||
/// - `k == 0`
|
||||
#[wasm_bindgen]
|
||||
pub fn search(&self, query: &[f32], k: u32) -> Result<Vec<SearchResult>, JsValue> {
|
||||
if k == 0 {
|
||||
return Err(JsValue::from_str("k must be > 0"));
|
||||
}
|
||||
let hits = self
|
||||
.inner
|
||||
.search(query, k as usize)
|
||||
.map_err(|e| JsValue::from_str(&format!("RabitqIndex.search: {e}")))?;
|
||||
|
||||
Ok(hits
|
||||
.into_iter()
|
||||
.map(|h| SearchResult {
|
||||
id: h.id as u32,
|
||||
distance: h.score,
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
/// Number of vectors indexed.
|
||||
#[wasm_bindgen(getter)]
|
||||
pub fn len(&self) -> u32 {
|
||||
self.inner.len() as u32
|
||||
}
|
||||
|
||||
/// True iff the index has zero vectors. Mirrors Rust's `is_empty`
|
||||
/// convention; exposed because `wasm-bindgen` getter for `len`
|
||||
/// returns u32, so callers can't `idx.len === 0` reliably.
|
||||
#[wasm_bindgen(getter, js_name = isEmpty)]
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.inner.len() == 0
|
||||
}
|
||||
}
|
||||
|
||||
/// Crate version string baked at build time.
|
||||
#[wasm_bindgen(js_name = version)]
|
||||
pub fn version() -> String {
|
||||
env!("CARGO_PKG_VERSION").to_string()
|
||||
}
|
||||
|
||||
// Tests for the WASM bindings live as `wasm_bindgen_test` and only run
|
||||
// in a wasm32 environment via `wasm-pack test`. Native tests can't
|
||||
// exercise the bindings because `wasm-bindgen 0.2.117` panics on
|
||||
// `JsValue::from_str` outside a wasm runtime.
|
||||
//
|
||||
// The inner numerical correctness is covered by `ruvector-rabitq`'s
|
||||
// own test suite; here we only verify the JS-facing surface.
|
||||
#[cfg(all(test, target_arch = "wasm32"))]
|
||||
mod wasm_tests {
|
||||
use super::*;
|
||||
use wasm_bindgen_test::*;
|
||||
|
||||
wasm_bindgen_test_configure!(run_in_browser);
|
||||
|
||||
#[wasm_bindgen_test]
|
||||
fn build_and_search() {
|
||||
let dim = 32usize;
|
||||
let n = 100usize;
|
||||
let mut vectors = vec![0.0f32; n * dim];
|
||||
for i in 0..n {
|
||||
for j in 0..dim {
|
||||
vectors[i * dim + j] = (i * 31 + j) as f32 / 100.0;
|
||||
}
|
||||
}
|
||||
let idx = RabitqIndex::build(&vectors, dim as u32, 42, 20).expect("build");
|
||||
assert_eq!(idx.len(), n as u32);
|
||||
assert!(!idx.is_empty());
|
||||
|
||||
let query: Vec<f32> = vectors[..dim].to_vec();
|
||||
let hits = idx.search(&query, 5).expect("search");
|
||||
assert_eq!(hits.len(), 5);
|
||||
assert_eq!(hits[0].id, 0);
|
||||
assert!(hits[0].distance < 1e-3);
|
||||
}
|
||||
|
||||
#[wasm_bindgen_test]
|
||||
fn version_is_nonempty() {
|
||||
assert!(!version().is_empty());
|
||||
}
|
||||
}
|
||||
|
|
@ -19,10 +19,15 @@ harness = false
|
|||
[dependencies]
|
||||
rand = { workspace = true }
|
||||
rand_distr = { workspace = true }
|
||||
rayon = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
|
||||
# rayon is native-only — wasm32 falls back to sequential iteration
|
||||
# in `from_vectors_parallel_with_rotation`. Output is bit-identical
|
||||
# because rotation is deterministic.
|
||||
[target.'cfg(not(target_arch = "wasm32"))'.dependencies]
|
||||
rayon = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
criterion = { workspace = true }
|
||||
|
|
|
|||
|
|
@ -665,7 +665,6 @@ impl RabitqPlusIndex {
|
|||
kind: RandomRotationKind,
|
||||
items: Vec<(usize, Vec<f32>)>,
|
||||
) -> Result<Self> {
|
||||
use rayon::prelude::*;
|
||||
let mut out = Self::new_with_rotation(dim, seed, rerank_factor, kind);
|
||||
for (_, v) in &items {
|
||||
if v.len() != dim {
|
||||
|
|
@ -675,11 +674,26 @@ impl RabitqPlusIndex {
|
|||
});
|
||||
}
|
||||
}
|
||||
// Phase 1: rotate + bit-pack every vector in parallel. The
|
||||
// rotation matrix is read-only so this is a pure data race
|
||||
// against nothing.
|
||||
// Phase 1: rotate + bit-pack every vector. On native we use rayon
|
||||
// parallel iteration (rotation matrix is read-only — no race). On
|
||||
// wasm32 (single-threaded) we fall back to sequential — output is
|
||||
// bit-identical because the rotation is deterministic, parallel
|
||||
// ordering doesn't affect bytes.
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
let encoded: Vec<(usize, Vec<u64>, f32, Vec<f32>)> = {
|
||||
use rayon::prelude::*;
|
||||
items
|
||||
.into_par_iter()
|
||||
.map(|(id, v)| {
|
||||
let (packed, _) = out.inner.encode_query_packed(&v);
|
||||
let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
(id, packed, norm, v)
|
||||
})
|
||||
.collect()
|
||||
};
|
||||
#[cfg(target_arch = "wasm32")]
|
||||
let encoded: Vec<(usize, Vec<u64>, f32, Vec<f32>)> = items
|
||||
.into_par_iter()
|
||||
.into_iter()
|
||||
.map(|(id, v)| {
|
||||
let (packed, _) = out.inner.encode_query_packed(&v);
|
||||
let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue