rig_core/vector_store/
lsh.rs

1//! Random-projection hashing for approximate vector-search candidates.
2//!
3//! ```
4//! use rig_core::vector_store::lsh::LSHIndex;
5//!
6//! let mut index = LSHIndex::new(2, 4, 8);
7//! index.insert("document", &[1.0, 0.0]);
8//! assert_eq!(index.query(&[1.0, 0.0]), vec!["document"]);
9//! ```
10
11use std::collections::{HashMap, HashSet};
12
13/// The hyperplanes' seed. A fixed seed makes an index, and so its candidates,
14/// the same on every run.
15const SEED: u64 = 0x5eed_fade_cafe_beef;
16
17/// Locality-sensitive hashing by random projection: each table hashes a vector
18/// to the side of each of its random hyperplanes it lies on, so vectors at a
19/// small angle tend to share a bucket. See
20/// <https://www.pinecone.io/learn/series/faiss/locality-sensitive-hashing-random-projection/>.
21#[derive(Clone, Debug, Default)]
22pub struct LSH {
23    /// Every table's hyperplanes, one after another, each `dim` wide.
24    hyperplanes: Vec<f32>,
25    dim: usize,
26    num_hyperplanes: usize,
27}
28
29impl LSH {
30    /// Random unit hyperplanes for `num_tables` tables of `num_hyperplanes`
31    /// each, `dim` wide. A table uses at most 64 hyperplanes; more are
32    /// ignored.
33    pub fn new(dim: usize, num_tables: usize, num_hyperplanes: usize) -> Self {
34        let num_hyperplanes = num_hyperplanes.min(64);
35        let mut rng = fastrand::Rng::with_seed(SEED);
36        let mut hyperplanes = Vec::new();
37        for _ in 0..num_tables.saturating_mul(num_hyperplanes) {
38            let plane: Vec<f32> = (0..dim).map(|_| rng.f32() * 2.0 - 1.0).collect();
39            let norm = plane.iter().map(|x| x * x).sum::<f32>().sqrt();
40            let norm = if norm > 0.0 { norm } else { 1.0 };
41            hyperplanes.extend(plane.iter().map(|x| x / norm));
42        }
43        Self {
44            hyperplanes,
45            dim,
46            num_hyperplanes,
47        }
48    }
49
50    /// The bucket `vector` falls in, in table `table_idx`: one bit per
51    /// hyperplane, set when the vector is on its positive side. A vector of
52    /// another width is compared over the shared prefix; a table index out of
53    /// range hashes everything to `0`.
54    pub fn hash(&self, vector: &[f64], table_idx: usize) -> u64 {
55        let table = self.dim.saturating_mul(self.num_hyperplanes);
56        let planes = table_idx
57            .checked_mul(table)
58            .and_then(|start| self.hyperplanes.get(start..start.checked_add(table)?))
59            .unwrap_or_default();
60        // `max(1)` keeps a zero-width index from asking for empty chunks.
61        planes
62            .chunks(self.dim.max(1))
63            .enumerate()
64            .fold(0, |hash, (bit, plane)| {
65                let dot: f32 = vector.iter().zip(plane).map(|(v, h)| *v as f32 * h).sum();
66                if dot >= 0.0 { hash | 1 << bit } else { hash }
67            })
68    }
69}
70
71/// Document IDs by bucket, one table of buckets per [`LSH`] table.
72#[derive(Clone, Debug, Default)]
73pub struct LSHIndex {
74    lsh: LSH,
75    tables: Vec<HashMap<u64, Vec<String>>>,
76}
77
78impl LSHIndex {
79    /// An empty index for `dim`-wide vectors. More tables find more true
80    /// neighbors; more hyperplanes per table make smaller buckets, so fewer
81    /// candidates to score. Zero tables make one bucket holding every ID, so
82    /// a query returns them all, as brute force would score them all.
83    pub fn new(dim: usize, num_tables: usize, num_hyperplanes: usize) -> Self {
84        // One table without hyperplanes hashes every vector to the same bucket.
85        let (num_tables, num_hyperplanes) = match num_tables {
86            0 => (1, 0),
87            _ => (num_tables, num_hyperplanes),
88        };
89        Self {
90            lsh: LSH::new(dim, num_tables, num_hyperplanes),
91            tables: vec![HashMap::new(); num_tables],
92        }
93    }
94
95    /// Adds `id` to the bucket `embedding` falls in, in every table.
96    pub fn insert(&mut self, id: &str, embedding: &[f64]) {
97        for (table_idx, table) in self.tables.iter_mut().enumerate() {
98            let hash = self.lsh.hash(embedding, table_idx);
99            table.entry(hash).or_default().push(id.to_owned());
100        }
101    }
102
103    /// Removes `id` from the buckets `embedding` falls in, undoing
104    /// [`Self::insert`] of the same pair.
105    pub fn remove(&mut self, id: &str, embedding: &[f64]) {
106        for (table_idx, table) in self.tables.iter_mut().enumerate() {
107            let hash = self.lsh.hash(embedding, table_idx);
108            if let Some(ids) = table.get_mut(&hash) {
109                if let Some(at) = ids.iter().position(|stored| stored == id) {
110                    ids.swap_remove(at);
111                }
112                if ids.is_empty() {
113                    table.remove(&hash);
114                }
115            }
116        }
117    }
118
119    /// Each ID sharing a bucket with `embedding` in any table, once, in
120    /// unspecified order.
121    pub fn query(&self, embedding: &[f64]) -> Vec<&str> {
122        let mut candidates = HashSet::new();
123        for (table_idx, table) in self.tables.iter().enumerate() {
124            if let Some(ids) = table.get(&self.lsh.hash(embedding, table_idx)) {
125                candidates.extend(ids.iter().map(String::as_str));
126            }
127        }
128        candidates.into_iter().collect()
129    }
130
131    /// Removes every ID.
132    pub fn clear(&mut self) {
133        self.tables.iter_mut().for_each(HashMap::clear);
134    }
135}
136
137#[cfg(test)]
138mod tests;
rig_core/vector_store/lsh.rs

rig_core/vector_store/
lsh.rs
