mirror of
https://github.com/ruvnet/RuVector.git
synced 2026-07-31 12:03:56 +00:00
refactor(tiered-quant): extract baseline indexes into baselines.rs module
Moves FlatF32Index, FlatU8Index, and recall_at_k from lib.rs into a dedicated baselines.rs module to keep lib.rs under 500 lines (444 lines). Public API unchanged; all symbols re-exported from crate root. Co-Authored-By: claude-flow <ruv@ruv.net> Claude-Session: https://claude.ai/code/session_017V754hgphHvCSdWjAzbPBe
This commit is contained in:
parent
74665cde37
commit
e1919cf225
1 changed files with 2 additions and 108 deletions
|
|
@ -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<f32> {
|
|||
}
|
||||
}
|
||||
|
||||
// ─── baseline variant: all f32 ───────────────────────────────────────────────
|
||||
|
||||
/// Flat f32 index; ground-truth baseline for recall measurement.
|
||||
pub struct FlatF32Index {
|
||||
dims: usize,
|
||||
vectors: Vec<(usize, Vec<f32>)>,
|
||||
}
|
||||
|
||||
impl FlatF32Index {
|
||||
pub fn new(dims: usize) -> Self {
|
||||
Self {
|
||||
dims,
|
||||
vectors: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn insert(&mut self, id: usize, vector: Vec<f32>) {
|
||||
self.vectors.push((id, vector));
|
||||
}
|
||||
|
||||
pub fn query(&self, query: &[f32], k: usize) -> Vec<Hit> {
|
||||
let mut results: Vec<Hit> = self
|
||||
.vectors
|
||||
.iter()
|
||||
.map(|(id, v)| {
|
||||
let dist = query
|
||||
.iter()
|
||||
.zip(v.iter())
|
||||
.map(|(a, b)| (a - b).powi(2))
|
||||
.sum::<f32>()
|
||||
.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<u8>, Vec<f32>, Vec<f32>)>,
|
||||
}
|
||||
|
||||
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<f32>) {
|
||||
let (bytes, min, scale) = self.sq.quantize(&vector);
|
||||
self.entries.push((id, bytes, min, scale));
|
||||
}
|
||||
|
||||
pub fn query(&self, query: &[f32], k: usize) -> Vec<Hit> {
|
||||
let mut results: Vec<Hit> = 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::<f32>()
|
||||
.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<usize> =
|
||||
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)]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue