feat(theta): add theta union

datasketches/src/lib.rs

@@ -54,3 +54,6 @@ pub mod hash_value;
 
 // private internal modules
 mod hash;
+/// Shared building blocks for Theta-family sketches.
+#[cfg(feature = "theta")]
+pub mod theta_common;

datasketches/src/theta/hash_table.rs

@@ -18,12 +18,12 @@
 use std::hash::Hash;
 use std::num::NonZeroU64;
 
-use crate::theta::raw_hash_table::RawHashTable;
-use crate::theta::raw_hash_table::RawHashTableEntry;
+use crate::theta_common::hash_table::RawHashTable;
+use crate::theta_common::hash_table::RawHashTableEntry;
 #[cfg(test)]
-pub(crate) use crate::theta::raw_hash_table::starting_sub_multiple;
+pub(crate) use crate::theta_common::hash_table::starting_sub_multiple;
 #[cfg(test)]
-pub(crate) use crate::theta::raw_hash_table::starting_theta_from_sampling_probability;
+pub(crate) use crate::theta_common::hash_table::starting_theta_from_sampling_probability;
 
 /// Specific hash table for theta sketch
 ///
@@ -34,16 +34,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 {

datasketches/src/theta/intersection.rs

@@ -122,7 +122,8 @@ impl ThetaIntersection {
                 self.table.hash_seed(),
                 self.table.is_empty(),
             );
-            for hash in sketch.iter() {
+            for entry in sketch.iter() {
+                let hash = entry.hash();
                 if !self.table.try_insert_hash(hash) {
                     return Err(Error::invalid_argument(
                         "Insert entries from sketch fail, possibly corrupted input sketch",
@@ -139,7 +140,8 @@ impl ThetaIntersection {
             let max_matches = self.table.num_retained().min(sketch.num_retained());
             let mut matched_entries = Vec::with_capacity(max_matches);
             let mut count = 0;
-            for hash in sketch.iter() {
+            for entry in sketch.iter() {
+                let hash = entry.hash();
                 if hash < self.table.theta() {
                     if self.table.contains_hash(hash) {
                         if matched_entries.len() == max_matches {
@@ -200,24 +202,15 @@ impl ThetaIntersection {
         self.is_valid
     }
 
-    /// Returns the intersection result as a compact theta sketch (ordered).
-    ///
-    /// # Panics
-    ///
-    /// Panics if called before the first [`update`](Self::update).
-    pub fn result(&self) -> CompactThetaSketch {
-        self.result_with_ordered(true)
-    }
-
     /// Returns the intersection result as a compact theta sketch.
     ///
     /// # Panics
     ///
     /// Panics if called before the first [`update`](Self::update).
-    pub fn result_with_ordered(&self, ordered: bool) -> CompactThetaSketch {
+    pub fn to_sketch(&self, ordered: bool) -> CompactThetaSketch {
         assert!(
             self.is_valid,
-            "ThetaIntersection::result() called before first update()"
+            "ThetaIntersection::to_sketch() called before first update()"
         );
         let mut hashes: Vec<u64> = self.table.iter().collect();
         if ordered {

datasketches/src/theta/mod.rs

@@ -42,29 +42,20 @@
 mod bit_pack;
 mod hash_table;
 mod intersection;
-mod raw_hash_table;
 mod serialization;
 mod sketch;
+mod union;
 
+pub use self::hash_table::ThetaEntry;
 pub use self::intersection::ThetaIntersection;
 pub use self::sketch::CompactThetaSketch;
 pub use self::sketch::ThetaSketch;
 pub use self::sketch::ThetaSketchBuilder;
 pub use self::sketch::ThetaSketchView;
-
-/// Maximum theta value (signed max for compatibility with Java)
-pub(crate) const MAX_THETA: u64 = i64::MAX as u64;
-/// Minimum log2 of K
-pub(crate) const MIN_LG_K: u8 = 5;
-/// Maximum log2 of K
-pub(crate) const MAX_LG_K: u8 = 26;
-/// Default log2 of K
-pub(crate) const DEFAULT_LG_K: u8 = 12;
-/// Resize threshold (0.5 = 50% load factor)
-pub(crate) const HASH_TABLE_RESIZE_THRESHOLD: f64 = 0.5;
-/// Rebuild threshold (15/16 = 93.75% load factor)
-pub(crate) const HASH_TABLE_REBUILD_THRESHOLD: f64 = 15.0 / 16.0;
-/// Stride hash bits (7 bits for stride calculation)
-pub(crate) const STRIDE_HASH_BITS: u8 = 7;
-/// Stride mask
-pub(crate) const STRIDE_MASK: u64 = (1 << STRIDE_HASH_BITS) - 1;
+pub use self::union::ThetaUnion;
+pub use self::union::ThetaUnionBuilder;
+pub(crate) use crate::theta_common::DEFAULT_LG_K;
+pub(crate) use crate::theta_common::HASH_TABLE_REBUILD_THRESHOLD;
+pub(crate) use crate::theta_common::MAX_LG_K;
+pub(crate) use crate::theta_common::MAX_THETA;
+pub(crate) use crate::theta_common::MIN_LG_K;

datasketches/src/theta/sketch.rs

@@ -42,46 +42,48 @@ use crate::theta::bit_pack::BitPacker;
 use crate::theta::bit_pack::BitUnpacker;
 use crate::theta::bit_pack::pack_bits_block;
 use crate::theta::bit_pack::unpack_bits_block;
+use crate::theta::hash_table::ThetaEntry;
 use crate::theta::hash_table::ThetaHashTable;
 use crate::theta::serialization;
 use crate::theta::serialization::V2_PREAMBLE_EMPTY;
 use crate::theta::serialization::V2_PREAMBLE_ESTIMATE;
 use crate::theta::serialization::V2_PREAMBLE_PRECISE;
-
-mod private {
-    use super::*;
-
-    // Sealed trait to prevent external implementations of ThetaSketchView.
-    pub trait Sealed {}
-
-    impl Sealed for ThetaSketch {}
-    impl Sealed for CompactThetaSketch {}
-}
+use crate::theta_common::sketch_view::RawThetaSketchView;
 
 /// Read-only view for Theta sketches.
 ///
 /// This trait provides a unified input abstraction for APIs that can accept either
 /// mutable [`ThetaSketch`] or immutable [`CompactThetaSketch`].
-pub trait ThetaSketchView: private::Sealed {
-    /// Returns the 16-bit seed hash.
-    fn seed_hash(&self) -> u16;
+pub trait ThetaSketchView: RawThetaSketchView<ThetaEntry> {}
 
-    /// Returns theta as `u64`.
-    fn theta64(&self) -> u64;
+impl<T: RawThetaSketchView<ThetaEntry>> ThetaSketchView for T {}
 
-    /// Returns true if this sketch is empty.
-    fn is_empty(&self) -> bool;
+impl crate::theta_common::sketch_view::private::Sealed for ThetaSketch {}
 
-    /// Returns an iterator over retained hash values.
-    fn iter(&self) -> impl Iterator<Item = u64> + '_;
+impl RawThetaSketchView<ThetaEntry> for ThetaSketch {
+    fn seed_hash(&self) -> u16 {
+        ThetaSketch::seed_hash(self)
+    }
+
+    fn theta64(&self) -> u64 {
+        ThetaSketch::theta64(self)
+    }
 
-    /// Returns number of retained hash values.
-    fn num_retained(&self) -> usize;
+    fn is_empty(&self) -> bool {
+        ThetaSketch::is_empty(self)
+    }
 
-    /// Returns whether retained entries are ordered in ascending order.
     fn is_ordered(&self) -> bool {
         false
     }
+
+    fn iter(&self) -> impl Iterator<Item = ThetaEntry> + '_ {
+        self.table.iter_entries().copied()
+    }
+
+    fn num_retained(&self) -> usize {
+        ThetaSketch::num_retained(self)
+    }
 }
 
 /// Mutable theta sketch for building from input data
@@ -191,7 +193,7 @@ impl ThetaSketch {
         self.table.reset();
     }
 
-    /// Return iterator over hash values
+    /// Return iterator over retained entries.
     ///
     /// # Examples
     ///
@@ -202,8 +204,8 @@ impl ThetaSketch {
     /// let mut iter = sketch.iter();
     /// assert!(iter.next().is_some());
     /// ```
-    pub fn iter(&self) -> impl Iterator<Item = u64> + '_ {
-        self.table.iter()
+    pub fn iter(&self) -> impl Iterator<Item = ThetaEntry> + '_ {
+        self.table.iter_entries().copied()
     }
 
     /// Return this sketch in compact (immutable) form.
@@ -220,7 +222,7 @@ impl ThetaSketch {
     /// assert_eq!(compact.num_retained(), 1);
     /// ```
     pub fn compact(&self, ordered: bool) -> CompactThetaSketch {
-        let mut entries: Vec<u64> = self.iter().collect();
+        let mut entries: Vec<u64> = self.iter().map(|entry| entry.hash()).collect();
 
         let empty = self.is_empty();
         let theta = if empty {
@@ -320,28 +322,6 @@ impl ThetaSketch {
     }
 }
 
-impl ThetaSketchView for ThetaSketch {
-    fn seed_hash(&self) -> u16 {
-        ThetaSketch::seed_hash(self)
-    }
-
-    fn theta64(&self) -> u64 {
-        ThetaSketch::theta64(self)
-    }
-
-    fn is_empty(&self) -> bool {
-        ThetaSketch::is_empty(self)
-    }
-
-    fn iter(&self) -> impl Iterator<Item = u64> + '_ {
-        ThetaSketch::iter(self)
-    }
-
-    fn num_retained(&self) -> usize {
-        ThetaSketch::num_retained(self)
-    }
-}
-
 /// Compact (immutable) theta sketch.
 ///
 /// This is the serialized-friendly form of a theta sketch: a compact array of retained hash values
@@ -420,9 +400,9 @@ impl CompactThetaSketch {
         self.seed_hash
     }
 
-    /// Return iterator over retained hash values.
-    pub fn iter(&self) -> impl Iterator<Item = u64> + '_ {
-        self.entries.iter().copied()
+    /// Return iterator over retained entries.
+    pub fn iter(&self) -> impl Iterator<Item = ThetaEntry> + '_ {
+        self.entries.iter().copied().map(ThetaEntry::new)
     }
 
     /// Returns the approximate lower error bound given the specified number of Standard Deviations.
@@ -900,7 +880,9 @@ impl CompactThetaSketch {
     }
 }
 
-impl ThetaSketchView for CompactThetaSketch {
+impl crate::theta_common::sketch_view::private::Sealed for CompactThetaSketch {}
+
+impl RawThetaSketchView<ThetaEntry> for CompactThetaSketch {
     fn seed_hash(&self) -> u16 {
         CompactThetaSketch::seed_hash(self)
     }
@@ -913,17 +895,17 @@ impl ThetaSketchView for CompactThetaSketch {
         CompactThetaSketch::is_empty(self)
     }
 
-    fn iter(&self) -> impl Iterator<Item = u64> + '_ {
-        CompactThetaSketch::iter(self)
+    fn is_ordered(&self) -> bool {
+        CompactThetaSketch::is_ordered(self)
+    }
+
+    fn iter(&self) -> impl Iterator<Item = ThetaEntry> + '_ {
+        self.entries.iter().copied().map(ThetaEntry::new)
     }
 
     fn num_retained(&self) -> usize {
         CompactThetaSketch::num_retained(self)
     }
-
-    fn is_ordered(&self) -> bool {
-        CompactThetaSketch::is_ordered(self)
-    }
 }
 
 /// Builder for ThetaSketch
@@ -1041,13 +1023,13 @@ mod tests {
     use super::*;
 
     fn sorted_theta_entries(sketch: &ThetaSketch) -> Vec<u64> {
-        let mut entries: Vec<u64> = sketch.iter().collect();
+        let mut entries: Vec<u64> = sketch.iter().map(|entry| entry.hash()).collect();
         entries.sort_unstable();
         entries
     }
 
     fn sorted_compact_entries(sketch: &CompactThetaSketch) -> Vec<u64> {
-        let mut entries: Vec<u64> = sketch.iter().collect();
+        let mut entries: Vec<u64> = sketch.iter().map(|entry| entry.hash()).collect();
         entries.sort_unstable();
         entries
     }