Skip to main content

datasketch_minhash_lsh/
minhash_lsh.rs

1use crate::error::MinHashingError;
2use crate::minhash::MinHash;
3use float_cmp::ApproxEq;
4use quadrature::integrate;
5use std::collections::{HashMap, HashSet};
6use std::hash::Hash;
7
8const _ALLOWED_INTEGRATE_ERR: f64 = 0.001;
9
10/// The weights configuring whether to prefer false positives or false negatives
11#[derive(Clone)]
12pub struct Weights(pub f64, pub f64);
13
14/// A part of a HashValue used in MinHashLsh
15#[derive(Clone, Debug, Hash, PartialEq, Eq)]
16pub struct HashValuePart(pub Vec<u64>);
17
18/// The LSH params for the number of bands and the band size
19#[derive(Clone, Debug)]
20pub struct LshParams {
21    pub b: usize,
22    pub r: usize,
23}
24
25impl LshParams {
26    pub fn find_optimal_params(threshold: f64, num_perm: usize, weights: &Weights) -> LshParams {
27        let Weights(false_positive_weight, false_negative_weight) = weights;
28        let mut min_error = f64::INFINITY;
29        let mut opt = LshParams { b: 0, r: 0 };
30        for b in 1..num_perm + 1 {
31            let max_r = num_perm / b;
32            for r in 1..max_r + 1 {
33                let false_pos = LshParams::false_positive_probability(threshold, b, r);
34                let false_neg = LshParams::false_negative_probability(threshold, b, r);
35                let error = false_pos * false_positive_weight + false_neg * false_negative_weight;
36                if error < min_error {
37                    min_error = error;
38                    opt = LshParams { b, r };
39                }
40            }
41        }
42        opt
43    }
44
45    fn false_positive_probability(threshold: f64, b: usize, r: usize) -> f64 {
46        let _probability =
47            |s| -> f64 { 1. - f64::powf(1. - f64::powi(s, r as i32) as f64, b as f64) };
48        integrate(_probability, 0.0, threshold, _ALLOWED_INTEGRATE_ERR).integral
49    }
50
51    fn false_negative_probability(threshold: f64, b: usize, r: usize) -> f64 {
52        let _probability =
53            |s| -> f64 { 1. - (1. - f64::powf(1. - f64::powi(s, r as i32), b as f64)) };
54        integrate(_probability, threshold, 1.0, _ALLOWED_INTEGRATE_ERR).integral
55    }
56}
57
58/// The MinHashLsh struct
59#[derive(Clone)]
60pub struct MinHashLsh<KeyType: Eq + Hash + Clone> {
61    num_perm: usize,
62    threshold: f64,
63    weights: Weights,
64    buffer_size: usize,
65    params: LshParams,
66    hash_tables: Vec<HashMap<HashValuePart, HashSet<KeyType>>>,
67    hash_ranges: Vec<(usize, usize)>,
68    keys: HashMap<KeyType, Vec<HashValuePart>>,
69}
70
71type Result<T> = std::result::Result<T, MinHashingError>;
72
73impl<KeyType: Eq + Hash + Clone> MinHashLsh<KeyType> {
74    /// Build a new MinHashLsh struct
75    pub fn new(
76        num_perm: usize,
77        weights: Option<Weights>,
78        threshold: Option<f64>,
79    ) -> Result<MinHashLsh<KeyType>> {
80        let threshold = match threshold {
81            Some(threshold) if !(0.0..=1.0).contains(&threshold) => {
82                return Err(MinHashingError::WrongThresholdInterval);
83            }
84            Some(threshold) => threshold,
85            _ => 0.9,
86        };
87        if num_perm < 2 {
88            return Err(MinHashingError::NumPermFuncsTooLow);
89        }
90
91        let weights = match weights {
92            Some(weights) => {
93                let Weights(left, right) = weights;
94                if !(0.0..=1.0).contains(&left) || !(0.0..=1.0).contains(&right) {
95                    return Err(MinHashingError::WrongWeightThreshold);
96                }
97                let sum_weights = left + right;
98                if !sum_weights.approx_eq(1.0, (0.0, 2)) {
99                    return Err(MinHashingError::UnexpectedSumWeight);
100                }
101                weights
102            }
103            _ => Weights(0.5, 0.5),
104        };
105        let params = LshParams::find_optimal_params(threshold, num_perm, &weights);
106
107        let hash_tables = (0..params.b).into_iter().map(|_| HashMap::new()).collect();
108        let hash_ranges = (0..params.b)
109            .into_iter()
110            .map(|i| (i * params.r, (i + 1) * params.r))
111            .collect();
112        Ok(MinHashLsh {
113            num_perm,
114            threshold,
115            weights,
116            buffer_size: 50_000,
117            params,
118            hash_tables,
119            hash_ranges,
120            keys: HashMap::<KeyType, Vec<HashValuePart>>::new(),
121        })
122    }
123
124    /// Check whether the MinHashLsh contains any MinHash structs
125    pub fn is_empty(&self) -> bool {
126        self.hash_tables.iter().any(|table| table.len() == 0)
127    }
128
129    /// Insert a new MinHash struct
130    pub fn insert(&mut self, key: KeyType, min_hash: &MinHash) -> Result<()> {
131        // TODO: We could also add optional checks whether the key is already present in index
132        // TODO: Why has the original implementation buffer params everywhere
133        if min_hash.hash_values.0.len() != self.num_perm {
134            return Err(MinHashingError::DifferentNumPermFuncs);
135        }
136        let mut hash_value_parts: Vec<HashValuePart> = self
137            .hash_ranges
138            .iter()
139            .map(|(start, end)| {
140                let hash_part = min_hash.hash_values.0[*start..*end].to_owned();
141                HashValuePart(hash_part)
142            })
143            .collect();
144        self.keys.insert(key.clone(), hash_value_parts.clone());
145        let hash_table_iter = &mut self.hash_tables.iter_mut();
146        let zipped_drain_iter = hash_value_parts.drain(..).zip(hash_table_iter);
147        for (hash_part, hash_table) in zipped_drain_iter {
148            hash_table
149                .entry(hash_part)
150                .or_insert_with(HashSet::new)
151                .insert(key.clone());
152        }
153        Ok(())
154    }
155
156    /// Checks whether a MinHash struct with a specific key is contained in the MinHashLsh
157    pub fn contains_key(&self, key: &KeyType) -> bool {
158        self.keys.contains_key(key)
159    }
160
161    /// Remove a MinHash struct with a specific key from the MinHashLsh
162    pub fn remove(&mut self, key: &KeyType) -> Result<()> {
163        if !self.keys.contains_key(key) {
164            return Err(MinHashingError::KeyDoesNotExist);
165        }
166        for (hash_part, table) in self
167            .keys
168            .get_mut(key)
169            .unwrap()
170            .iter_mut()
171            .zip(&mut self.hash_tables)
172        {
173            table.get_mut(hash_part).unwrap().remove(key);
174            if let Some(set) = table.get(hash_part) {
175                if set.is_empty() {
176                    table.remove(hash_part);
177                }
178            }
179        }
180        self.keys.remove(key);
181        Ok(())
182    }
183
184    /// Get the number of MinHash structs contained in the MinHashLsh
185    pub fn get_counts(&self) -> Vec<HashMap<HashValuePart, usize>> {
186        self.hash_tables
187            .iter()
188            .map(|table| {
189                table
190                    .iter()
191                    .map(|(key, value)| (key.clone(), value.len()))
192                    .collect()
193            })
194            .collect()
195    }
196
197    /// Query for candidates potentially within a jaccard-distance corresponding to the configured
198    /// threshold
199    pub fn query(&mut self, min_hash: &MinHash) -> Result<HashSet<KeyType>> {
200        if min_hash.hash_values.0.len() != self.num_perm {
201            return Err(MinHashingError::DifferentNumPermFuncs);
202        }
203        let unique_candidates = self
204            .hash_ranges
205            .iter()
206            .zip(&self.hash_tables)
207            .flat_map(|(range, table)| {
208                let (start, end) = range;
209                let hash_part = min_hash.hash_values.0[*start..*end].to_owned();
210                table.get(&HashValuePart(hash_part))
211            })
212            .flatten()
213            .cloned()
214            .collect();
215        Ok(unique_candidates)
216    }
217}
218
219#[cfg(test)]
220mod test {
221    use super::*;
222    use crate::minhash::MinHash;
223
224    #[test]
225    fn test_init() -> Result<()> {
226        let lsh = <MinHashLsh<&str>>::new(128, None, Some(0.8))?;
227        assert!(lsh.is_empty());
228        let LshParams { b: b1, r: r1 } = lsh.params;
229        let lsh = <MinHashLsh<&str>>::new(128, Some(Weights(0.2, 0.8)), Some(0.8))?;
230        let LshParams { b: b2, r: r2 } = lsh.params;
231        assert!(b1 < b2);
232        assert!(r1 > r2);
233        Ok(())
234    }
235
236    #[test]
237    fn test_insert() -> Result<()> {
238        let mut lsh = <MinHashLsh<&str>>::new(128, None, Some(0.5))?;
239        let mut m1 = <MinHash>::new(128, Some(0));
240        m1.update(&"a");
241        let mut m2 = <MinHash>::new(128, Some(0));
242        m2.update(&"b");
243        lsh.insert("a", &m1)?;
244        lsh.insert("b", &m2)?;
245        for table in &lsh.hash_tables {
246            assert!(table.len() >= 1);
247            let table_values: HashSet<_> = table.values().flatten().collect();
248            assert!(table_values.contains(&"a"));
249            assert!(table_values.contains(&"b"));
250        }
251        assert!(lsh.contains_key(&"a"));
252        assert!(lsh.contains_key(&"b"));
253        let a_keys_content = lsh.keys.get(&"a").unwrap();
254        for (index, hash_part) in a_keys_content.iter().enumerate() {
255            assert!(lsh.hash_tables[index][hash_part].contains(&"a"));
256        }
257        Ok(())
258    }
259
260    #[test]
261    fn test_query() -> Result<()> {
262        let mut lsh = <MinHashLsh<&str>>::new(16, None, Some(0.5))?;
263        let mut m1 = <MinHash>::new(16, Some(0));
264        m1.update(&"a");
265        let mut m2 = <MinHash>::new(16, Some(0));
266        m2.update(&"b");
267        lsh.insert("a", &m1)?;
268        lsh.insert("b", &m2)?;
269        let result = lsh.query(&m1)?;
270        assert!(result.contains(&"a"));
271        let result = lsh.query(&m2)?;
272        assert!(result.contains(&"b"));
273        assert!(result.len() <= 2);
274
275        let m3 = <MinHash>::new(18, Some(0));
276        let result = std::panic::catch_unwind(|| {
277            lsh.clone().query(&m3).unwrap();
278        });
279        assert!(result.is_err());
280        Ok(())
281    }
282
283    #[test]
284    fn test_remove() -> Result<()> {
285        let mut lsh = <MinHashLsh<&str>>::new(16, None, Some(0.5))?;
286        let mut m1 = <MinHash>::new(16, Some(0));
287        m1.update(&"a");
288        let mut m2 = <MinHash>::new(16, Some(0));
289        m2.update(&"b");
290        lsh.insert("a", &m1)?;
291        lsh.insert("b", &m2)?;
292
293        lsh.remove(&"a")?;
294        assert!(!lsh.keys.contains_key("&a"));
295        for table in lsh.hash_tables {
296            for value in table.keys() {
297                assert!(table[value].len() > 0);
298                assert!(!table[value].contains(&"a"))
299            }
300        }
301        Ok(())
302    }
303
304    #[test]
305    fn test_get_counts() -> Result<()> {
306        let mut lsh = <MinHashLsh<&str>>::new(16, None, Some(0.5))?;
307        let mut m1 = <MinHash>::new(16, Some(0));
308        m1.update(&"a");
309        let mut m2 = <MinHash>::new(16, Some(0));
310        m2.update(&"b");
311        lsh.insert("a", &m1)?;
312        lsh.insert("b", &m2)?;
313
314        let counts = lsh.get_counts();
315        assert_eq!(counts.len(), lsh.params.b);
316        for table in &counts {
317            assert_eq!(table.values().sum::<usize>(), 2);
318        }
319        Ok(())
320    }
321
322    #[test]
323    fn example_eg1() -> Result<()> {
324        let set1: HashSet<&'static str> = [
325            "minhash",
326            "is",
327            "a",
328            "probabilistic",
329            "data",
330            "structure",
331            "for",
332            "estimating",
333            "the",
334            "similarity",
335            "between",
336            "datasets",
337        ]
338        .iter()
339        .cloned()
340        .collect();
341        let set2: HashSet<&'static str> = [
342            "minhash",
343            "is",
344            "a",
345            "probability",
346            "data",
347            "structure",
348            "for",
349            "estimating",
350            "the",
351            "similarity",
352            "between",
353            "documents",
354        ]
355        .iter()
356        .cloned()
357        .collect();
358        let set3: HashSet<&'static str> = [
359            "minhash",
360            "is",
361            "probability",
362            "data",
363            "structure",
364            "for",
365            "estimating",
366            "the",
367            "similarity",
368            "between",
369            "documents",
370        ]
371        .iter()
372        .cloned()
373        .collect();
374
375        let n_projections = 128;
376        let mut m1 = <MinHash>::new(n_projections, Some(0));
377        let mut m2 = <MinHash>::new(n_projections, Some(0));
378        let mut m3 = <MinHash>::new(n_projections, Some(0));
379        for d in set1 {
380            m1.update(&d);
381        }
382        for d in set2 {
383            m2.update(&d);
384        }
385        for d in set3 {
386            m3.update(&d);
387        }
388
389        // Create LSHindex
390        let mut lsh = <MinHashLsh<&str>>::new(128, None, Some(0.5))?;
391        lsh.insert(&"m2", &m2)?;
392        lsh.insert(&"m3", &m3)?;
393        let result = lsh.query(&m1);
394        println!(
395            "Approximate neighbours with Jaccard similarity > 0.5: {:?}",
396            result
397        );
398        Ok(())
399    }
400}