feat: add theta sketch (part 2)

datasketches/src/common/mod.rs

@@ -0,0 +1,27 @@
+// 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.
+
+//! Data structures and functions that may be used across all the sketch families.
+
+// public common components for datasketches crate
+mod num_std_dev;
+mod resize;
+pub use self::num_std_dev::NumStdDev;
+pub use self::resize::ResizeFactor;
+
+// private to datasketches crate
+pub(crate) mod binomial_bounds;

datasketches/src/common/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 const fn tail_probability(&self) -> f64 {
+        DELTA_OF_NUM_STD_DEVS[*self as usize]
+    }
+}

datasketches/src/resize.rs → datasketches/src/common/resize.rs

@@ -36,7 +36,7 @@
 /// # Examples
 ///
 /// ```
-/// # use datasketches::ResizeFactor;
+/// # use datasketches::common::ResizeFactor;
 /// let factor = ResizeFactor::X4;
 /// assert_eq!(factor.value(), 4);
 /// assert_eq!(factor.lg_value(), 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

@@ -23,8 +23,8 @@
 use super::aux_map::AuxMap;
 use crate::codec::SketchBytes;
 use crate::codec::SketchSlice;
+use crate::common::NumStdDev;
 use crate::error::Error;
-use crate::hll::NumStdDev;
 use crate::hll::estimator::HipEstimator;
 use crate::hll::get_slot;
 use crate::hll::get_value;

datasketches/src/hll/array6.rs

@@ -23,8 +23,8 @@
 
 use crate::codec::SketchBytes;
 use crate::codec::SketchSlice;
+use crate::common::NumStdDev;
 use crate::error::Error;
-use crate::hll::NumStdDev;
 use crate::hll::estimator::HipEstimator;
 use crate::hll::get_slot;
 use crate::hll::get_value;

datasketches/src/hll/array8.rs

@@ -22,8 +22,8 @@
 
 use crate::codec::SketchBytes;
 use crate::codec::SketchSlice;
+use crate::common::NumStdDev;
 use crate::error::Error;
-use crate::hll::NumStdDev;
 use crate::hll::estimator::HipEstimator;
 use crate::hll::get_slot;
 use crate::hll::get_value;

datasketches/src/hll/container.rs

@@ -20,8 +20,8 @@
 //! Provides a simple array-based storage for coupons (hash values) with
 //! cubic interpolation-based cardinality estimation and confidence bounds.
 
+use crate::common::NumStdDev;
 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;

datasketches/src/hll/estimator.rs

@@ -22,6 +22,7 @@
 //! This is more accurate than the standard HLL estimator, especially for
 //! moderate cardinalities.
 
+use crate::common::NumStdDev;
 use crate::hll::composite_interpolation;
 use crate::hll::cubic_interpolation;
 use crate::hll::harmonic_numbers;
@@ -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::common::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

@@ -23,9 +23,9 @@
 use std::hash::Hash;
 
 use crate::codec::SketchSlice;
+use crate::common::NumStdDev;
 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;

datasketches/src/hll/union.rs

@@ -30,9 +30,9 @@
 
 use std::hash::Hash;
 
+use crate::common::NumStdDev;
 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;

datasketches/src/lib.rs

@@ -31,6 +31,7 @@
 compile_error!("datasketches does not support big-endian targets");
 
 pub mod bloom;
+pub mod common;
 pub mod countmin;
 pub mod error;
 pub mod frequencies;
@@ -40,6 +41,3 @@ pub mod theta;
 
 mod codec;
 mod hash;
-mod resize;
-
-pub use self::resize::ResizeFactor;

datasketches/src/theta/hash_table.rs

@@ -17,7 +17,7 @@
 
 use std::hash::Hash;
 
-use crate::ResizeFactor;
+use crate::common::ResizeFactor;
 use crate::hash::MurmurHash3X64128;
 
 /// Maximum theta value (signed max for compatibility with Java)
@@ -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::ResizeFactor;
+use crate::common::NumStdDev;
+use crate::common::ResizeFactor;
+use crate::common::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::common::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::common::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::common::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::common::NumStdDev;
 use datasketches::hll::HllSketch;
 use datasketches::hll::HllType;
-use datasketches::hll::NumStdDev;
 
 #[test]
 fn test_basic_update() {