ruvector_core/index/
flat.rs

1//! Flat (brute-force) index for baseline and small datasets
2
3use crate::distance::distance;
4use crate::error::Result;
5use crate::index::VectorIndex;
6use crate::types::{DistanceMetric, SearchResult, VectorId};
7use dashmap::DashMap;
8use rayon::prelude::*;
9
10/// Flat index using brute-force search
11pub struct FlatIndex {
12    vectors: DashMap<VectorId, Vec<f32>>,
13    metric: DistanceMetric,
14    dimensions: usize,
15}
16
17impl FlatIndex {
18    /// Create a new flat index
19    pub fn new(dimensions: usize, metric: DistanceMetric) -> Self {
20        Self {
21            vectors: DashMap::new(),
22            metric,
23            dimensions,
24        }
25    }
26}
27
28impl VectorIndex for FlatIndex {
29    fn add(&mut self, id: VectorId, vector: Vec<f32>) -> Result<()> {
30        self.vectors.insert(id, vector);
31        Ok(())
32    }
33
34    fn search(&self, query: &[f32], k: usize) -> Result<Vec<SearchResult>> {
35        // Parallel distance calculation
36        let mut results: Vec<_> = self
37            .vectors
38            .iter()
39            .par_bridge()
40            .map(|entry| {
41                let id = entry.key().clone();
42                let vector = entry.value();
43                let dist = distance(query, vector, self.metric)?;
44                Ok((id, dist))
45            })
46            .collect::<Result<Vec<_>>>()?;
47
48        // Sort by distance and take top k
49        results.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
50        results.truncate(k);
51
52        Ok(results
53            .into_iter()
54            .map(|(id, score)| SearchResult {
55                id,
56                score,
57                vector: None,
58                metadata: None,
59            })
60            .collect())
61    }
62
63    fn remove(&mut self, id: &VectorId) -> Result<bool> {
64        Ok(self.vectors.remove(id).is_some())
65    }
66
67    fn len(&self) -> usize {
68        self.vectors.len()
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75
76    #[test]
77    fn test_flat_index() -> Result<()> {
78        let mut index = FlatIndex::new(3, DistanceMetric::Euclidean);
79
80        index.add("v1".to_string(), vec![1.0, 0.0, 0.0])?;
81        index.add("v2".to_string(), vec![0.0, 1.0, 0.0])?;
82        index.add("v3".to_string(), vec![0.0, 0.0, 1.0])?;
83
84        let query = vec![1.0, 0.0, 0.0];
85        let results = index.search(&query, 2)?;
86
87        assert_eq!(results.len(), 2);
88        assert_eq!(results[0].id, "v1");
89        assert!(results[0].score < 0.01);
90
91        Ok(())
92    }
93}
ruvector_core/index/flat.rs

ruvector_core/index/
flat.rs
