diff --git a/src/vector.rs b/src/vector.rs index 4b0fd9d..7abb060 100644 --- a/src/vector.rs +++ b/src/vector.rs @@ -1 +1,28 @@ -// implemented in a later task +/// Exact cosine search. 90k×1024 f32 ≈ 370 MB held in RAM; a query +/// is 90k dot products, well under 50 ms. No ANN dependency (YAGNI). +pub struct VectorIndex { + vectors: Vec>, + norms: Vec, +} + +impl VectorIndex { + pub fn from_vectors(vectors: Vec>) -> Self { + let norms = vectors.iter().map(|v| dot(v, v).sqrt().max(1e-8)).collect(); + Self { vectors, norms } + } + + /// (entry_index, cosine_similarity) descending. + pub fn top_k(&self, query: &[f32], k: usize) -> Vec<(usize, f32)> { + let qn = dot(query, query).sqrt().max(1e-8); + let mut scored: Vec<(usize, f32)> = self.vectors.iter().enumerate() + .map(|(i, v)| (i, dot(v, query) / (self.norms[i] * qn))) + .collect(); + scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); + scored.truncate(k); + scored + } +} + +fn dot(a: &[f32], b: &[f32]) -> f32 { + a.iter().zip(b).map(|(x, y)| x * y).sum() +} diff --git a/tests/vector_tests.rs b/tests/vector_tests.rs new file mode 100644 index 0000000..67756fa --- /dev/null +++ b/tests/vector_tests.rs @@ -0,0 +1,14 @@ +use alpha_id::vector::VectorIndex; + +#[test] +fn cosine_topk_ranks_nearest_first() { + let vi = VectorIndex::from_vectors(vec![ + vec![1.0, 0.0], // id 0 + vec![0.0, 1.0], // id 1 + vec![0.9, 0.1], // id 2 + ]); + let hits = vi.top_k(&[1.0, 0.0], 2); + assert_eq!(hits[0].0, 0); + assert_eq!(hits[1].0, 2); + assert!(hits[0].1 > hits[1].1); +}