research(nightly): soar-ivf — SOAR orthogonality-amplified IVF spilling (NeurIPS 2023)

Cargo.toml

@@ -18,6 +18,7 @@ exclude = ["crates/micro-hnsw-wasm", "crates/ruvector-hyperbolic-hnsw", "crates/
     # land in iters 92-97.
     "crates/ruos-thermal"]
 members = [
+    "crates/ruvector-soar",
     "crates/ruvector-acorn",
     "crates/ruvector-acorn-wasm",
     "crates/ruvector-rabitq",

crates/ruvector-soar/Cargo.toml

@@ -0,0 +1,30 @@
+[package]
+name = "ruvector-soar"
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+authors.workspace = true
+repository.workspace = true
+description = "SOAR-IVF: Spilling with Orthogonality-Amplified Residuals for high-recall approximate nearest-neighbour search — NeurIPS 2023"
+
+[[bin]]
+name = "soar-demo"
+path = "src/main.rs"
+
+[[bench]]
+name = "soar_bench"
+harness = false
+
+[dependencies]
+rand = { workspace = true }
+rand_distr = { workspace = true }
+thiserror = { workspace = true }
+serde = { workspace = true }
+serde_json = { workspace = true }
+
+[target.'cfg(not(target_arch = "wasm32"))'.dependencies]
+rayon = { workspace = true }
+
+[dev-dependencies]
+criterion = { workspace = true }

crates/ruvector-soar/benches/soar_bench.rs

@@ -0,0 +1,81 @@
+use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
+use rand::SeedableRng;
+use rand_distr::{Distribution, Normal, Uniform};
+use ruvector_soar::{IndexKind, SoarConfig, SoarIndex};
+
+fn gen_data(n: usize, d: usize, seed: u64) -> Vec<Vec<f32>> {
+    let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
+    let range = Uniform::new(-1.0f32, 1.0);
+    let noise = Normal::new(0.0f64, 0.3).unwrap();
+    let centroids: Vec<Vec<f32>> = (0..20)
+        .map(|_| (0..d).map(|_| range.sample(&mut rng)).collect())
+        .collect();
+    (0..n)
+        .map(|i| {
+            let c = &centroids[i % 20];
+            c.iter()
+                .map(|&x| x + noise.sample(&mut rng) as f32)
+                .collect()
+        })
+        .collect()
+}
+
+fn bench_search(c: &mut Criterion) {
+    let n = 5_000;
+    let d = 128;
+    let nq = 50;
+    let k = 10;
+    let corpus = gen_data(n, d, 1);
+    let queries = gen_data(nq, d, 2);
+
+    let mut group = c.benchmark_group("soar_search");
+
+    for &nprobe in &[4usize, 8, 16] {
+        // IVF-PQ
+        let ivf = SoarIndex::build(
+            corpus.clone(),
+            SoarConfig {
+                kind: IndexKind::IvfPq,
+                nlist: 64,
+                nprobe,
+                m_pq: 8,
+                ..Default::default()
+            },
+        )
+        .unwrap();
+        group.bench_with_input(BenchmarkId::new("IVF-PQ", nprobe), &nprobe, |b, _| {
+            b.iter(|| {
+                for q in &queries {
+                    let _ = ivf.search(q, k).unwrap();
+                }
+            })
+        });
+
+        // SOAR-IVF-PQ
+        let soar = SoarIndex::build(
+            corpus.clone(),
+            SoarConfig {
+                kind: IndexKind::SoarIvfPq,
+                nlist: 64,
+                nprobe,
+                m_pq: 8,
+                lambda: 1.0,
+                n_secondary_candidates: 10,
+                ..Default::default()
+            },
+        )
+        .unwrap();
+        group.bench_with_input(BenchmarkId::new("SOAR-IVF-PQ", nprobe), &nprobe, |b, _| {
+            b.iter(|| {
+                for q in &queries {
+                    let _ = soar.search(q, k).unwrap();
+                }
+            })
+        });
+    }
+
+    group.finish();
+}
+
+criterion_group!(benches, bench_search);
+criterion_main!(benches);

crates/ruvector-soar/src/error.rs

@@ -0,0 +1,18 @@
+use thiserror::Error;
+
+#[derive(Debug, Error)]
+pub enum SoarError {
+    #[error("dimension mismatch: expected {expected}, got {actual}")]
+    DimensionMismatch { expected: usize, actual: usize },
+
+    #[error("empty dataset")]
+    Empty,
+
+    #[error("invalid config: {0}")]
+    InvalidConfig(String),
+
+    #[error("index not trained")]
+    NotTrained,
+}
+
+pub type Result<T> = std::result::Result<T, SoarError>;