datafusion_functions_aggregate/
hyperloglog.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! # HyperLogLog
19//!
20//! `hyperloglog` is a module that contains a modified version
21//! of [redis's implementation](https://github.com/redis/redis/blob/4930d19e70c391750479951022e207e19111eb55/src/hyperloglog.c)
22//! with some modification based on strong assumption of usage
23//! within datafusion, so that function can
24//! be efficiently implemented.
25//!
26//! Specifically, like Redis's version, this HLL structure uses
27//! 2**14 = 16384 registers, which means the standard error is
28//! 1.04/(16384**0.5) = 0.8125%. Unlike Redis, the register takes
29//! up full [`u8`] size instead of a raw int* and thus saves some
30//! tricky bit shifting techniques used in the original version.
31//! This results in a memory usage increase from 12Kib to 16Kib.
32//! Also only the dense version is adopted, so there's no automatic
33//! conversion, largely to simplify the code.
34//!
35//! This module also borrows some code structure from [pdatastructs.rs](https://github.com/crepererum/pdatastructs.rs/blob/3997ed50f6b6871c9e53c4c5e0f48f431405fc63/src/hyperloglog.rs).
36
37use std::hash::BuildHasher;
38use std::hash::Hash;
39use std::marker::PhantomData;
40
41/// The greater is P, the smaller the error.
42const HLL_P: usize = 14_usize;
43/// The number of bits of the hash value used determining the number of leading zeros
44const HLL_Q: usize = 64_usize - HLL_P;
45pub(crate) const NUM_REGISTERS: usize = 1_usize << HLL_P;
46/// Mask to obtain index into the registers
47const HLL_P_MASK: u64 = (NUM_REGISTERS as u64) - 1;
48
49#[derive(Clone, Debug)]
50pub(crate) struct HyperLogLog<T>
51where
52    T: Hash + ?Sized,
53{
54    registers: [u8; NUM_REGISTERS],
55    phantom: PhantomData<T>,
56}
57
58pub(crate) use datafusion_common::hash_utils::HLL_RANDOM_STATE as HLL_HASH_STATE;
59
60impl<T> Default for HyperLogLog<T>
61where
62    T: Hash + ?Sized,
63{
64    fn default() -> Self {
65        Self::new()
66    }
67}
68
69impl<T> HyperLogLog<T>
70where
71    T: Hash + ?Sized,
72{
73    /// Creates a new, empty HyperLogLog.
74    pub fn new() -> Self {
75        let registers = [0; NUM_REGISTERS];
76        Self::new_with_registers(registers)
77    }
78
79    /// Creates a HyperLogLog from already populated registers
80    /// note that this method should not be invoked in untrusted environment
81    /// because the internal structure of registers are not examined.
82    pub(crate) fn new_with_registers(registers: [u8; NUM_REGISTERS]) -> Self {
83        Self {
84            registers,
85            phantom: PhantomData,
86        }
87    }
88
89    /// The HLL hash state is shared through `datafusion_common::hash_utils`
90    /// so sketches remain compatible across accumulators.
91    #[inline]
92    fn hash_value(&self, obj: &T) -> u64 {
93        HLL_HASH_STATE.hash_one(obj)
94    }
95
96    /// Adds an element to the HyperLogLog.
97    pub fn add(&mut self, obj: &T) {
98        let hash = self.hash_value(obj);
99        self.add_hashed(hash);
100    }
101
102    /// Adds a pre-computed hash value directly to the HyperLogLog.
103    ///
104    /// The hash should be computed using [`HLL_HASH_STATE`], the same hasher used
105    /// by [`Self::add`].
106    #[inline]
107    pub(crate) fn add_hashed(&mut self, hash: u64) {
108        let index = (hash & HLL_P_MASK) as usize;
109        let p = ((hash >> HLL_P) | (1_u64 << HLL_Q)).trailing_zeros() + 1;
110        self.registers[index] = self.registers[index].max(p as u8);
111    }
112
113    /// Get the register histogram (each value in register index into
114    /// the histogram; u32 is enough because we only have 2**14=16384 registers
115    #[inline]
116    fn get_histogram(&self) -> [u32; HLL_Q + 2] {
117        let mut histogram = [0; HLL_Q + 2];
118        // hopefully this can be unrolled
119        for r in self.registers {
120            histogram[r as usize] += 1;
121        }
122        histogram
123    }
124
125    /// Merge the other [`HyperLogLog`] into this one
126    pub fn merge(&mut self, other: &HyperLogLog<T>) {
127        assert!(
128            self.registers.len() == other.registers.len(),
129            "unexpected got unequal register size, expect {}, got {}",
130            self.registers.len(),
131            other.registers.len()
132        );
133        for i in 0..self.registers.len() {
134            self.registers[i] = self.registers[i].max(other.registers[i]);
135        }
136    }
137
138    /// Guess the number of unique elements seen by the HyperLogLog.
139    pub fn count(&self) -> usize {
140        count_from_histogram(&self.get_histogram())
141    }
142}
143
144/// Compute `index` and `rho` (register value) for a precomputed hash, exactly as
145/// [`HyperLogLog::add_hashed`] does.
146#[inline]
147pub(crate) fn register_for_hash(hash: u64) -> (usize, u8) {
148    let index = (hash & HLL_P_MASK) as usize;
149    let rho = (((hash >> HLL_P) | (1_u64 << HLL_Q)).trailing_zeros() + 1) as u8;
150    (index, rho)
151}
152
153/// Estimate the cardinality of a set of precomputed hashes without
154/// materializing a full [`NUM_REGISTERS`]-byte register array.
155///
156/// This is equivalent to adding every hash to a fresh [`HyperLogLog`] via
157/// [`HyperLogLog::add_hashed`] and calling [`HyperLogLog::count`], but only does
158/// work proportional to the number of hashes. It is used to cheaply estimate the
159/// many small groups produced by a high-cardinality `GROUP BY`, where allocating
160/// and scanning a 16 KiB sketch per group would dominate the runtime.
161///
162/// `hashes` may contain duplicates (duplicate hashes are idempotent).
163pub(crate) fn count_from_hashes(hashes: &[u64]) -> usize {
164    if hashes.is_empty() {
165        return 0;
166    }
167    // For each touched register index keep the maximum rho. Sorting by
168    // (index, rho) groups equal indices together with the max rho last.
169    let mut idx_rho: Vec<(usize, u8)> =
170        hashes.iter().map(|&hash| register_for_hash(hash)).collect();
171    idx_rho.sort_unstable();
172
173    let mut histogram = [0u32; HLL_Q + 2];
174    let mut touched = 0u32;
175    let mut i = 0;
176    while i < idx_rho.len() {
177        let index = idx_rho[i].0;
178        let mut max_rho = idx_rho[i].1;
179        i += 1;
180        while i < idx_rho.len() && idx_rho[i].0 == index {
181            max_rho = idx_rho[i].1; // ascending rho => last is the max
182            i += 1;
183        }
184        histogram[max_rho as usize] += 1;
185        touched += 1;
186    }
187    // All remaining registers are still zero.
188    histogram[0] = NUM_REGISTERS as u32 - touched;
189    count_from_histogram(&histogram)
190}
191
192/// Apply the HyperLogLog cardinality estimator to a register histogram.
193#[inline]
194fn count_from_histogram(histogram: &[u32; HLL_Q + 2]) -> usize {
195    let m = NUM_REGISTERS as f64;
196    let mut z = m * hll_tau((m - histogram[HLL_Q + 1] as f64) / m);
197    for i in histogram[1..=HLL_Q].iter().rev() {
198        z += *i as f64;
199        z *= 0.5;
200    }
201    z += m * hll_sigma(histogram[0] as f64 / m);
202    (0.5 / 2_f64.ln() * m * m / z).round() as usize
203}
204
205/// Helper function sigma as defined in
206/// "New cardinality estimation algorithms for HyperLogLog sketches"
207/// Otmar Ertl, arXiv:1702.01284
208#[inline]
209fn hll_sigma(x: f64) -> f64 {
210    if x == 1. {
211        f64::INFINITY
212    } else {
213        let mut y = 1.0;
214        let mut z = x;
215        let mut x = x;
216        loop {
217            x *= x;
218            let z_prime = z;
219            z += x * y;
220            y += y;
221            if z_prime == z {
222                break;
223            }
224        }
225        z
226    }
227}
228
229/// Helper function tau as defined in
230/// "New cardinality estimation algorithms for HyperLogLog sketches"
231/// Otmar Ertl, arXiv:1702.01284
232#[inline]
233fn hll_tau(x: f64) -> f64 {
234    if x == 0.0 || x == 1.0 {
235        0.0
236    } else {
237        let mut y = 1.0;
238        let mut z = 1.0 - x;
239        let mut x = x;
240        loop {
241            x = x.sqrt();
242            let z_prime = z;
243            y *= 0.5;
244            z -= (1.0 - x).powi(2) * y;
245            if z_prime == z {
246                break;
247            }
248        }
249        z / 3.0
250    }
251}
252
253impl<T> AsRef<[u8]> for HyperLogLog<T>
254where
255    T: Hash + ?Sized,
256{
257    fn as_ref(&self) -> &[u8] {
258        &self.registers
259    }
260}
261
262impl<T> Extend<T> for HyperLogLog<T>
263where
264    T: Hash,
265{
266    fn extend<S: IntoIterator<Item = T>>(&mut self, iter: S) {
267        for elem in iter {
268            self.add(&elem);
269        }
270    }
271}
272
273impl<'a, T> Extend<&'a T> for HyperLogLog<T>
274where
275    T: 'a + Hash + ?Sized,
276{
277    fn extend<S: IntoIterator<Item = &'a T>>(&mut self, iter: S) {
278        for elem in iter {
279            self.add(elem);
280        }
281    }
282}
283
284#[cfg(test)]
285mod tests {
286    use super::{HyperLogLog, NUM_REGISTERS};
287
288    fn compare_with_delta(got: usize, expected: usize) {
289        let expected = expected as f64;
290        let diff = (got as f64) - expected;
291        let diff = diff.abs() / expected;
292        // times 6 because we want the tests to be stable
293        // so we allow a rather large margin of error
294        // this is adopted from redis's unit test version as well
295        let margin = 1.04 / ((NUM_REGISTERS as f64).sqrt()) * 6.0;
296        assert!(
297            diff <= margin,
298            "{} is not near {} percent of {} which is ({}, {})",
299            got,
300            margin,
301            expected,
302            expected * (1.0 - margin),
303            expected * (1.0 + margin)
304        );
305    }
306
307    macro_rules! sized_number_test {
308        ($SIZE: expr, $T: tt) => {{
309            let mut hll = HyperLogLog::<$T>::new();
310            for i in 0..$SIZE {
311                hll.add(&i);
312            }
313            compare_with_delta(hll.count(), $SIZE);
314        }};
315    }
316
317    macro_rules! typed_large_number_test {
318        ($SIZE: expr) => {{
319            sized_number_test!($SIZE, u64);
320            sized_number_test!($SIZE, u128);
321            sized_number_test!($SIZE, i64);
322            sized_number_test!($SIZE, i128);
323        }};
324    }
325
326    macro_rules! typed_number_test {
327        ($SIZE: expr) => {{
328            sized_number_test!($SIZE, u16);
329            sized_number_test!($SIZE, u32);
330            sized_number_test!($SIZE, i16);
331            sized_number_test!($SIZE, i32);
332            typed_large_number_test!($SIZE);
333        }};
334    }
335
336    #[test]
337    fn test_empty() {
338        let hll = HyperLogLog::<u64>::new();
339        assert_eq!(hll.count(), 0);
340    }
341
342    #[test]
343    fn test_one() {
344        let mut hll = HyperLogLog::<u64>::new();
345        hll.add(&1);
346        assert_eq!(hll.count(), 1);
347    }
348
349    #[test]
350    fn test_number_100() {
351        typed_number_test!(100);
352    }
353
354    #[test]
355    fn test_number_1k() {
356        typed_number_test!(1_000);
357    }
358
359    #[test]
360    fn test_number_10k() {
361        typed_number_test!(10_000);
362    }
363
364    #[test]
365    fn test_number_100k() {
366        typed_large_number_test!(100_000);
367    }
368
369    #[test]
370    fn test_number_1m() {
371        typed_large_number_test!(1_000_000);
372    }
373
374    #[test]
375    fn test_u8() {
376        let mut hll = HyperLogLog::<[u8]>::new();
377        for i in 0..1000 {
378            let s = i.to_string();
379            let b = s.as_bytes();
380            hll.add(b);
381        }
382        compare_with_delta(hll.count(), 1000);
383    }
384
385    #[test]
386    fn test_string() {
387        let mut hll = HyperLogLog::<String>::new();
388        hll.extend((0..1000).map(|i| i.to_string()));
389        compare_with_delta(hll.count(), 1000);
390    }
391
392    #[test]
393    fn test_empty_merge() {
394        let mut hll = HyperLogLog::<u64>::new();
395        hll.merge(&HyperLogLog::<u64>::new());
396        assert_eq!(hll.count(), 0);
397    }
398
399    #[test]
400    fn test_merge_overlapped() {
401        let mut hll = HyperLogLog::<String>::new();
402        hll.extend((0..1000).map(|i| i.to_string()));
403
404        let mut other = HyperLogLog::<String>::new();
405        other.extend((0..1000).map(|i| i.to_string()));
406
407        hll.merge(&other);
408        compare_with_delta(hll.count(), 1000);
409    }
410
411    #[test]
412    fn test_repetition() {
413        let mut hll = HyperLogLog::<u32>::new();
414        for i in 0..1_000_000 {
415            hll.add(&(i % 1000));
416        }
417        compare_with_delta(hll.count(), 1000);
418    }
419}
datafusion_functions_aggregate/hyperloglog.rs

datafusion_functions_aggregate/
hyperloglog.rs
