diff --git a/crates/ruvector-tiered-quant/src/lib.rs b/crates/ruvector-tiered-quant/src/lib.rs index f4acb2137..5aebb9aeb 100644 --- a/crates/ruvector-tiered-quant/src/lib.rs +++ b/crates/ruvector-tiered-quant/src/lib.rs @@ -15,6 +15,7 @@ //! ## Trait surface //! [`TieredIndex`] is the main trait. [`HotWarmColdIndex`] is the concrete implementation. +pub mod baselines; pub mod dataset; pub mod quantizer; pub mod tier; @@ -23,6 +24,7 @@ use std::collections::BinaryHeap; // ─── public re-exports ──────────────────────────────────────────────────────── +pub use baselines::{recall_at_k, FlatF32Index, FlatU8Index}; pub use quantizer::{BinaryQuantizer, ScalarQuantizer}; pub use tier::{TierKind, TierStats, VectorRecord}; @@ -291,114 +293,6 @@ fn decode_to_f32(tier: &TierKind, dims: usize) -> Vec { } } -// ─── baseline variant: all f32 ─────────────────────────────────────────────── - -/// Flat f32 index; ground-truth baseline for recall measurement. -pub struct FlatF32Index { - dims: usize, - vectors: Vec<(usize, Vec)>, -} - -impl FlatF32Index { - pub fn new(dims: usize) -> Self { - Self { - dims, - vectors: Vec::new(), - } - } - - pub fn insert(&mut self, id: usize, vector: Vec) { - self.vectors.push((id, vector)); - } - - pub fn query(&self, query: &[f32], k: usize) -> Vec { - let mut results: Vec = self - .vectors - .iter() - .map(|(id, v)| { - let dist = query - .iter() - .zip(v.iter()) - .map(|(a, b)| (a - b).powi(2)) - .sum::() - .sqrt(); - Hit { id: *id, dist } - }) - .collect(); - results.sort_by(|a, b| a.dist.partial_cmp(&b.dist).unwrap()); - results.truncate(k); - results - } - - pub fn len(&self) -> usize { - self.vectors.len() - } -} - -// ─── all-quantized variant: all u8 ─────────────────────────────────────────── - -/// All-warm index: every vector stored as u8 scalar quantized. No tiering. -pub struct FlatU8Index { - dims: usize, - sq: ScalarQuantizer, - entries: Vec<(usize, Vec, Vec, Vec)>, -} - -impl FlatU8Index { - pub fn new(dims: usize) -> Self { - Self { - dims, - sq: ScalarQuantizer::new(dims), - entries: Vec::new(), - } - } - - pub fn insert(&mut self, id: usize, vector: Vec) { - let (bytes, min, scale) = self.sq.quantize(&vector); - self.entries.push((id, bytes, min, scale)); - } - - pub fn query(&self, query: &[f32], k: usize) -> Vec { - let mut results: Vec = self - .entries - .iter() - .map(|(id, bytes, min, scale)| { - let dist: f32 = query - .iter() - .enumerate() - .map(|(i, &q)| { - let r = min[i] + bytes[i] as f32 * scale[i]; - (q - r).powi(2) - }) - .sum::() - .sqrt(); - Hit { id: *id, dist } - }) - .collect(); - results.sort_by(|a, b| a.dist.partial_cmp(&b.dist).unwrap()); - results.truncate(k); - results - } - - pub fn len(&self) -> usize { - self.entries.len() - } -} - -// ─── recall helper ─────────────────────────────────────────────────────────── - -/// Compute recall@k: fraction of ground-truth top-k found in candidate top-k. -pub fn recall_at_k(ground_truth: &[Hit], candidates: &[Hit], k: usize) -> f32 { - let gt_ids: std::collections::HashSet = - ground_truth.iter().take(k).map(|h| h.id).collect(); - let found = candidates - .iter() - .take(k) - .filter(|h| gt_ids.contains(&h.id)) - .count(); - found as f32 / k.min(ground_truth.len()) as f32 -} - // ─── tests ─────────────────────────────────────────────────────────────────── #[cfg(test)]