libdd_ddsketch/

lib.rs

// Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
// SPDX-License-Identifier: Apache-2.0

//! This crate defines a minimal implementation of DDSketch.
//!
//! DDSketch is a data sketch used to generate percentiles over streaming data using constant
//! memory. A DDSketch is essentially a histogram that partitions the range of positive values into
//! an infinite number of indexed bins whose size grows exponentially. It keeps track of the number
//! of values (or possibly floating-point weights) added to each bin. Negative values are
//! partitioned like positive values, symmetrically to zero. The value zero as well as its close
//! neighborhood that would be mapped to extreme bin indexes is mapped to a specific counter.

#![cfg_attr(not(test), deny(clippy::panic))]
#![cfg_attr(not(test), deny(clippy::unwrap_used))]
#![cfg_attr(not(test), deny(clippy::expect_used))]
#![cfg_attr(not(test), deny(clippy::todo))]
#![cfg_attr(not(test), deny(clippy::unimplemented))]

use std::collections::{HashMap, VecDeque};

use prost::Message;

/// Protobuf representation of DDSketch - generated by build.rs
#[rustfmt::skip]
pub mod pb;

/// This is a minimal DDSketch implementation
///
/// This implementation only supports a part of the standard (which is also only the parts dd
/// backend supports :shrug:)
/// - max length contiguous bin store, with lower bin collapse behavior.
/// - Positive or zero values
///
/// The default sketch has a 1% relative accuracy, and only accepts positive points
///
/// See <https://github.com/DataDog/sketches-go> for the reference implementation
#[derive(Debug, Default, Clone)]
pub struct DDSketch {
    store: LowCollapsingDenseStore, // Store the weight of each bin
    zero_count: f64,                // Store the weight of the bin of value 0
    mapping: LogMapping,            // Bin-Value mapping
}

impl DDSketch {
    /// Return an iterator over `(value, weight)` pair for each bin
    pub fn ordered_bins(&self) -> Vec<(f64, f64)> {
        let mut bins: Vec<_> = std::iter::once((0.0, self.zero_count))
            .chain(self.store.bins().map(|(b, v)| (self.mapping.value(b), v)))
            .collect();
        bins.sort_by(|a, b| a.0.total_cmp(&b.0));
        bins
    }

    // Return the number of points in the sketch
    pub fn count(&self) -> f64 {
        self.zero_count + self.store.bins.iter().sum::<f64>()
    }

    /// Add a point with value `point` to the sketch
    /// `point` must be positive
    pub fn add(&mut self, point: f64) -> Result<(), Box<dyn std::error::Error>> {
        self.add_with_count(point, 1.0)
    }

    /// Add `count` point with value `point` to the sketch
    /// `count` and `point` must be positive
    pub fn add_with_count(
        &mut self,
        point: f64,
        count: f64,
    ) -> Result<(), Box<dyn std::error::Error>> {
        if count.is_nan() || count.is_infinite() {
            return Err("count is invalid".into());
        }
        if point < 0.0 || point.is_nan() || point.is_infinite() {
            return Err("point is invalid".into());
        } else if point < self.mapping.min_indexable_value {
            self.zero_count += count;
        } else {
            let index = self.mapping.index(point);
            *self.store.bin_mut(index) += count;
        }
        Ok(())
    }

    /// Return a protobuf of the sketch
    pub fn into_pb(self) -> pb::DdSketch {
        let contiguous_bins: Vec<f64> = self.store.bins.into();
        pb::DdSketch {
            mapping: Some(pb::IndexMapping {
                gamma: self.mapping.gamma,
                index_offset: self.mapping.index_offset,
                interpolation: pb::index_mapping::Interpolation::None.into(),
            }),
            positive_values: Some(pb::Store {
                bin_counts: HashMap::new(),
                contiguous_bin_counts: contiguous_bins,
                contiguous_bin_index_offset: self.store.offset,
            }),
            zero_count: self.zero_count,
            negative_values: Some(pb::Store {
                bin_counts: HashMap::new(),
                contiguous_bin_counts: Vec::new(),
                contiguous_bin_index_offset: 0,
            }),
        }
    }

    /// Return a serialized protobuf of the sketch
    pub fn encode_to_vec(self) -> Vec<u8> {
        self.into_pb().encode_to_vec()
    }
}

/// A store mapping the bin indexes to their respective weights
///
/// Stores the weights as contiguousBinCounts, only the bins within `offset` and the highest
/// non empty bin are stored.
///
/// Stores the weights of a contiguous range of bins containing all non-empty bins. The range start
/// at index `offset` and end at index `offset + bins.len()`
///
/// The range of stored bins is updated when accessing the index of a bin out of the range
/// with [`Self::bin_idx_to_bin_idx()`]. If the `max_size` is reached the lower bins are
/// collapsed together to free space.
#[derive(Debug, Clone)]
struct LowCollapsingDenseStore {
    bins: VecDeque<f64>,
    offset: i32,
    max_size: i32,
}

impl LowCollapsingDenseStore {
    fn new(max_size: i32) -> Option<Self> {
        if max_size < 0 {
            return None;
        }
        Some(Self {
            bins: VecDeque::new(),
            offset: 0,
            max_size,
        })
    }

    /// Return an iterator over the bins
    /// The iterator yields pairs `(bin_index, count)`
    fn bins(&self) -> impl Iterator<Item = (i32, f64)> + '_ {
        self.bins
            .iter()
            .enumerate()
            .map(|(i, &v)| (i as i32 + self.offset, v))
    }

    /// Return a mutable reference to the bin at index `bin_index`
    fn bin_mut(&mut self, bin_index: i32) -> &mut f64 {
        let store_index = self.bin_idx_to_store_idx(bin_index);
        &mut self.bins[store_index]
    }

    /// Return the `store_index` of the bin at index `bin_index`
    fn bin_idx_to_store_idx(&mut self, bin_index: i32) -> usize {
        if self.bins.is_empty() {
            // If the bins are empty, start them at the index
            self.offset = bin_index;
            self.bins.push_back(0.0);
            return 0;
        }

        // General case
        // Bucket lower than the stored range
        if bin_index < self.offset {
            let additional_low_bins = self.offset - bin_index;
            debug_assert!(additional_low_bins >= 0);

            let additional_low_bins = std::cmp::min(
                additional_low_bins as usize,
                self.max_size as usize - self.bins.len(),
            );

            self.bins.reserve(additional_low_bins);
            for _ in 0..additional_low_bins {
                self.bins.push_front(0.0);
            }

            self.offset -= additional_low_bins as i32;
            0
        }
        // Bucket higher than the stored range
        else if self.offset + self.bins.len() as i32 <= bin_index {
            let bin_range_size = bin_index - self.offset + 1; // Number of bucket to store

            if bin_range_size > self.max_size {
                self.collapse_low_bins(bin_range_size - self.max_size);
            }
            debug_assert!(self.bins.len() as i32 <= self.max_size);

            let store_index = bin_index - self.offset;
            for _ in 0..(store_index - self.bins.len() as i32 + 1) {
                self.bins.push_back(0.0);
            }
            store_index as usize
        }
        // Bucket within the stored range
        else {
            (bin_index - self.offset) as usize
        }
    }

    /// Collapse the `bin_number` lowest bins
    fn collapse_low_bins(&mut self, bin_number: i32) {
        let mut count = 0.0;
        for _ in 0..bin_number {
            count += self.bins.pop_front().unwrap_or(0.0);
        }
        if let Some(lowest_bin) = self.bins.front_mut() {
            *lowest_bin += count;
        } else {
            self.bins.push_front(count);
        }
        self.offset += bin_number;
    }
}

impl Default for LowCollapsingDenseStore {
    fn default() -> Self {
        #[allow(clippy::unwrap_used)]
        Self::new(2048).unwrap()
    }
}

/// Logarithmic mapping of bucket index to value
#[derive(Debug, Clone, Copy)]
struct LogMapping {
    gamma: f64,
    multiplier: f64,
    min_indexable_value: f64,
    index_offset: f64,
}

impl LogMapping {
    fn new(gamma: f64, offset: f64) -> Option<Self> {
        if gamma <= 1.0 {
            return None;
        }
        let multiplier = Self::multiplier_from_gamma(gamma);
        Some(Self {
            gamma,
            multiplier,
            min_indexable_value: max(
                // So that the value representing the lowest bucket is >= std::f64::MIN_POSITIVE
                f64::MIN_POSITIVE * gamma,
                // Minimum value so that index >= i32::MIN
                ((i32::MIN as f64 - offset) / multiplier + 1.0).exp(),
            )?,
            index_offset: offset,
        })
    }

    /// Returns the multiplier used to convert ln to base-gamma logarithm
    fn multiplier_from_gamma(gamma: f64) -> f64 {
        1.0 / gamma.ln()
    }

    /// Returns the relative accuracy guaranteed by the mapping
    fn relative_accuracy(&self) -> f64 {
        1.0 - 2.0 / (1.0 + self.gamma)
    }

    /// Returns the index of the bucket containing `value`
    fn index(&self, value: f64) -> i32 {
        (value.ln() * self.multiplier + self.index_offset).floor() as i32
    }

    /// Returns the value representing the bucket at `index`
    fn value(&self, index: i32) -> f64 {
        ((index as f64 - self.index_offset) / self.multiplier).exp()
            * (1.0 + self.relative_accuracy())
    }
}

impl Default for LogMapping {
    fn default() -> Self {
        const RELATIVE_ACCURACY: f64 = 0.007751937984496124;
        const GAMMA: f64 = (1.0 + RELATIVE_ACCURACY) / (1.0 - RELATIVE_ACCURACY);

        const BACKEND_SKETCH_MIN_VALUE: f64 = 1e-9;
        // offset used in datadog's backend for sketches
        let offset: f64 = (1.0 - (BACKEND_SKETCH_MIN_VALUE.ln() / GAMMA.ln()).floor()) + 0.5;

        #[allow(clippy::unwrap_used)]
        Self::new(GAMMA, offset).unwrap()
    }
}

fn max(a: f64, b: f64) -> Option<f64> {
    if a.is_nan() || b.is_nan() {
        None
    } else if a > b {
        Some(a)
    } else {
        Some(b)
    }
}

#[cfg(test)]
mod test {
    use prost::Message;

    use super::*;

    macro_rules! assert_within {
        ($x:expr, $y:expr, $tolerance:expr) => {
            let diff = $x - $y;
            assert!(
                -$tolerance < diff && diff < $tolerance,
                "x: {} y: {}",
                $x,
                $y,
            );
        };
    }

    #[test]
    fn test_exponential_mapping_within_tolerances() {
        let mapping = LogMapping::default();

        let values: &[f64] = &[1e-30, 0.1, 2.0, 10.0, 25.0, 10000.0];
        for &value in values {
            let index = mapping.index(value);
            let value_bucket = mapping.value(index);

            assert_within!(value_bucket / value, 1.0, 0.01);
        }
    }

    #[test]
    fn test_exponential_mapping_relative_accuracy() {
        let mapping = LogMapping::default();

        assert_within!(
            mapping.relative_accuracy(),
            0.007751937984496138,
            f64::EPSILON
        );
    }

    #[test]
    fn test_sketch_add() {
        let mut sketch = DDSketch::default();
        let points: &[f64] = &[0.0, 1e-5, 0.1, 2.0, 10.0, 25.0, 10000.0];
        for (i, &point) in points.iter().enumerate() {
            assert!(sketch.add_with_count(point, i as f64 + 1.0).is_ok());
        }

        dbg!(sketch.store.bins.len(), sketch.store.offset);

        for (i, (value, count)) in sketch
            .ordered_bins()
            .into_iter()
            .filter(|(_, p)| *p != 0.0)
            .enumerate()
        {
            if points[i] == 0.0 {
                assert_within!(value, 0.0, f64::EPSILON);
                assert_within!(count, i as f64 + 1.0, f64::EPSILON);
            } else {
                assert_within!(value / points[i], 1.0, 0.01);
                assert_within!(count, i as f64 + 1.0, f64::EPSILON);
            }
        }
    }

    #[test]
    fn test_skecth_add_negative() {
        let mut sketch = DDSketch::default();
        assert!(sketch.add(-1.0).is_err());
    }

    #[test]
    fn test_skecth_add_nan() {
        let mut sketch = DDSketch::default();
        assert!(sketch.add(f64::NAN).is_err());
    }

    #[test]
    fn test_sketch_count_add_with_count() {
        let mut sketch = DDSketch::default();
        assert_within!(sketch.count(), 0.0, f64::EPSILON);

        let points: &[f64] = &[0.0, 1e-30, 0.1, 2.0, 10.0, 25.0, 10000.0];
        for (i, &point) in points.iter().enumerate() {
            assert!(sketch.add_with_count(point, i as f64).is_ok());
        }

        assert_within!(sketch.count(), 21.0, f64::EPSILON);
    }

    #[test]
    fn test_sketch_count_add() {
        let mut sketch = DDSketch::default();
        assert_within!(sketch.count(), 0.0, f64::EPSILON);

        let points: &[f64] = &[0.0, 1e-30, 0.1, 2.0, 10.0, 25.0, 10000.0];
        for &point in points.iter() {
            assert!(sketch.add(point).is_ok());
        }
        assert_within!(sketch.count(), 7.0, f64::EPSILON);

        assert!(sketch.add(1.0).is_ok());
        assert_within!(sketch.count(), 8.0, f64::EPSILON);

        for n in 0..100 {
            assert!(sketch.add(n as f64).is_ok());
        }
        assert_within!(sketch.count(), 108.0, f64::EPSILON);
    }

    #[test]
    fn test_skecth_encode() {
        let mut sketch = DDSketch::default();
        let points: &[f64] = &[0.0, 1e-30, 0.1, 2.0, 10.0, 25.0, 10000.0];
        for (i, &point) in points.iter().enumerate() {
            assert!(sketch.add_with_count(point, i as f64).is_ok());
        }

        let pb_sketch = sketch.into_pb().encode_to_vec();
        assert!(!pb_sketch.is_empty());
    }

    #[test]
    fn test_low_collapsing_store() {
        let mut store = LowCollapsingDenseStore::new(5).unwrap();

        // Test initial push up to capacity
        for i in 0..5 {
            *store.bin_mut(i + 10) = 1.0;
        }
        for (i, b) in store.bins().enumerate() {
            assert_eq!(b.0, i as i32 + 10);
            assert_eq!(b.1, 1.0)
        }

        // Indexing existing bins
        for i in 0..5 {
            *store.bin_mut(i + 10) += 1.0;
        }
        for (i, b) in store.bins().enumerate() {
            assert_eq!(b.0, i as i32 + 10);
            assert_eq!(b.1, 2.0)
        }
    }

    #[test]
    fn test_low_collapsing_store_low_bins_are_collapsed() {
        let mut store = LowCollapsingDenseStore::new(5).unwrap();

        // Test initial push up to capacity to max
        for i in 0..5 {
            *store.bin_mut(i + 10) = 1.0;
        }

        // Indexing low bins at max capacity
        for i in 0..3 {
            *store.bin_mut(i) += 1.0;
        }
        for (i, b) in store.bins().enumerate() {
            assert_eq!(b.0, i as i32 + 10);
            if i == 0 {
                assert_eq!(b.1, 4.0)
            } else {
                assert_eq!(b.1, 1.0)
            }
        }

        // Indexing higer bins collapses lower bins
        *store.bin_mut(15) = 1.0;
        for (i, b) in store.bins().enumerate() {
            assert_eq!(b.0, i as i32 + 11);
            if i == 0 {
                assert_eq!(b.1, 5.0)
            } else {
                assert_eq!(b.1, 1.0)
            }
        }
    }

    #[test]
    fn test_low_collapsing_store_up_expansion() {
        let mut store = LowCollapsingDenseStore::new(3).unwrap();

        *store.bin_mut(1) = 1.0;
        *store.bin_mut(3) = 1.0;
        assert_eq!(
            store.bins().collect::<Vec<_>>(),
            &[(1, 1.0), (2, 0.0), (3, 1.0)]
        )
    }

    #[test]
    fn test_low_collapsing_store_down_expansion() {
        let mut store = LowCollapsingDenseStore::new(3).unwrap();

        *store.bin_mut(3) = 1.0;
        *store.bin_mut(1) = 1.0;
        assert_eq!(
            store.bins().collect::<Vec<_>>(),
            &[(1, 1.0), (2, 0.0), (3, 1.0)]
        )
    }
}