datafusion-functions-aggregate 55.0.0

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! # HyperLogLog
//!
//! `hyperloglog` is a module that contains a modified version
//! of [redis's implementation](https://github.com/redis/redis/blob/4930d19e70c391750479951022e207e19111eb55/src/hyperloglog.c)
//! with some modification based on strong assumption of usage
//! within datafusion, so that function can
//! be efficiently implemented.
//!
//! Specifically, like Redis's version, this HLL structure uses
//! 2**14 = 16384 registers, which means the standard error is
//! 1.04/(16384**0.5) = 0.8125%. Unlike Redis, the register takes
//! up full [`u8`] size instead of a raw int* and thus saves some
//! tricky bit shifting techniques used in the original version.
//! This results in a memory usage increase from 12Kib to 16Kib.
//! Also only the dense version is adopted, so there's no automatic
//! conversion, largely to simplify the code.
//!
//! This module also borrows some code structure from [pdatastructs.rs](https://github.com/crepererum/pdatastructs.rs/blob/3997ed50f6b6871c9e53c4c5e0f48f431405fc63/src/hyperloglog.rs).

use std::hash::BuildHasher;
use std::hash::Hash;
use std::marker::PhantomData;

/// The greater is P, the smaller the error.
const HLL_P: usize = 14_usize;
/// The number of bits of the hash value used determining the number of leading zeros
const HLL_Q: usize = 64_usize - HLL_P;
pub(crate) const NUM_REGISTERS: usize = 1_usize << HLL_P;
/// Mask to obtain index into the registers
const HLL_P_MASK: u64 = (NUM_REGISTERS as u64) - 1;

#[derive(Clone, Debug)]
pub(crate) struct HyperLogLog<T>
where
    T: Hash + ?Sized,
{
    registers: [u8; NUM_REGISTERS],
    phantom: PhantomData<T>,
}

pub(crate) use datafusion_common::hash_utils::HLL_RANDOM_STATE as HLL_HASH_STATE;

impl<T> Default for HyperLogLog<T>
where
    T: Hash + ?Sized,
{
    fn default() -> Self {
        Self::new()
    }
}

impl<T> HyperLogLog<T>
where
    T: Hash + ?Sized,
{
    /// Creates a new, empty HyperLogLog.
    pub fn new() -> Self {
        let registers = [0; NUM_REGISTERS];
        Self::new_with_registers(registers)
    }

    /// Creates a HyperLogLog from already populated registers
    /// note that this method should not be invoked in untrusted environment
    /// because the internal structure of registers are not examined.
    pub(crate) fn new_with_registers(registers: [u8; NUM_REGISTERS]) -> Self {
        Self {
            registers,
            phantom: PhantomData,
        }
    }

    /// The HLL hash state is shared through `datafusion_common::hash_utils`
    /// so sketches remain compatible across accumulators.
    #[inline]
    fn hash_value(&self, obj: &T) -> u64 {
        HLL_HASH_STATE.hash_one(obj)
    }

    /// Adds an element to the HyperLogLog.
    pub fn add(&mut self, obj: &T) {
        let hash = self.hash_value(obj);
        self.add_hashed(hash);
    }

    /// Adds a pre-computed hash value directly to the HyperLogLog.
    ///
    /// The hash should be computed using [`HLL_HASH_STATE`], the same hasher used
    /// by [`Self::add`].
    #[inline]
    pub(crate) fn add_hashed(&mut self, hash: u64) {
        let index = (hash & HLL_P_MASK) as usize;
        let p = ((hash >> HLL_P) | (1_u64 << HLL_Q)).trailing_zeros() + 1;
        self.registers[index] = self.registers[index].max(p as u8);
    }

    /// Get the register histogram (each value in register index into
    /// the histogram; u32 is enough because we only have 2**14=16384 registers
    #[inline]
    fn get_histogram(&self) -> [u32; HLL_Q + 2] {
        let mut histogram = [0; HLL_Q + 2];
        // hopefully this can be unrolled
        for r in self.registers {
            histogram[r as usize] += 1;
        }
        histogram
    }

    /// Merge the other [`HyperLogLog`] into this one
    pub fn merge(&mut self, other: &HyperLogLog<T>) {
        assert!(
            self.registers.len() == other.registers.len(),
            "unexpected got unequal register size, expect {}, got {}",
            self.registers.len(),
            other.registers.len()
        );
        for i in 0..self.registers.len() {
            self.registers[i] = self.registers[i].max(other.registers[i]);
        }
    }

    /// Guess the number of unique elements seen by the HyperLogLog.
    pub fn count(&self) -> usize {
        count_from_histogram(&self.get_histogram())
    }
}

/// Compute `index` and `rho` (register value) for a precomputed hash, exactly as
/// [`HyperLogLog::add_hashed`] does.
#[inline]
pub(crate) fn register_for_hash(hash: u64) -> (usize, u8) {
    let index = (hash & HLL_P_MASK) as usize;
    let rho = (((hash >> HLL_P) | (1_u64 << HLL_Q)).trailing_zeros() + 1) as u8;
    (index, rho)
}

/// Estimate the cardinality of a set of precomputed hashes without
/// materializing a full [`NUM_REGISTERS`]-byte register array.
///
/// This is equivalent to adding every hash to a fresh [`HyperLogLog`] via
/// [`HyperLogLog::add_hashed`] and calling [`HyperLogLog::count`], but only does
/// work proportional to the number of hashes. It is used to cheaply estimate the
/// many small groups produced by a high-cardinality `GROUP BY`, where allocating
/// and scanning a 16 KiB sketch per group would dominate the runtime.
///
/// `hashes` may contain duplicates (duplicate hashes are idempotent).
pub(crate) fn count_from_hashes(hashes: &[u64]) -> usize {
    if hashes.is_empty() {
        return 0;
    }
    // For each touched register index keep the maximum rho. Sorting by
    // (index, rho) groups equal indices together with the max rho last.
    let mut idx_rho: Vec<(usize, u8)> =
        hashes.iter().map(|&hash| register_for_hash(hash)).collect();
    idx_rho.sort_unstable();

    let mut histogram = [0u32; HLL_Q + 2];
    let mut touched = 0u32;
    let mut i = 0;
    while i < idx_rho.len() {
        let index = idx_rho[i].0;
        let mut max_rho = idx_rho[i].1;
        i += 1;
        while i < idx_rho.len() && idx_rho[i].0 == index {
            max_rho = idx_rho[i].1; // ascending rho => last is the max
            i += 1;
        }
        histogram[max_rho as usize] += 1;
        touched += 1;
    }
    // All remaining registers are still zero.
    histogram[0] = NUM_REGISTERS as u32 - touched;
    count_from_histogram(&histogram)
}

/// Apply the HyperLogLog cardinality estimator to a register histogram.
#[inline]
fn count_from_histogram(histogram: &[u32; HLL_Q + 2]) -> usize {
    let m = NUM_REGISTERS as f64;
    let mut z = m * hll_tau((m - histogram[HLL_Q + 1] as f64) / m);
    for i in histogram[1..=HLL_Q].iter().rev() {
        z += *i as f64;
        z *= 0.5;
    }
    z += m * hll_sigma(histogram[0] as f64 / m);
    (0.5 / 2_f64.ln() * m * m / z).round() as usize
}

/// Helper function sigma as defined in
/// "New cardinality estimation algorithms for HyperLogLog sketches"
/// Otmar Ertl, arXiv:1702.01284
#[inline]
fn hll_sigma(x: f64) -> f64 {
    if x == 1. {
        f64::INFINITY
    } else {
        let mut y = 1.0;
        let mut z = x;
        let mut x = x;
        loop {
            x *= x;
            let z_prime = z;
            z += x * y;
            y += y;
            if z_prime == z {
                break;
            }
        }
        z
    }
}

/// Helper function tau as defined in
/// "New cardinality estimation algorithms for HyperLogLog sketches"
/// Otmar Ertl, arXiv:1702.01284
#[inline]
fn hll_tau(x: f64) -> f64 {
    if x == 0.0 || x == 1.0 {
        0.0
    } else {
        let mut y = 1.0;
        let mut z = 1.0 - x;
        let mut x = x;
        loop {
            x = x.sqrt();
            let z_prime = z;
            y *= 0.5;
            z -= (1.0 - x).powi(2) * y;
            if z_prime == z {
                break;
            }
        }
        z / 3.0
    }
}

impl<T> AsRef<[u8]> for HyperLogLog<T>
where
    T: Hash + ?Sized,
{
    fn as_ref(&self) -> &[u8] {
        &self.registers
    }
}

impl<T> Extend<T> for HyperLogLog<T>
where
    T: Hash,
{
    fn extend<S: IntoIterator<Item = T>>(&mut self, iter: S) {
        for elem in iter {
            self.add(&elem);
        }
    }
}

impl<'a, T> Extend<&'a T> for HyperLogLog<T>
where
    T: 'a + Hash + ?Sized,
{
    fn extend<S: IntoIterator<Item = &'a T>>(&mut self, iter: S) {
        for elem in iter {
            self.add(elem);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{HyperLogLog, NUM_REGISTERS};

    fn compare_with_delta(got: usize, expected: usize) {
        let expected = expected as f64;
        let diff = (got as f64) - expected;
        let diff = diff.abs() / expected;
        // times 6 because we want the tests to be stable
        // so we allow a rather large margin of error
        // this is adopted from redis's unit test version as well
        let margin = 1.04 / ((NUM_REGISTERS as f64).sqrt()) * 6.0;
        assert!(
            diff <= margin,
            "{} is not near {} percent of {} which is ({}, {})",
            got,
            margin,
            expected,
            expected * (1.0 - margin),
            expected * (1.0 + margin)
        );
    }

    macro_rules! sized_number_test {
        ($SIZE: expr, $T: tt) => {{
            let mut hll = HyperLogLog::<$T>::new();
            for i in 0..$SIZE {
                hll.add(&i);
            }
            compare_with_delta(hll.count(), $SIZE);
        }};
    }

    macro_rules! typed_large_number_test {
        ($SIZE: expr) => {{
            sized_number_test!($SIZE, u64);
            sized_number_test!($SIZE, u128);
            sized_number_test!($SIZE, i64);
            sized_number_test!($SIZE, i128);
        }};
    }

    macro_rules! typed_number_test {
        ($SIZE: expr) => {{
            sized_number_test!($SIZE, u16);
            sized_number_test!($SIZE, u32);
            sized_number_test!($SIZE, i16);
            sized_number_test!($SIZE, i32);
            typed_large_number_test!($SIZE);
        }};
    }

    #[test]
    fn test_empty() {
        let hll = HyperLogLog::<u64>::new();
        assert_eq!(hll.count(), 0);
    }

    #[test]
    fn test_one() {
        let mut hll = HyperLogLog::<u64>::new();
        hll.add(&1);
        assert_eq!(hll.count(), 1);
    }

    #[test]
    fn test_number_100() {
        typed_number_test!(100);
    }

    #[test]
    fn test_number_1k() {
        typed_number_test!(1_000);
    }

    #[test]
    fn test_number_10k() {
        typed_number_test!(10_000);
    }

    #[test]
    fn test_number_100k() {
        typed_large_number_test!(100_000);
    }

    #[test]
    fn test_number_1m() {
        typed_large_number_test!(1_000_000);
    }

    #[test]
    fn test_u8() {
        let mut hll = HyperLogLog::<[u8]>::new();
        for i in 0..1000 {
            let s = i.to_string();
            let b = s.as_bytes();
            hll.add(b);
        }
        compare_with_delta(hll.count(), 1000);
    }

    #[test]
    fn test_string() {
        let mut hll = HyperLogLog::<String>::new();
        hll.extend((0..1000).map(|i| i.to_string()));
        compare_with_delta(hll.count(), 1000);
    }

    #[test]
    fn test_empty_merge() {
        let mut hll = HyperLogLog::<u64>::new();
        hll.merge(&HyperLogLog::<u64>::new());
        assert_eq!(hll.count(), 0);
    }

    #[test]
    fn test_merge_overlapped() {
        let mut hll = HyperLogLog::<String>::new();
        hll.extend((0..1000).map(|i| i.to_string()));

        let mut other = HyperLogLog::<String>::new();
        other.extend((0..1000).map(|i| i.to_string()));

        hll.merge(&other);
        compare_with_delta(hll.count(), 1000);
    }

    #[test]
    fn test_repetition() {
        let mut hll = HyperLogLog::<u32>::new();
        for i in 0..1_000_000 {
            hll.add(&(i % 1000));
        }
        compare_with_delta(hll.count(), 1000);
    }
}