datasketches 0.3.0-rc.2

// 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.

use std::hash::Hash;
use std::hash::Hasher;

use crate::codec::SketchBytes;
use crate::codec::SketchSlice;
use crate::codec::assert::ensure_preamble_longs_in;
use crate::codec::assert::ensure_serial_version_is;
use crate::codec::assert::insufficient_data;
use crate::codec::family::Family;
use crate::countmin::CountMinValue;
use crate::countmin::UnsignedCountMinValue;
use crate::countmin::serialization::FLAGS_IS_EMPTY;
use crate::countmin::serialization::LONG_SIZE_BYTES;
use crate::countmin::serialization::PREAMBLE_LONGS_SHORT;
use crate::countmin::serialization::SERIAL_VERSION;
use crate::error::Error;
use crate::hash::DEFAULT_UPDATE_SEED;
use crate::hash::MurmurHash3X64128;
use crate::hash::compute_seed_hash;

const MAX_TABLE_ENTRIES: usize = 1 << 30;

/// Count-Min sketch for estimating item frequencies.
///
/// The sketch provides upper and lower bounds on estimated item frequencies
/// with configurable relative error and confidence.
#[derive(Debug, Clone, PartialEq)]
pub struct CountMinSketch<T: CountMinValue> {
    num_hashes: u8,
    num_buckets: u32,
    seed: u64,
    seed_hash: u16,
    total_weight: T,
    counts: Vec<T>,
    hash_seeds: Vec<u64>,
}

impl<T: CountMinValue> CountMinSketch<T> {
    /// Creates a new Count-Min sketch with the default seed.
    ///
    /// # Panics
    ///
    /// Panics if `num_hashes` is 0, `num_buckets` is less than 3, or the
    /// total table size exceeds the supported limit.
    ///
    /// # Examples
    ///
    /// ```
    /// # use datasketches::countmin::CountMinSketch;
    /// let sketch = CountMinSketch::<i64>::new(4, 128);
    /// assert_eq!(sketch.num_buckets(), 128);
    /// ```
    pub fn new(num_hashes: u8, num_buckets: u32) -> Self {
        Self::with_seed(num_hashes, num_buckets, DEFAULT_UPDATE_SEED)
    }

    /// Creates a new Count-Min sketch with the provided seed.
    ///
    /// # Panics
    ///
    /// Panics if any of:
    /// * `num_hashes` is 0
    /// * `num_buckets` is less than 3
    /// * the total table size exceeds the supported limit
    /// * the computed seed hash is zero
    ///
    /// # Examples
    ///
    /// ```
    /// # use datasketches::countmin::CountMinSketch;
    /// let sketch = CountMinSketch::<i64>::with_seed(4, 64, 42);
    /// assert_eq!(sketch.seed(), 42);
    /// ```
    pub fn with_seed(num_hashes: u8, num_buckets: u32, seed: u64) -> Self {
        let entries = entries_for_config(num_hashes, num_buckets);
        Self::make(num_hashes, num_buckets, seed, entries)
    }

    /// Returns the number of hash functions used by the sketch.
    pub fn num_hashes(&self) -> u8 {
        self.num_hashes
    }

    /// Returns the number of buckets per hash function.
    pub fn num_buckets(&self) -> u32 {
        self.num_buckets
    }

    /// Returns the seed used by the sketch.
    pub fn seed(&self) -> u64 {
        self.seed
    }

    /// Returns the total weight inserted into the sketch.
    pub fn total_weight(&self) -> T {
        self.total_weight
    }

    /// Returns the relative error (epsilon) implied by the number of buckets.
    pub fn relative_error(&self) -> f64 {
        std::f64::consts::E / self.num_buckets as f64
    }

    /// Returns true if the sketch has not seen any updates.
    pub fn is_empty(&self) -> bool {
        self.total_weight == T::ZERO
    }

    /// Suggests the number of buckets to achieve the given relative error.
    ///
    /// # Panics
    ///
    /// Panics if `relative_error` is negative.
    pub fn suggest_num_buckets(relative_error: f64) -> u32 {
        assert!(relative_error >= 0.0, "relative_error must be at least 0");
        (std::f64::consts::E / relative_error).ceil() as u32
    }

    /// Suggests the number of hashes to achieve the given confidence.
    ///
    /// # Panics
    ///
    /// Panics if `confidence` is not in `[0, 1]`.
    pub fn suggest_num_hashes(confidence: f64) -> u8 {
        assert!(
            (0.0..=1.0).contains(&confidence),
            "confidence must be between 0 and 1.0 (inclusive)"
        );
        if confidence == 1.0 {
            return 127;
        }
        let hashes = (1.0 / (1.0 - confidence)).ln().ceil();
        hashes.min(127.0) as u8
    }

    /// Updates the sketch with a single occurrence of the item.
    ///
    /// # Examples
    ///
    /// ```
    /// # use datasketches::countmin::CountMinSketch;
    /// let mut sketch = CountMinSketch::<i64>::new(4, 128);
    /// sketch.update("apple");
    /// assert!(sketch.estimate("apple") >= 1);
    /// ```
    pub fn update<I: Hash>(&mut self, item: I) {
        self.update_with_weight(item, T::ONE);
    }

    /// Updates the sketch with the given item and weight.
    ///
    /// # Examples
    ///
    /// ```
    /// # use datasketches::countmin::CountMinSketch;
    /// let mut sketch = CountMinSketch::<i64>::new(4, 128);
    /// sketch.update_with_weight("banana", 3);
    /// assert!(sketch.estimate("banana") >= 3);
    /// ```
    pub fn update_with_weight<I: Hash>(&mut self, item: I, weight: T) {
        if weight == T::ZERO {
            return;
        }
        let abs_weight = weight.abs();
        self.total_weight = self.total_weight.add(abs_weight);
        let num_buckets = self.num_buckets as usize;
        for (row, seed) in self.hash_seeds.iter().enumerate() {
            let bucket = self.bucket_index(&item, *seed);
            let index = row * num_buckets + bucket;
            self.counts[index] = self.counts[index].add(weight);
        }
    }

    /// Returns the estimated frequency of the given item.
    ///
    /// # Examples
    ///
    /// ```
    /// # use datasketches::countmin::CountMinSketch;
    /// let mut sketch = CountMinSketch::<i64>::new(4, 128);
    /// sketch.update_with_weight("pear", 2);
    /// assert!(sketch.estimate("pear") >= 2);
    /// ```
    pub fn estimate<I: Hash>(&self, item: I) -> T {
        let num_buckets = self.num_buckets as usize;
        let mut min = T::MAX;
        for (row, seed) in self.hash_seeds.iter().enumerate() {
            let bucket = self.bucket_index(&item, *seed);
            let index = row * num_buckets + bucket;
            let value = self.counts[index];
            if value < min {
                min = value;
            }
        }
        min
    }

    /// Returns the lower bound on the true frequency of the given item.
    pub fn lower_bound<I: Hash>(&self, item: I) -> T {
        self.estimate(item)
    }

    /// Returns the upper bound on the true frequency of the given item.
    pub fn upper_bound<I: Hash>(&self, item: I) -> T {
        let estimate = self.estimate(item);
        let error = T::from_f64(self.relative_error() * self.total_weight.to_f64());
        estimate.add(error)
    }

    /// Merges another sketch into this one.
    ///
    /// # Panics
    ///
    /// Panics if the sketches have incompatible configurations.
    ///
    /// # Examples
    ///
    /// ```
    /// # use datasketches::countmin::CountMinSketch;
    /// let mut left = CountMinSketch::<i64>::new(4, 128);
    /// let mut right = CountMinSketch::<i64>::new(4, 128);
    ///
    /// left.update("apple");
    /// right.update_with_weight("banana", 2);
    ///
    /// left.merge(&right);
    /// assert!(left.estimate("banana") >= 2);
    /// ```
    pub fn merge(&mut self, other: &CountMinSketch<T>) {
        if std::ptr::eq(self, other) {
            panic!("Cannot merge a sketch with itself.");
        }
        assert_eq!(self.num_hashes, other.num_hashes);
        assert_eq!(self.num_buckets, other.num_buckets);
        assert_eq!(self.seed, other.seed);
        assert_eq!(self.counts.len(), other.counts.len());
        let counts_len = self.counts.len();
        for i in 0..counts_len {
            self.counts[i] = self.counts[i].add(other.counts[i]);
        }
        self.total_weight = self.total_weight.add(other.total_weight);
    }

    /// Serializes this sketch into the DataSketches Count-Min format.
    ///
    /// # Examples
    ///
    /// ```
    /// # use datasketches::countmin::CountMinSketch;
    /// # let mut sketch = CountMinSketch::<i64>::new(4, 128);
    /// # sketch.update("apple");
    /// let bytes = sketch.serialize();
    /// let decoded = CountMinSketch::<i64>::deserialize(&bytes).unwrap();
    /// assert!(decoded.estimate("apple") >= 1);
    /// ```
    pub fn serialize(&self) -> Vec<u8> {
        let header_size = PREAMBLE_LONGS_SHORT as usize * LONG_SIZE_BYTES;
        let value_size = LONG_SIZE_BYTES;
        let payload_size = if self.is_empty() {
            0
        } else {
            value_size + (self.counts.len() * value_size)
        };
        let mut bytes = SketchBytes::with_capacity(header_size + payload_size);

        bytes.write_u8(PREAMBLE_LONGS_SHORT);
        bytes.write_u8(SERIAL_VERSION);
        bytes.write_u8(Family::COUNTMIN.id);
        bytes.write_u8(if self.is_empty() { FLAGS_IS_EMPTY } else { 0 });
        bytes.write_u32_le(0); // unused

        bytes.write_u32_le(self.num_buckets);
        bytes.write_u8(self.num_hashes);
        debug_assert_eq!(self.seed_hash, compute_seed_hash(self.seed));
        bytes.write_u16_le(self.seed_hash);
        bytes.write_u8(0);

        if self.is_empty() {
            return bytes.into_bytes();
        }

        bytes.write(&self.total_weight.to_bytes());
        for count in &self.counts {
            bytes.write(&count.to_bytes());
        }
        bytes.into_bytes()
    }

    /// Deserializes a sketch from bytes using the default seed.
    ///
    /// # Examples
    ///
    /// ```
    /// # use datasketches::countmin::CountMinSketch;
    /// # let mut sketch = CountMinSketch::<i64>::new(4, 64);
    /// # sketch.update("apple");
    /// # let bytes = sketch.serialize();
    /// let decoded = CountMinSketch::<i64>::deserialize(&bytes).unwrap();
    /// assert!(decoded.estimate("apple") >= 1);
    /// ```
    pub fn deserialize(bytes: &[u8]) -> Result<Self, Error> {
        Self::deserialize_with_seed(bytes, DEFAULT_UPDATE_SEED)
    }

    /// Deserializes a sketch from bytes using the provided seed.
    ///
    /// # Examples
    ///
    /// ```
    /// # use datasketches::countmin::CountMinSketch;
    /// # let mut sketch = CountMinSketch::<i64>::with_seed(4, 64, 7);
    /// # sketch.update("apple");
    /// # let bytes = sketch.serialize();
    /// let decoded = CountMinSketch::<i64>::deserialize_with_seed(&bytes, 7).unwrap();
    /// assert!(decoded.estimate("apple") >= 1);
    /// ```
    pub fn deserialize_with_seed(bytes: &[u8], seed: u64) -> Result<Self, Error> {
        fn read_value<T: CountMinValue>(
            cursor: &mut SketchSlice<'_>,
            tag: &'static str,
        ) -> Result<T, Error> {
            let mut bs = [0u8; 8];
            cursor.read_exact(&mut bs).map_err(insufficient_data(tag))?;
            T::try_from_bytes(bs)
        }

        let mut cursor = SketchSlice::new(bytes);
        let preamble_longs = cursor
            .read_u8()
            .map_err(insufficient_data("preamble_longs"))?;
        let serial_version = cursor
            .read_u8()
            .map_err(insufficient_data("serial_version"))?;
        let family_id = cursor.read_u8().map_err(insufficient_data("family_id"))?;
        let flags = cursor.read_u8().map_err(insufficient_data("flags"))?;
        cursor
            .read_u32_le()
            .map_err(insufficient_data("<unused>"))?;

        Family::COUNTMIN.validate_id(family_id)?;
        ensure_serial_version_is(SERIAL_VERSION, serial_version)?;
        ensure_preamble_longs_in(&[PREAMBLE_LONGS_SHORT], preamble_longs)?;

        let num_buckets = cursor
            .read_u32_le()
            .map_err(insufficient_data("num_buckets"))?;
        let num_hashes = cursor.read_u8().map_err(insufficient_data("num_hashes"))?;
        let seed_hash = cursor
            .read_u16_le()
            .map_err(insufficient_data("seed_hash"))?;
        cursor.read_u8().map_err(insufficient_data("unused8"))?;

        let expected_seed_hash = compute_seed_hash(seed);
        if seed_hash != expected_seed_hash {
            return Err(Error::deserial(format!(
                "incompatible seed hash: expected {expected_seed_hash}, got {seed_hash}",
            )));
        }

        let entries = entries_for_config_checked(num_hashes, num_buckets)?;
        let mut sketch = Self::make(num_hashes, num_buckets, seed, entries);
        if (flags & FLAGS_IS_EMPTY) != 0 {
            return Ok(sketch);
        }

        sketch.total_weight = read_value(&mut cursor, "total_weight")?;
        for count in &mut sketch.counts {
            *count = read_value(&mut cursor, "counts")?;
        }
        Ok(sketch)
    }

    fn make(num_hashes: u8, num_buckets: u32, seed: u64, entries: usize) -> Self {
        let counts = vec![T::ZERO; entries];
        let seed_hash = compute_seed_hash(seed);
        let hash_seeds = make_hash_seeds(seed, num_hashes);
        CountMinSketch {
            num_hashes,
            num_buckets,
            seed,
            seed_hash,
            total_weight: T::ZERO,
            counts,
            hash_seeds,
        }
    }

    fn bucket_index<I: Hash>(&self, item: &I, seed: u64) -> usize {
        let mut hasher = MurmurHash3X64128::with_seed(seed);
        item.hash(&mut hasher);
        let (h1, _) = hasher.finish128();
        (h1 % self.num_buckets as u64) as usize
    }
}

impl<T: UnsignedCountMinValue> CountMinSketch<T> {
    /// Divides every counter by two, truncating toward zero.
    ///
    /// Useful for exponential decay where counts represent recent activity.
    ///
    /// # Examples
    ///
    /// ```
    /// # use datasketches::countmin::CountMinSketch;
    /// let mut sketch = CountMinSketch::<u64>::new(4, 128);
    /// sketch.update_with_weight("apple", 3);
    /// sketch.halve();
    /// assert!(sketch.estimate("apple") >= 1);
    /// ```
    pub fn halve(&mut self) {
        for c in &mut self.counts {
            *c = c.halve()
        }
        self.total_weight = self.total_weight.halve();
    }

    /// Multiplies every counter by `decay` and truncates back into `T`.
    ///
    /// Values are truncated toward zero after multiplication; choose `decay` in `(0, 1]`.
    /// The total weight is scaled by the same factor to keep bounds consistent.
    ///
    /// # Examples
    ///
    /// ```
    /// # use datasketches::countmin::CountMinSketch;
    /// let mut sketch = CountMinSketch::<u64>::new(4, 128);
    /// sketch.update_with_weight("apple", 3);
    /// sketch.decay(0.5);
    /// assert!(sketch.estimate("apple") >= 1);
    /// ```
    ///
    /// # Panics
    ///
    /// Panics if `decay` is not finite or is outside `(0, 1]`.
    pub fn decay(&mut self, decay: f64) {
        assert!(decay > 0.0 && decay <= 1.0, "decay must be within (0, 1]");
        for c in &mut self.counts {
            *c = c.decay(decay)
        }
        self.total_weight = self.total_weight.decay(decay);
    }
}

fn entries_for_config(num_hashes: u8, num_buckets: u32) -> usize {
    assert!(num_hashes > 0, "num_hashes must be at least 1");
    assert!(num_buckets >= 3, "num_buckets must be at least 3");
    let entries = (num_hashes as usize)
        .checked_mul(num_buckets as usize)
        .expect("num_hashes * num_buckets overflows usize");
    assert!(
        entries < MAX_TABLE_ENTRIES,
        "num_hashes * num_buckets must be < {}",
        MAX_TABLE_ENTRIES
    );
    entries
}

fn entries_for_config_checked(num_hashes: u8, num_buckets: u32) -> Result<usize, Error> {
    if num_hashes == 0 {
        return Err(Error::deserial("num_hashes must be at least 1"));
    }
    if num_buckets < 3 {
        return Err(Error::deserial("num_buckets must be at least 3"));
    }
    let entries = (num_hashes as usize)
        .checked_mul(num_buckets as usize)
        .ok_or_else(|| Error::deserial("num_hashes * num_buckets overflows usize"))?;
    if entries >= MAX_TABLE_ENTRIES {
        return Err(Error::deserial(format!(
            "num_hashes * num_buckets must be < {MAX_TABLE_ENTRIES}",
        )));
    }
    Ok(entries)
}

fn make_hash_seeds(seed: u64, num_hashes: u8) -> Vec<u64> {
    let mut seeds = Vec::with_capacity(num_hashes as usize);
    for i in 0..num_hashes {
        // Derive per-row hash seeds deterministically from the sketch seed.
        let mut hasher = MurmurHash3X64128::with_seed(seed);
        hasher.write(&u64::from(i).to_le_bytes());
        let (h1, _) = hasher.finish128();
        seeds.push(h1);
    }
    seeds
}