polars_utils/
f2_sketch.rs

1const NUM_BUCKETS: usize = 64;
2
3/// Estimates the second frequency moment F2 = sum_k f_k^2 of a stream of hashed
4/// keys, where f_k is the number of times key k was inserted.
5///
6/// This is a single row of a Count-Sketch (the "fast AMS" sketch): every hash
7/// adds +1 or -1 to one of 64 buckets, and the sum of squared buckets is an
8/// unbiased estimate of F2 with a relative standard error of ~18%.
9#[derive(Clone)]
10pub struct F2Sketch {
11    buckets: [i64; NUM_BUCKETS],
12    num_inserts: u64,
13}
14
15impl Default for F2Sketch {
16    fn default() -> Self {
17        Self::new()
18    }
19}
20
21impl F2Sketch {
22    pub fn new() -> Self {
23        Self {
24            buckets: [0; NUM_BUCKETS],
25            num_inserts: 0,
26        }
27    }
28
29    /// Add a new hash to the sketch.
30    #[inline]
31    pub fn insert(&mut self, h: u64) {
32        // An arbitrary odd number different from the one in the cardinality
33        // sketch, so the two use independent bits.
34        const ARBITRARY_ODD: u64 = 0xd6e8feb86659fd93;
35        let x = h.wrapping_mul(ARBITRARY_ODD);
36        let idx = (x >> 58) as usize;
37        let negative = ((x >> 57) & 1) as i64;
38        self.buckets[idx] += 1 - 2 * negative;
39        self.num_inserts += 1;
40    }
41
42    pub fn num_inserts(&self) -> u64 {
43        self.num_inserts
44    }
45
46    pub fn estimate(&self) -> f64 {
47        self.buckets.iter().map(|c| (*c as f64) * (*c as f64)).sum()
48    }
49}
polars_utils/f2_sketch.rs

polars_utils/
f2_sketch.rs
