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:
Claude 2026-07-31 07:41:03 +00:00
parent 74665cde37
commit e1919cf225
No known key found for this signature in database

View file

@ -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)]