feat: add theta sketch (part 2)

datasketches/src/error.rs

@@ -91,6 +91,10 @@ impl Error {
 
 // Convenience constructors for deserialization errors
 impl Error {
+    pub(crate) fn invalid_argument(msg: impl Into<String>) -> Self {
+        Self::new(ErrorKind::InvalidArgument, msg)
+    }
+
     pub(crate) fn deserial(msg: impl Into<String>) -> Self {
         Self::new(ErrorKind::InvalidData, msg)
     }

datasketches/src/hll/array4.rs

@@ -24,10 +24,10 @@ use super::aux_map::AuxMap;
 use crate::codec::SketchBytes;
 use crate::codec::SketchSlice;
 use crate::error::Error;
-use crate::hll::NumStdDev;
 use crate::hll::estimator::HipEstimator;
 use crate::hll::get_slot;
 use crate::hll::get_value;
+use crate::num_std_dev::NumStdDev;
 
 const AUX_TOKEN: u8 = 15;
 

datasketches/src/hll/array6.rs

@@ -24,10 +24,10 @@
 use crate::codec::SketchBytes;
 use crate::codec::SketchSlice;
 use crate::error::Error;
-use crate::hll::NumStdDev;
 use crate::hll::estimator::HipEstimator;
 use crate::hll::get_slot;
 use crate::hll::get_value;
+use crate::num_std_dev::NumStdDev;
 
 const VAL_MASK_6: u16 = 0x3F; // 6 bits: 0b0011_1111
 

datasketches/src/hll/array8.rs

@@ -23,10 +23,10 @@
 use crate::codec::SketchBytes;
 use crate::codec::SketchSlice;
 use crate::error::Error;
-use crate::hll::NumStdDev;
 use crate::hll::estimator::HipEstimator;
 use crate::hll::get_slot;
 use crate::hll::get_value;
+use crate::num_std_dev::NumStdDev;
 
 /// Core Array8 data structure - one byte per slot, no packing
 #[derive(Debug, Clone, PartialEq)]

datasketches/src/hll/container.rs

@@ -21,10 +21,10 @@
 //! cubic interpolation-based cardinality estimation and confidence bounds.
 
 use crate::hll::COUPON_RSE;
-use crate::hll::NumStdDev;
 use crate::hll::coupon_mapping::X_ARR;
 use crate::hll::coupon_mapping::Y_ARR;
 use crate::hll::cubic_interpolation::using_x_and_y_tables;
+use crate::num_std_dev::NumStdDev;
 
 /// Sentinel value indicating an empty coupon slot
 pub const COUPON_EMPTY: u32 = 0;

datasketches/src/hll/estimator.rs

@@ -25,6 +25,7 @@
 use crate::hll::composite_interpolation;
 use crate::hll::cubic_interpolation;
 use crate::hll::harmonic_numbers;
+use crate::num_std_dev::NumStdDev;
 
 /// HIP estimator with KxQ registers for improved cardinality estimation
 ///
@@ -324,22 +325,6 @@ fn inv_pow2(value: u8) -> f64 {
     }
 }
 
-/// Number of standard deviations for confidence bounds
-///
-/// This enum specifies the number of standard deviations to use when computing
-/// upper and lower bounds for cardinality estimates. Higher values provide wider
-/// confidence intervals with greater certainty that the true cardinality falls
-/// within the bounds.
-#[repr(u8)]
-pub enum NumStdDev {
-    /// One standard deviation (\~68% confidence interval)
-    One = 1,
-    /// Two standard deviations (\~95% confidence interval)
-    Two = 2,
-    /// Three standard deviations (\~99.7% confidence interval)
-    Three = 3,
-}
-
 /// Get relative error for HLL estimates
 ///
 /// This matches the implementation in datasketches-cpp HllUtil.hpp and RelativeErrorTables.hpp

datasketches/src/hll/mod.rs

@@ -77,7 +77,7 @@
 //! ```rust
 //! # use datasketches::hll::HllSketch;
 //! # use datasketches::hll::HllType;
-//! # use datasketches::hll::NumStdDev;
+//! # use datasketches::NumStdDev;
 //! let mut sketch = HllSketch::new(12, HllType::Hll8);
 //! sketch.update("apple");
 //! let upper = sketch.upper_bound(NumStdDev::Two);
@@ -124,7 +124,6 @@ mod serialization;
 mod sketch;
 mod union;
 
-pub use self::estimator::NumStdDev;
 pub use self::sketch::HllSketch;
 pub use self::union::HllUnion;
 

datasketches/src/hll/sketch.rs

@@ -25,7 +25,6 @@ use std::hash::Hash;
 use crate::codec::SketchSlice;
 use crate::error::Error;
 use crate::hll::HllType;
-use crate::hll::NumStdDev;
 use crate::hll::RESIZE_DENOMINATOR;
 use crate::hll::RESIZE_NUMERATOR;
 use crate::hll::array4::Array4;
@@ -37,6 +36,7 @@ use crate::hll::hash_set::HashSet;
 use crate::hll::list::List;
 use crate::hll::mode::Mode;
 use crate::hll::serialization::*;
+use crate::num_std_dev::NumStdDev;
 
 /// A HyperLogLog sketch.
 ///

datasketches/src/hll/union.rs

@@ -32,12 +32,12 @@ use std::hash::Hash;
 
 use crate::hll::HllSketch;
 use crate::hll::HllType;
-use crate::hll::NumStdDev;
 use crate::hll::array4::Array4;
 use crate::hll::array6::Array6;
 use crate::hll::array8::Array8;
 use crate::hll::mode::Mode;
 use crate::hll::pack_coupon;
+use crate::num_std_dev::NumStdDev;
 
 /// An HLL Union for combining multiple HLL sketches.
 ///

datasketches/src/lib.rs

@@ -38,8 +38,11 @@ pub mod hll;
 pub mod tdigest;
 pub mod theta;
 
+mod binomial_bounds;
 mod codec;
 mod hash;
+mod num_std_dev;
 mod resize;
 
+pub use self::num_std_dev::NumStdDev;
 pub use self::resize::ResizeFactor;

datasketches/src/num_std_dev.rs

@@ -0,0 +1,53 @@
+// 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.
+
+//! Standard deviation enums for confidence bounds
+//!
+//! This module provides types for specifying confidence levels when computing
+//! upper and lower bounds for sketch estimates.
+
+#[allow(clippy::excessive_precision)]
+const DELTA_OF_NUM_STD_DEVS: [f64; 4] = [
+    0.5000000000000000000, // = 0.5 (1 + erf(0))
+    0.1586553191586026479, // = 0.5 (1 + erf((-1/sqrt(2))))
+    0.0227502618904135701, // = 0.5 (1 + erf((-2/sqrt(2))))
+    0.0013498126861731796, // = 0.5 (1 + erf((-3/sqrt(2))))
+];
+
+/// Number of standard deviations for confidence bounds
+///
+/// This enum specifies the number of standard deviations to use when computing
+/// upper and lower bounds for cardinality estimates. Higher values provide wider
+/// confidence intervals with greater certainty that the true cardinality falls
+/// within the bounds.
+#[repr(u8)]
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum NumStdDev {
+    /// One standard deviation (\~68% confidence interval)
+    One = 1,
+    /// Two standard deviations (\~95% confidence interval)
+    Two = 2,
+    /// Three standard deviations (\~99.7% confidence interval)
+    Three = 3,
+}
+
+impl NumStdDev {
+    /// Returns the tail probability (delta) for this confidence level
+    pub fn tail_probability(&self) -> f64 {
+        DELTA_OF_NUM_STD_DEVS[*self as usize]
+    }
+}

datasketches/src/theta/hash_table.rs

@@ -355,7 +355,7 @@ mod tests {
         assert_ne!(hash1, hash2);
 
         // With low theta, some hashes should be filtered
-        table.theta = 0;
+        table.theta = 1;
         let hash3 = table.hash_and_screen("test3");
         assert_eq!(hash3, 0);
     }

datasketches/src/theta/sketch.rs

@@ -22,7 +22,9 @@
 
 use std::hash::Hash;
 
+use crate::NumStdDev;
 use crate::ResizeFactor;
+use crate::binomial_bounds;
 use crate::hash::DEFAULT_UPDATE_SEED;
 use crate::theta::hash_table::DEFAULT_LG_K;
 use crate::theta::hash_table::MAX_LG_K;
@@ -170,6 +172,79 @@ impl ThetaSketch {
     pub fn iter(&self) -> impl Iterator<Item = u64> + '_ {
         self.table.iter()
     }
+
+    /// Returns the approximate lower error bound given the specified number of Standard Deviations.
+    ///
+    /// # Arguments
+    ///
+    /// * `num_std_dev` - The number of standard deviations for confidence bounds.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use datasketches::NumStdDev;
+    /// use datasketches::theta::ThetaSketch;
+    ///
+    /// let mut sketch = ThetaSketch::builder().lg_k(12).build();
+    /// for i in 0..10000 {
+    ///     sketch.update(i);
+    /// }
+    ///
+    /// let estimate = sketch.estimate();
+    /// let lower_bound = sketch.lower_bound(NumStdDev::Two);
+    /// let upper_bound = sketch.upper_bound(NumStdDev::Two);
+    ///
+    /// assert!(lower_bound <= estimate);
+    /// assert!(estimate <= upper_bound);
+    /// ```
+    pub fn lower_bound(&self, num_std_dev: NumStdDev) -> f64 {
+        if !self.is_estimation_mode() {
+            return self.num_retained() as f64;
+        }
+        // This is safe because sampling_probability is guaranteed to be > 0,
+        // so theta will always be > 0, and binomial_bounds will never fail
+        binomial_bounds::lower_bound(self.num_retained() as u64, self.theta(), num_std_dev)
+            .expect("theta should always be valid")
+    }
+
+    /// Returns the approximate upper error bound given the specified number of Standard Deviations.
+    ///
+    /// # Arguments
+    ///
+    /// * `num_std_dev` - The number of standard deviations for confidence bounds.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use datasketches::NumStdDev;
+    /// use datasketches::theta::ThetaSketch;
+    ///
+    /// let mut sketch = ThetaSketch::builder().lg_k(12).build();
+    /// for i in 0..10000 {
+    ///     sketch.update(i);
+    /// }
+    ///
+    /// let estimate = sketch.estimate();
+    /// let lower_bound = sketch.lower_bound(NumStdDev::Two);
+    /// let upper_bound = sketch.upper_bound(NumStdDev::Two);
+    ///
+    /// assert!(lower_bound <= estimate);
+    /// assert!(estimate <= upper_bound);
+    /// ```
+    pub fn upper_bound(&self, num_std_dev: NumStdDev) -> f64 {
+        if !self.is_estimation_mode() {
+            return self.num_retained() as f64;
+        }
+        // This is safe because sampling_probability is guaranteed to be > 0,
+        // so theta will always be > 0, and binomial_bounds will never fail
+        binomial_bounds::upper_bound(
+            self.num_retained() as u64,
+            self.theta(),
+            num_std_dev,
+            self.is_empty(),
+        )
+        .expect("theta should always be valid")
+    }
 }
 
 /// Builder for ThetaSketch
@@ -226,9 +301,12 @@ impl ThetaSketchBuilder {
 
     /// Set sampling probability p.
     ///
+    /// The sampling probability controls the fraction of hashed values that are retained.
+    /// Must be greater than 0 to ensure valid theta values for bound calculations.
+    ///
     /// # Panics
     ///
-    /// If p is not in range [0.0, 1.0]
+    /// Panics if p is not in range (0.0, 1.0]
     ///
     /// # Examples
     ///
@@ -238,8 +316,8 @@ impl ThetaSketchBuilder {
     /// ```
     pub fn sampling_probability(mut self, probability: f32) -> Self {
         assert!(
-            (0.0..=1.0).contains(&probability),
-            "p must be in [0.0, 1.0], got {probability}"
+            (0.0..=1.0).contains(&probability) && probability > 0.0,
+            "sampling_probability must be in (0.0, 1.0], got {probability}"
         );
         self.sampling_probability = probability;
         self

datasketches/tests/hll_union_test.rs

@@ -27,10 +27,10 @@
 //!
 //! This mirrors the testing strategy used in hll_update_test.rs
 
+use datasketches::NumStdDev;
 use datasketches::hll::HllSketch;
 use datasketches::hll::HllType;
 use datasketches::hll::HllUnion;
-use datasketches::hll::NumStdDev;
 
 #[test]
 fn test_union_basic_operations() {

datasketches/tests/hll_update_test.rs

@@ -15,9 +15,9 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use datasketches::NumStdDev;
 use datasketches::hll::HllSketch;
 use datasketches::hll::HllType;
-use datasketches::hll::NumStdDev;
 
 #[test]
 fn test_basic_update() {