feat(theta): add theta union (#145)

Co-authored-by: tison <wander4096@gmail.com>

Author: ZENOTME

datasketches/src/bloom/sketch.rs

@@ -564,7 +564,7 @@ impl BloomFilter {
 
     /// Returns the estimated size of the filter in bytes
     pub fn estimated_size(&self) -> usize {
-        std::mem::size_of::<Self>() + self.bit_array.len() * std::mem::size_of::<u64>()
+        size_of::<Self>() + self.bit_array.len() * size_of::<u64>()
     }
 }
 

datasketches/src/common/mod.rs

@@ -22,7 +22,5 @@ mod resize;
 pub use self::num_std_dev::NumStdDev;
 pub use self::resize::ResizeFactor;
 
-#[cfg(feature = "theta")]
-pub(crate) mod binomial_bounds;
 #[cfg(any(feature = "cpc", feature = "hll"))]
 pub(crate) mod inv_pow2_table;

datasketches/src/cpc/pair_table.rs

@@ -249,6 +249,6 @@ impl PairTable {
 
     /// Returns the estimated size of the heap allocations in bytes
     pub fn estimated_size(&self) -> usize {
-        self.slots.capacity() * std::mem::size_of::<u32>()
+        self.slots.capacity() * size_of::<u32>()
     }
 }

datasketches/src/cpc/sketch.rs

@@ -460,7 +460,7 @@ impl CpcSketch {
                 .map(|t| t.estimated_size())
                 .unwrap_or(0);
 
-        std::mem::size_of::<Self>() + heap_size
+        size_of::<Self>() + heap_size
     }
 }
 

datasketches/src/hll/aux_map.rs

@@ -229,7 +229,7 @@ impl AuxMap {
 
     /// Returns the estimated size of the heap allocations in bytes
     pub fn estimated_size(&self) -> usize {
-        self.entries.len() * std::mem::size_of::<Coupon>()
+        self.entries.len() * size_of::<Coupon>()
     }
 }
 

datasketches/src/hll/container.rs

@@ -138,6 +138,6 @@ impl Container {
 
     /// Returns the estimated size of the heap allocations in bytes
     pub fn estimated_size(&self) -> usize {
-        self.coupons.len() * std::mem::size_of::<Coupon>()
+        self.coupons.len() * size_of::<Coupon>()
     }
 }

datasketches/src/hll/sketch.rs

@@ -434,7 +434,7 @@ impl HllSketch {
             Mode::Array8(arr) => arr.estimated_size(),
         };
 
-        std::mem::size_of::<Self>() + heap_size
+        size_of::<Self>() + heap_size
     }
 }
 

datasketches/src/lib.rs

@@ -45,6 +45,8 @@ pub mod hll;
 pub mod tdigest;
 #[cfg(feature = "theta")]
 pub mod theta;
+#[cfg(feature = "theta")]
+pub mod thetacommon;
 
 // common modules
 pub mod codec;

datasketches/src/tdigest/sketch.rs

@@ -796,9 +796,9 @@ impl TDigestMut {
 
     /// Returns the estimated size of the sketch in bytes
     pub fn estimated_size(&self) -> usize {
-        std::mem::size_of::<Self>()
-            + self.centroids.capacity() * std::mem::size_of::<Centroid>()
-            + self.buffer.capacity() * std::mem::size_of::<f64>()
+        size_of::<Self>()
+            + self.centroids.capacity() * size_of::<Centroid>()
+            + self.buffer.capacity() * size_of::<f64>()
     }
 }
 
@@ -1011,7 +1011,7 @@ impl TDigest {
 
     /// Returns the estimated size of the sketch in bytes
     pub fn estimated_size(&self) -> usize {
-        std::mem::size_of::<Self>() + self.centroids.capacity() * std::mem::size_of::<Centroid>()
+        size_of::<Self>() + self.centroids.capacity() * size_of::<Centroid>()
     }
 }
 

datasketches/src/theta/hash_table.rs

@@ -18,12 +18,8 @@
 use std::hash::Hash;
 use std::num::NonZeroU64;
 
-use crate::theta::raw_hash_table::RawHashTable;
-use crate::theta::raw_hash_table::RawHashTableEntry;
-#[cfg(test)]
-pub(crate) use crate::theta::raw_hash_table::starting_sub_multiple;
-#[cfg(test)]
-pub(crate) use crate::theta::raw_hash_table::starting_theta_from_sampling_probability;
+use crate::thetacommon::RawHashTableEntry;
+use crate::thetacommon::hash_table::RawHashTable;
 
 /// Specific hash table for theta sketch
 ///
@@ -34,16 +30,22 @@ pub(crate) use crate::theta::raw_hash_table::starting_theta_from_sampling_probab
 ///   update the theta to the k-th smallest entry.
 pub(super) type ThetaHashTable = RawHashTable<ThetaEntry>;
 
+/// A retained entry in a Theta sketch.
 #[derive(Debug, Clone, Copy)]
-pub(crate) struct ThetaEntry {
+pub struct ThetaEntry {
     hash: NonZeroU64,
 }
 
 impl ThetaEntry {
-    fn new(hash: u64) -> Self {
+    pub(crate) fn new(hash: u64) -> Self {
         let hash = NonZeroU64::new(hash).expect("hash must be non-zero");
         Self { hash }
     }
+
+    /// Return the hash used as this entry's key.
+    pub fn hash(&self) -> u64 {
+        self.hash.get()
+    }
 }
 
 impl RawHashTableEntry for ThetaEntry {
@@ -90,18 +92,20 @@ mod tests {
     use super::*;
     use crate::common::ResizeFactor;
     use crate::hash::DEFAULT_UPDATE_SEED;
-    use crate::theta::MAX_THETA;
-    use crate::theta::MIN_LG_K;
+    use crate::thetacommon::constants::MAX_THETA;
+    use crate::thetacommon::constants::MIN_LG_K;
+    use crate::thetacommon::hash_table::starting_sub_multiple;
+    use crate::thetacommon::hash_table::starting_theta_from_sampling_probability;
 
     #[test]
     fn test_new_hash_table() {
         let table = ThetaHashTable::new(8, ResizeFactor::X8, 1.0, DEFAULT_UPDATE_SEED);
 
         assert_eq!(
-            table.lg_cur_size,
+            table.lg_cur_size(),
             starting_sub_multiple(8 + 1, MIN_LG_K, ResizeFactor::X8.lg_value())
         );
-        assert_eq!(table.theta, starting_theta_from_sampling_probability(1.0));
+        assert_eq!(table.theta(), starting_theta_from_sampling_probability(1.0));
         assert_eq!(table.num_retained(), 0);
         assert!(table.is_empty());
         assert_eq!(table.iter().count(), 0);
@@ -119,7 +123,7 @@ mod tests {
         assert_ne!(hash1, hash2);
 
         // With low theta, update should be screened out.
-        table.theta = 1;
+        table.set_theta(1);
         assert!(!table.try_insert("test3"));
     }
 
@@ -136,7 +140,7 @@ mod tests {
         assert_eq!(table.num_retained(), 1);
 
         // Force screening and verify insertion fails
-        table.theta = 0;
+        table.set_theta(1); // Set theta to a low value to screen out new entries
         assert!(!table.try_insert("screened"));
         assert_eq!(table.num_retained(), 1);
         assert!(!table.is_empty());
@@ -174,7 +178,7 @@ mod tests {
         {
             let mut table = ThetaHashTable::new(8, ResizeFactor::X2, 1.0, DEFAULT_UPDATE_SEED);
 
-            assert_eq!(table.entries.len(), 32);
+            assert_eq!(table.num_entries(), 32);
 
             // Insert enough values to trigger resize (50% threshold)
             // Capacity = 32 * 0.5 = 16
@@ -183,14 +187,14 @@ mod tests {
             // Table should have resized and all values should be inserted
             assert!(table.num_retained() > 0);
             assert_eq!(table.num_retained(), inserted);
-            assert_eq!(table.entries.len(), 64);
+            assert_eq!(table.num_entries(), 64);
         }
 
         // Test different resize factors
         {
             let mut table = ThetaHashTable::new(8, ResizeFactor::X4, 1.0, DEFAULT_UPDATE_SEED);
 
-            assert_eq!(table.entries.len(), 32);
+            assert_eq!(table.num_entries(), 32);
 
             // Insert enough values to trigger resize (50% threshold)
             // Capacity = 32 * 0.5 = 16
@@ -199,17 +203,17 @@ mod tests {
             // Table should have resized and all values should be inserted
             assert!(table.num_retained() > 0);
             assert_eq!(table.num_retained(), inserted);
-            assert_eq!(table.entries.len(), 128);
+            assert_eq!(table.num_entries(), 128);
         }
     }
 
     #[test]
     fn test_rebuild() {
         let mut table = ThetaHashTable::new(5, ResizeFactor::X8, 1.0, DEFAULT_UPDATE_SEED);
 
-        assert_eq!(table.lg_cur_size, 6);
-        assert_eq!(table.entries.len(), 64);
-        assert_eq!(table.theta, MAX_THETA);
+        assert_eq!(table.lg_cur_size(), 6);
+        assert_eq!(table.num_entries(), 64);
+        assert_eq!(table.theta(), MAX_THETA);
 
         // Insert many values to trigger rebuild
         for i in 0..100 {
@@ -228,9 +232,9 @@ mod tests {
             let _ = table.try_insert(format!("value_{}", i));
         }
 
-        assert_eq!(table.lg_cur_size, 6);
-        assert!(table.entries.len() >= 64);
-        assert!(table.theta < new_theta);
+        assert_eq!(table.lg_cur_size(), 6);
+        assert!(table.num_entries() >= 64);
+        assert!(table.theta() < new_theta);
     }
 
     #[test]
@@ -274,8 +278,8 @@ mod tests {
     fn test_reset() {
         let mut table = ThetaHashTable::new(8, ResizeFactor::X8, 1.0, DEFAULT_UPDATE_SEED);
         let init_theta = table.theta();
-        let init_lg_cur = table.lg_cur_size;
-        let init_entries = table.entries.len();
+        let init_lg_cur = table.lg_cur_size();
+        let init_entries = table.num_entries();
 
         // Insert some values
         for i in 0..10 {
@@ -291,8 +295,8 @@ mod tests {
         assert!(table.is_empty());
         assert_eq!(table.num_retained(), 0);
         assert_eq!(table.theta(), init_theta);
-        assert_eq!(table.lg_cur_size, init_lg_cur);
-        assert_eq!(table.entries.len(), init_entries);
+        assert_eq!(table.lg_cur_size(), init_lg_cur);
+        assert_eq!(table.num_entries(), init_entries);
         assert_eq!(table.iter().count(), 0);
     }