polars_utils/
min_hash_sketch.rs

1const NUM_ENTRIES: usize = 256;
2
3#[derive(Clone, Copy, Debug)]
4struct Entry {
5    tag: u32,
6    count: u32,
7}
8
9/// A sample of distinct keys chosen by hash, along with how often each key
10/// occurred. Whether a key is sampled does not depend on how often it occurs.
11#[derive(Clone)]
12pub struct MinHashSketch {
13    entries: Box<[Entry; NUM_ENTRIES]>,
14}
15
16impl Default for MinHashSketch {
17    fn default() -> Self {
18        Self::new()
19    }
20}
21
22impl MinHashSketch {
23    pub fn new() -> Self {
24        let empty = Entry {
25            tag: u32::MAX,
26            count: 0,
27        };
28        Self {
29            entries: vec![empty; NUM_ENTRIES].try_into().unwrap(),
30        }
31    }
32
33    /// Add a new hash to the sketch.
34    #[inline]
35    pub fn insert(&mut self, h: u64) {
36        const IDX_ODD: u64 = 0x85b0fc582ca15d65; // Randomly chosen.
37        const TAG_ODD: u64 = 0x4a0da407caf4e9ad; // Randomly chosen.
38        let idx = (h.wrapping_mul(IDX_ODD) >> (64 - NUM_ENTRIES.trailing_zeros())) as usize;
39        let tag = (h.wrapping_mul(TAG_ODD) >> 32) as u32;
40        let entry = &mut self.entries[idx];
41        if tag < entry.tag {
42            entry.tag = tag;
43            entry.count = 0;
44        }
45        entry.count += (tag == entry.tag) as u32;
46    }
47
48    pub fn combine(&mut self, other: &MinHashSketch) {
49        for (entry, other) in self.entries.iter_mut().zip(other.entries.iter()) {
50            if other.tag < entry.tag {
51                *entry = *other;
52            } else if other.tag == entry.tag {
53                entry.count += other.count;
54            }
55        }
56    }
57
58    /// Estimates the number of distinct keys as a fraction of the number of
59    /// rows, weighting each distinct key equally.
60    pub fn key_ratio(&self) -> Option<f64> {
61        let mut nonempty = 0u64;
62        let mut total = 0u64;
63        for entry in self.entries.iter() {
64            nonempty += (entry.count > 0) as u64;
65            total += entry.count as u64;
66        }
67        (nonempty > 0).then(|| nonempty as f64 / total as f64)
68    }
69}
polars_utils/min_hash_sketch.rs

polars_utils/
min_hash_sketch.rs
