vrr/ruvector
2
0
Fork 0
mirror of https://github.com/ruvnet/RuVector.git synced 2026-07-09 17:28:42 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

style(rabitq): cargo fmt pass to satisfy Rustfmt CI

Browse source 
Pure whitespace changes from `cargo fmt -p ruvector-rabitq`. No
behaviour changes. Keeps the CI Rustfmt check green.

  cargo fmt   -p ruvector-rabitq -- --check  ✓ clean
  cargo test  -p ruvector-rabitq --release   ✓ 20 passed
  cargo clippy -p ruvector-rabitq --release --all-targets -- -D warnings  ✓ clean

Co-Authored-By: claude-flow <ruv@ruv.net>
This commit is contained in:
ruvnet 2026-04-23 14:48:08 -04:00
parent 34b85f1e01
commit 4c2646094b
5 changed files with 146 additions and 39 deletions
Download patch file Download diff file Expand all files Collapse all files

4
crates/ruvector-rabitq/benches/rabitq_bench.rs
View file

@ -32,7 +32,9 @@ fn clustered(n: usize, d: usize, n_clusters: usize, seed: u64) -> Vec<Vec<f32>>
(0..n)
.map(|_| {
let c = &centroids[rng.gen_range(0..n_clusters)];
c.iter().map(|&x| x + noise.sample(&mut rng) as f32).collect()
c.iter()
.map(|&x| x + noise.sample(&mut rng) as f32)
.collect()
})
.collect()
}

122
crates/ruvector-rabitq/src/index.rs
View file

@ -114,20 +114,31 @@ impl TopK {
#[inline]
fn push_raw(&mut self, id: usize, score: f32, pos: usize) {
if self.heap.len() < self.k {
self.heap.push(HeapEntry { id, score, pos: pos as u32 });
self.heap.push(HeapEntry {
id,
score,
pos: pos as u32,
});
return;
}
let worst = self.heap.peek().unwrap().score;
if score.total_cmp(&worst) == Ordering::Less {
self.heap.pop();
self.heap.push(HeapEntry { id, score, pos: pos as u32 });
self.heap.push(HeapEntry {
id,
score,
pos: pos as u32,
});
}
}
fn into_sorted_asc(self) -> Vec<SearchResult> {
let mut v: Vec<SearchResult> = self
.heap
.into_iter()
.map(|e| SearchResult { id: e.id, score: e.score })
.map(|e| SearchResult {
id: e.id,
score: e.score,
})
.collect();
v.sort_unstable_by(|a, b| cmp_score_asc(a.score, b.score));
v
@ -154,13 +165,19 @@ pub struct FlatF32Index {
impl FlatF32Index {
pub fn new(dim: usize) -> Self {
Self { dim, vectors: Vec::new() }
Self {
dim,
vectors: Vec::new(),
}
}
}
#[inline]
fn sq_l2(a: &[f32], b: &[f32]) -> f32 {
a.iter().zip(b.iter()).map(|(&x, &y)| (x - y) * (x - y)).sum()
a.iter()
.zip(b.iter())
.map(|(&x, &y)| (x - y) * (x - y))
.sum()
}
impl AnnIndex for FlatF32Index {
@ -257,7 +274,9 @@ fn build_last_word_mask(dim: usize) -> u64 {
fn build_cos_lut(dim: usize) -> Vec<f32> {
use std::f32::consts::PI;
let d = dim as f32;
(0..=dim).map(|b| (PI * (1.0 - b as f32 / d)).cos()).collect()
(0..=dim)
.map(|b| (PI * (1.0 - b as f32 / d)).cos())
.collect()
}
impl RabitqIndex {
@ -308,7 +327,11 @@ impl RabitqIndex {
/// around [`Self::encode_query_packed`] that boxes the result.
pub fn encode_query(&self, q: &[f32]) -> BinaryCode {
let (words, norm) = self.encode_query_packed(q);
BinaryCode { words, norm, dim: self.dim }
BinaryCode {
words,
norm,
dim: self.dim,
}
}
/// Prepare a query for the asymmetric estimator — returns the rotated
@ -341,7 +364,11 @@ impl RabitqIndex {
let words = self.packed[s..s + self.n_words].to_vec();
(
self.ids[i] as usize,
BinaryCode { words, norm: self.norms[i], dim: self.dim },
BinaryCode {
words,
norm: self.norms[i],
dim: self.dim,
},
)
})
.collect()
@ -469,7 +496,10 @@ impl AnnIndex for RabitqIndex {
let results = self.symmetric_scan_topk(&q_packed, q_norm, k);
Ok(results
.into_iter()
.map(|(_, id, score)| SearchResult { id: id as usize, score })
.map(|(_, id, score)| SearchResult {
id: id as usize,
score,
})
.collect())
}
@ -533,7 +563,9 @@ impl AnnIndex for RabitqPlusIndex {
// Binary-code scan via the tuned SoA loop.
let (q_packed, q_norm) = self.inner.encode_query_packed(query);
let cand = self.inner.symmetric_scan_topk(&q_packed, q_norm, candidates);
let cand = self
.inner
.symmetric_scan_topk(&q_packed, q_norm, candidates);
// Exact rerank on the candidate set — `pos` is the row index.
let k_eff = k.min(cand.len());
@ -552,8 +584,7 @@ impl AnnIndex for RabitqPlusIndex {
self.inner.dim()
}
fn memory_bytes(&self) -> usize {
self.inner.memory_bytes()
+ self.originals.len() * (self.inner.dim * 4 + 24)
self.inner.memory_bytes() + self.originals.len() * (self.inner.dim * 4 + 24)
}
}
@ -636,7 +667,10 @@ impl AnnIndex for RabitqAsymIndex {
let mut out: Vec<SearchResult> = cand
.into_iter()
.take(k_eff)
.map(|(_, id, score)| SearchResult { id: id as usize, score })
.map(|(_, id, score)| SearchResult {
id: id as usize,
score,
})
.collect();
out.sort_unstable_by(|a, b| cmp_score_asc(a.score, b.score));
return Ok(out);
@ -687,12 +721,18 @@ mod tests {
use rand::{Rng as _, SeedableRng as _};
let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
let centroids: Vec<Vec<f32>> = (0..n_clusters)
.map(|_| (0..d).map(|_| rng.gen::<f32>() * 4.0 - 2.0).collect::<Vec<_>>())
.map(|_| {
(0..d)
.map(|_| rng.gen::<f32>() * 4.0 - 2.0)
.collect::<Vec<_>>()
})
.collect();
(0..n)
.map(|_| {
let c = &centroids[rng.gen_range(0..n_clusters)];
c.iter().map(|&x| x + (rng.gen::<f32>() - 0.5) * 0.3).collect()
c.iter()
.map(|&x| x + (rng.gen::<f32>() - 0.5) * 0.3)
.collect()
})
.collect()
}
@ -736,11 +776,20 @@ mod tests {
for q in &queries {
let e: std::collections::HashSet<usize> =
exact.search(q, k).unwrap().iter().map(|r| r.id).collect();
hits += idx.search(q, k).unwrap().iter().filter(|r| e.contains(&r.id)).count();
hits += idx
.search(q, k)
.unwrap()
.iter()
.filter(|r| e.contains(&r.id))
.count();
}
let recall = hits as f64 / (nq * k) as f64;
// k/n = 1% is random chance; we want the estimator to beat it by ≥ 10×.
assert!(recall > 0.20, "recall@10={:.1}% — not above 20 % baseline", recall * 100.0);
assert!(
recall > 0.20,
"recall@10={:.1}% — not above 20 % baseline",
recall * 100.0
);
}
#[test]
@ -764,10 +813,19 @@ mod tests {
for q in &queries {
let e: std::collections::HashSet<usize> =
exact.search(q, k).unwrap().iter().map(|r| r.id).collect();
hits += idx.search(q, k).unwrap().iter().filter(|r| e.contains(&r.id)).count();
hits += idx
.search(q, k)
.unwrap()
.iter()
.filter(|r| e.contains(&r.id))
.count();
}
let recall = hits as f64 / (nq * k) as f64;
assert!(recall > 0.90, "rerank×5 recall@10={:.1}% < 90 %", recall * 100.0);
assert!(
recall > 0.90,
"rerank×5 recall@10={:.1}% < 90 %",
recall * 100.0
);
}
/// Asymmetric (f32 query × 1-bit db) should *equal or beat* symmetric on
@ -797,8 +855,18 @@ mod tests {
for q in &queries {
let e: std::collections::HashSet<usize> =
exact.search(q, k).unwrap().iter().map(|r| r.id).collect();
sh += sym.search(q, k).unwrap().iter().filter(|r| e.contains(&r.id)).count();
ah += asym.search(q, k).unwrap().iter().filter(|r| e.contains(&r.id)).count();
sh += sym
.search(q, k)
.unwrap()
.iter()
.filter(|r| e.contains(&r.id))
.count();
ah += asym
.search(q, k)
.unwrap()
.iter()
.filter(|r| e.contains(&r.id))
.count();
}
let sr = sh as f64 / (nq * k) as f64;
let ar = ah as f64 / (nq * k) as f64;
@ -831,7 +899,12 @@ mod tests {
for q in &queries {
let e: std::collections::HashSet<usize> =
exact.search(q, k).unwrap().iter().map(|r| r.id).collect();
hits += idx.search(q, k).unwrap().iter().filter(|r| e.contains(&r.id)).count();
hits += idx
.search(q, k)
.unwrap()
.iter()
.filter(|r| e.contains(&r.id))
.count();
}
let r = hits as f64 / (nq * k) as f64;
assert!(r > 0.80, "D=100 rerank×5 recall={:.1}% < 80 %", r * 100.0);
@ -870,7 +943,10 @@ mod tests {
// Pure 1-bit index MUST report less than flat — it holds no originals.
assert!(rqb < f, "RabitqIndex {rqb} should be < Flat {f}");
// Plus-index reports MORE than flat — it holds originals AND codes.
assert!(rqpb > f, "RabitqPlusIndex {rqpb} should be > Flat {f} (rerank stores both)");
assert!(
rqpb > f,
"RabitqPlusIndex {rqpb} should be > Flat {f} (rerank stores both)"
);
}
#[test]

16
crates/ruvector-rabitq/src/main.rs
View file

@ -190,7 +190,13 @@ fn run_scale(n: usize, d: usize, n_clusters: usize, nq: usize, seed: u64, k_max:
print_header();
let rows = [
measure("FlatF32 (exact)", &flat, &queries, &truth, k_max),
measure("RaBitQ 1-bit (sym, no rerank)", &rq, &queries, &truth, k_max),
measure(
"RaBitQ 1-bit (sym, no rerank)",
&rq,
&queries,
&truth,
k_max,
),
measure("RaBitQ+ (sym, rerank×5)", &rq_p5, &queries, &truth, k_max),
measure("RaBitQ+ (sym, rerank×20)", &rq_p20, &queries, &truth, k_max),
measure("RaBitQ-Asym (no rerank)", &rq_a1, &queries, &truth, k_max),
@ -255,7 +261,13 @@ fn main() {
let (rq_p5, _) = build_plus(d, 123, 5, db);
print_header();
print_row(&measure("FlatF32", &flat, &queries, &truth, k_max));
print_row(&measure("RaBitQ+ sym ×5 (D=100)", &rq_p5, &queries, &truth, k_max));
print_row(&measure(
"RaBitQ+ sym ×5 (D=100)",
&rq_p5,
&queries,
&truth,
k_max,
));
println!("\nAll numbers reproducible with the seeds above.");
}

38
crates/ruvector-rabitq/src/quantize.rs
View file

@ -136,11 +136,7 @@ impl BinaryCode {
/// `q_rotated` must be length `self.dim`; caller pre-normalises and
/// pre-rotates the query once per search (amortised across n candidates).
#[inline]
pub fn estimated_sq_distance_asymmetric(
&self,
q_rotated_unit: &[f32],
q_norm: f32,
) -> f32 {
pub fn estimated_sq_distance_asymmetric(&self, q_rotated_unit: &[f32], q_norm: f32) -> f32 {
debug_assert_eq!(q_rotated_unit.len(), self.dim);
let d = self.dim;
let inv_sqrt_d = 1.0 / (d as f32).sqrt();
@ -190,15 +186,22 @@ mod tests {
#[test]
fn xnor_self_is_all_ones() {
let v: Vec<f32> = (0..64).map(|i| if i % 2 == 0 { 1.0 } else { -1.0 }).collect();
let v: Vec<f32> = (0..64)
.map(|i| if i % 2 == 0 { 1.0 } else { -1.0 })
.collect();
let code = BinaryCode::encode(&v, 1.0);
let agreement = code.xnor_popcount(&code);
assert_eq!(agreement, 64, "self-agreement should be D=64, got {agreement}");
assert_eq!(
agreement, 64,
"self-agreement should be D=64, got {agreement}"
);
}
#[test]
fn xnor_opposite_is_zero() {
let v: Vec<f32> = (0..64).map(|i| if i % 2 == 0 { 1.0 } else { -1.0 }).collect();
let v: Vec<f32> = (0..64)
.map(|i| if i % 2 == 0 { 1.0 } else { -1.0 })
.collect();
let neg_v: Vec<f32> = v.iter().map(|&x| -x).collect();
let code = BinaryCode::encode(&v, 1.0);
let code_neg = BinaryCode::encode(&neg_v, 1.0);
@ -212,16 +215,24 @@ mod tests {
#[test]
fn masked_popcount_handles_non_aligned_dim() {
// D=100 → 2 u64 words, 28 padding bits.
let v: Vec<f32> = (0..100).map(|i| if i % 2 == 0 { 1.0 } else { -1.0 }).collect();
let v: Vec<f32> = (0..100)
.map(|i| if i % 2 == 0 { 1.0 } else { -1.0 })
.collect();
let neg_v: Vec<f32> = v.iter().map(|&x| -x).collect();
let code = BinaryCode::encode(&v, 1.0);
let code_neg = BinaryCode::encode(&neg_v, 1.0);
// Raw would read 0 matches + 28 padding matches = 28 (wrong).
let raw = code.xnor_popcount(&code_neg);
assert_eq!(raw, 28, "raw xnor should count padding as matches (bug demo)");
assert_eq!(
raw, 28,
"raw xnor should count padding as matches (bug demo)"
);
// Masked must report 0 matches.
let masked = code.masked_xnor_popcount(&code_neg);
assert_eq!(masked, 0, "masked xnor must ignore padding bits; got {masked}");
assert_eq!(
masked, 0,
"masked xnor must ignore padding bits; got {masked}"
);
// Self-compare: every real bit matches, padding is masked.
let self_masked = code.masked_xnor_popcount(&code);
assert_eq!(self_masked, 100);
@ -246,7 +257,10 @@ mod tests {
let code = BinaryCode::encode(&unit, 1.0);
let est = code.estimated_sq_distance(&code);
// Symmetric Charikar estimator on the same code: cos(π·(1−D/D))=1 → est=0.
assert!(est.abs() < 1e-5, "self sq-distance estimate too large: {est}");
assert!(
est.abs() < 1e-5,
"self sq-distance estimate too large: {est}"
);
}
#[test]

5
crates/ruvector-rabitq/src/rotation.rs
View file

@ -24,7 +24,10 @@ impl RandomRotation {
let mut m: Vec<Vec<f32>> = (0..dim)
.map(|_| {
(0..dim)
.map(|_| <StandardNormal as Distribution<f64>>::sample(&StandardNormal, &mut rng) as f32)
.map(|_| {
<StandardNormal as Distribution<f64>>::sample(&StandardNormal, &mut rng)
as f32
})
.collect()
})
.collect();