feat(tuple): add tuple union operation

Author: ariesdevil

datasketches/src/thetacommon/hash_table.rs

@@ -32,12 +32,6 @@ use crate::thetacommon::constants::STRIDE_MASK;
 ///
 /// The entry type supplies the retained hash and any sketch-specific payload. The table owns all
 /// theta screening, probing, resizing, rebuilding, trimming, and logical-empty state.
-///
-/// It maintains an array with capacity up to 2^lg_max_size:
-/// * Before it reaches the max capacity, it will extend the array based on resize_factor.
-/// * After it reaches the capacity bigger than 2^lg_nom_size, every time the number of entries
-///   exceeds the threshold, it will rebuild the table: only keep the min 2^lg_nom_size entries and
-///   update the theta to the k-th smallest entry.
 #[derive(Debug)]
 pub struct RawHashTable<E> {
     lg_cur_size: u8,
@@ -132,16 +126,6 @@ where
 
     /// Inserts or updates the entry slot for a pre-hashed key.
     ///
-    /// The callback `f` is invoked with the current entry for `hash`:
-    /// * `Some(existing)` if the key is already retained. The callback should modify it in place;
-    ///   its return value is ignored.
-    /// * `None` if the key is new. The callback returns `Some(entry)` to insert it, or `None` to
-    ///   decline insertion.
-    ///
-    /// Using a single callback ensures any captured update value is consumed exactly once, so it
-    /// works for both the update sketch (folding an update value) and set operations (merging an
-    /// incoming entry) without requiring the value to be `Copy` or `Clone`.
-    ///
     /// Returns true if a new entry was created, false otherwise (existing key, declined insertion,
     /// or a hash screened out by theta).
     pub fn upsert_entry<F>(&mut self, hash: u64, f: F) -> bool
@@ -302,7 +286,6 @@ where
     }
 
     /// Set empty flag.
-    #[allow(dead_code)] // used by the set operations; tuple set operations arrive in follow-ups
     pub fn set_empty(&mut self, is_empty: bool) {
         self.is_empty = is_empty;
     }

datasketches/src/thetacommon/mod.rs

@@ -20,7 +20,7 @@
 pub(crate) mod binomial_bounds;
 pub(crate) mod constants;
 pub(crate) mod hash_table;
-#[cfg(feature = "theta")]
+#[cfg(any(feature = "theta", feature = "tuple"))]
 pub(crate) mod union;
 
 /// An entry retained by a Theta sketch family hash table.

datasketches/src/tuple/hash_table.rs

@@ -52,6 +52,11 @@ impl<S> TupleEntry<S> {
     pub fn summary(&self) -> &S {
         &self.summary
     }
+
+    /// Returns a mutable reference to the summary stored in this entry.
+    pub(super) fn summary_mut(&mut self) -> &mut S {
+        &mut self.summary
+    }
 }
 
 /// Specific hash table for tuple sketch.

datasketches/src/tuple/mod.rs

@@ -22,8 +22,9 @@
 //! mechanics (theta screening, resize, rebuild to nominal size k) mirror the Theta sketch, with the
 //! added requirement that colliding keys merge their summaries.
 //!
-//! The behavior of a summary (how to create and update it) is supplied externally through
-//! policy objects ([`SummaryUpdatePolicy`]) rather than being baked into the summary type.
+//! The behavior of a summary (how to create, update, and combine it) is supplied externally through
+//! policy objects ([`SummaryUpdatePolicy`] and [`SummaryCombinePolicy`]) rather than being baked
+//! into the summary type.
 //!
 //! # Usage
 //!
@@ -39,13 +40,18 @@ mod policy;
 mod serde;
 mod serialization;
 mod sketch;
+mod union;
 
 pub use self::hash_table::TupleEntry;
+pub use self::policy::DefaultUnionPolicy;
 pub use self::policy::DefaultUpdatePolicy;
+pub use self::policy::SummaryCombinePolicy;
 pub use self::policy::SummaryUpdatePolicy;
 pub use self::serde::PrimitiveSummarySerde;
 pub use self::serde::SummarySerde;
 pub use self::sketch::CompactTupleSketch;
 pub use self::sketch::TupleSketch;
 pub use self::sketch::TupleSketchBuilder;
 pub use self::sketch::TupleSketchView;
+pub use self::union::TupleUnion;
+pub use self::union::TupleUnionBuilder;

datasketches/src/tuple/policy.rs

@@ -25,6 +25,22 @@
 
 use std::ops::AddAssign;
 
+use crate::thetacommon::union::RawThetaUnionPolicy;
+use crate::tuple::hash_table::TupleEntry;
+
+/// Adapts a [`SummaryCombinePolicy`] to the raw union's entry-merge policy.
+#[derive(Debug)]
+pub(crate) struct CombinePolicyAdapter<P>(pub(crate) P);
+
+impl<S, P> RawThetaUnionPolicy<TupleEntry<S>> for CombinePolicyAdapter<P>
+where
+    P: SummaryCombinePolicy<S>,
+{
+    fn merge(&self, existing: &mut TupleEntry<S>, incoming: TupleEntry<S>) {
+        self.0.combine(existing.summary_mut(), incoming.summary());
+    }
+}
+
 /// Defines how a summary is created and how update values are folded into it.
 ///
 /// This is used by the update tuple sketch. `S` is the stored summary type and `U` is the type of
@@ -40,11 +56,20 @@ pub trait SummaryUpdatePolicy<S, U> {
     fn update(&self, summary: &mut S, value: U);
 }
 
+/// Defines how two summaries that share the same key are combined.
+///
+/// This is used by both union and intersection. Each operator is given its own policy instance,
+/// because the two operations may combine summaries differently for the same summary type.
+pub trait SummaryCombinePolicy<S> {
+    /// Combines `other` into `summary` in place.
+    fn combine(&self, summary: &mut S, other: &S);
+}
+
 /// Default update policy for summaries that are default-constructible and additive.
 ///
-/// This is the convenience policy used when no custom policy is supplied, equivalent to C++
-/// `default_tuple_update_policy` (which folds updates with `summary += update`). It is available
-/// for any summary type `S` and update type `U` where `S: Default + AddAssign<U>`.
+/// This is the convenience policy used when no custom policy is supplied; it folds updates with
+/// `summary += update`. It is available for any summary type `S` and update type `U` where
+/// `S: Default + AddAssign<U>`.
 #[derive(Debug, Default, Clone, Copy)]
 pub struct DefaultUpdatePolicy;
 
@@ -61,6 +86,26 @@ where
     }
 }
 
+/// Default combine policy for additive summaries, used by the union when no custom policy is given.
+///
+/// It combines two summaries with `summary += other` and is available for any summary type `S`
+/// where `S: AddAssign<&S>`.
+///
+/// There is intentionally no default combine policy for the intersection: how to combine summaries
+/// of the keys present in both inputs is application-specific, so the intersection always requires
+/// an explicit policy.
+#[derive(Debug, Default, Clone, Copy)]
+pub struct DefaultUnionPolicy;
+
+impl<S> SummaryCombinePolicy<S> for DefaultUnionPolicy
+where
+    for<'a> S: AddAssign<&'a S>,
+{
+    fn combine(&self, summary: &mut S, other: &S) {
+        *summary += other;
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -89,6 +134,12 @@ mod tests {
         }
     }
 
+    impl SummaryCombinePolicy<u64> for MaxPolicy {
+        fn combine(&self, summary: &mut u64, other: &u64) {
+            self.update(summary, *other);
+        }
+    }
+
     #[test]
     fn custom_update_policy_keeps_max() {
         let policy = MaxPolicy;
@@ -98,4 +149,22 @@ mod tests {
         policy.update(&mut summary, 2);
         assert_eq!(summary, 7);
     }
+
+    #[test]
+    fn custom_combine_policy_keeps_max() {
+        let policy = MaxPolicy;
+        let mut summary = 5u64;
+        policy.combine(&mut summary, &10);
+        policy.combine(&mut summary, &7);
+        assert_eq!(summary, 10);
+    }
+
+    #[test]
+    fn default_union_policy_combines_additively() {
+        let policy = DefaultUnionPolicy;
+        let mut summary = 5u64;
+        policy.combine(&mut summary, &10);
+        policy.combine(&mut summary, &7);
+        assert_eq!(summary, 22);
+    }
 }