feat: expose estimated_size for stateful set operations

CHANGELOG.md

@@ -18,6 +18,7 @@ All significant changes to this project will be documented in this file.
 * `FrequentItemsSketch` now supports borrowed-key updates via `update_ref` and `update_with_count_ref`, allowing sketches such as `FrequentItemsSketch<String>` to update from `&str` without allocating on existing-key hits. Frequency queries also accept borrowed key forms matching `Borrow<Q>`.
 * `FrequentItemsSketch` no longer requires item types to implement `Clone` for core updates, queries, and serialization. Custom `FrequentItemValue` implementations can now be non-`Clone`; APIs that return or merge owned items still require `Clone`.
 * `CountMinSketch` and `FrequentItemsSketch` now expose `estimated_size()`, reporting the in-memory footprint of the sketch in bytes, following the other sketches.
+* The stateful set operations `HllUnion`, `CpcUnion`, `ThetaUnion`, `ThetaIntersection`, `TupleUnion`, and `TupleIntersection` now expose `estimated_size()`, reporting the in-memory footprint of the operator's internal state in bytes.
 
 ### Bug fixes
 

datasketches/src/cpc/union.rs

@@ -335,6 +335,16 @@ impl CpcUnion {
             }
         }
     }
+
+    /// Returns the estimated size of the union in bytes.
+    pub fn estimated_size(&self) -> usize {
+        // The state's inline size is already covered by size_of::<Self>().
+        let heap_size = match &self.state {
+            UnionState::Accumulator(sketch) => sketch.estimated_size() - size_of::<CpcSketch>(),
+            UnionState::BitMatrix(matrix) => matrix.capacity() * size_of::<u64>(),
+        };
+        size_of::<Self>() + heap_size
+    }
 }
 
 // testing methods

datasketches/src/hll/union.rs

@@ -336,6 +336,12 @@ impl HllUnion {
     pub fn lower_bound(&self, num_std_dev: NumStdDev) -> f64 {
         self.gadget.lower_bound(num_std_dev)
     }
+
+    /// Returns the estimated size of the union in bytes.
+    pub fn estimated_size(&self) -> usize {
+        // The gadget's inline size is already covered by size_of::<Self>().
+        size_of::<Self>() - size_of::<HllSketch>() + self.gadget.estimated_size()
+    }
 }
 
 /// Convert a coupon mode (List or Set) to Hll8 target type

datasketches/src/thetafamily/common/intersection.rs

@@ -229,6 +229,11 @@ where
         self.has_result
     }
 
+    /// Returns the estimated size of the heap allocations in bytes.
+    pub fn estimated_size(&self) -> usize {
+        self.table.estimated_size()
+    }
+
     /// Return the current intersection state as compact-sketch parts.
     pub fn result(&self, ordered: bool) -> CompactSketchParts<E>
     where

datasketches/src/thetafamily/common/union.rs

@@ -152,6 +152,11 @@ where
         self.table.reset();
         self.union_theta = self.table.theta();
     }
+
+    /// Returns the estimated size of the heap allocations in bytes.
+    pub fn estimated_size(&self) -> usize {
+        self.table.estimated_size()
+    }
 }
 
 #[cfg(test)]

datasketches/src/thetafamily/theta/intersection.rs

@@ -67,6 +67,11 @@ impl ThetaIntersection {
         self.state.has_result()
     }
 
+    /// Returns the estimated size of the intersection in bytes.
+    pub fn estimated_size(&self) -> usize {
+        size_of::<Self>() + self.state.estimated_size()
+    }
+
     /// Returns the intersection result as a compact theta sketch.
     ///
     /// # Panics

datasketches/src/thetafamily/theta/union.rs

@@ -66,6 +66,11 @@ impl ThetaUnion {
     pub fn reset(&mut self) {
         self.state.reset();
     }
+
+    /// Returns the estimated size of the union in bytes.
+    pub fn estimated_size(&self) -> usize {
+        size_of::<Self>() + self.state.estimated_size()
+    }
 }
 
 /// Builder for [`ThetaUnion`].

datasketches/src/thetafamily/tuple/intersection.rs

@@ -130,6 +130,11 @@ where
         self.state.has_result()
     }
 
+    /// Returns the estimated size of the intersection in bytes.
+    pub fn estimated_size(&self) -> usize {
+        size_of::<Self>() + self.state.estimated_size()
+    }
+
     /// Returns the intersection result as a compact Tuple sketch.
     ///
     /// If `ordered` is true, retained entries are sorted ascending by hash.

datasketches/src/thetafamily/tuple/union.rs

@@ -114,6 +114,11 @@ where
     pub fn reset(&mut self) {
         self.state.reset();
     }
+
+    /// Returns the estimated size of the union in bytes.
+    pub fn estimated_size(&self) -> usize {
+        size_of::<Self>() + self.state.estimated_size()
+    }
 }
 
 /// Builder for [`TupleUnion`].

datasketches/tests/cpc_test/union.rs

@@ -179,3 +179,16 @@ fn test_lg_k_too_small() {
 fn test_lg_k_too_large() {
     CpcSketch::new(27);
 }
+
+#[test]
+fn test_union_estimated_size() {
+    let mut union = CpcUnion::new(11);
+    assert_eq!(union.estimated_size(), 112);
+
+    let mut sketch = CpcSketch::new(11);
+    for i in 0..1000 {
+        sketch.update(i);
+    }
+    union.update(&sketch);
+    assert_eq!(union.estimated_size(), 16496);
+}

datasketches/tests/hll_test/union.rs

@@ -617,3 +617,16 @@ fn test_union_validation() {
     union.reset();
     assert_eq!(union.lg_max_k(), 15, "lg_max_k should persist after reset");
 }
+
+#[test]
+fn test_union_estimated_size() {
+    let mut union = HllUnion::new(10);
+    assert_eq!(union.estimated_size(), 128);
+
+    let mut sketch = HllSketch::new(10, HllType::Hll8);
+    for i in 0..1000 {
+        sketch.update(i);
+    }
+    union.update(&sketch);
+    assert_eq!(union.estimated_size(), 1120);
+}

datasketches/tests/theta_test/intersection.rs

@@ -317,3 +317,13 @@ fn test_seed_mismatch_non_empty_returns_error() {
     let mut i = ThetaIntersection::with_seed(123);
     assert!(i.update(&s).is_err());
 }
+
+#[test]
+fn test_intersection_estimated_size() {
+    let mut intersection = ThetaIntersection::default();
+    assert_eq!(intersection.estimated_size(), 72);
+
+    let sketch = sketch_with_range(0, 1000);
+    intersection.update(&sketch).unwrap();
+    assert_eq!(intersection.estimated_size(), 16456);
+}

datasketches/tests/theta_test/union.rs

@@ -687,3 +687,13 @@ fn test_corner_case_union_states() {
         assert_eq!(compact_result.is_empty(), expected_empty);
     }
 }
+
+#[test]
+fn test_union_estimated_size() {
+    let mut union = ThetaUnionBuilder::default().build();
+    assert_eq!(union.estimated_size(), 1096);
+
+    let sketch = sketch_with_range(12, 0, 1000);
+    union.update(&sketch).unwrap();
+    assert_eq!(union.estimated_size(), 65608);
+}

datasketches/tests/tuple_test/intersection.rs

@@ -179,3 +179,13 @@ fn estimation_bounds_cover_the_true_intersection() {
         "expected 25000 in [{lower}, {upper}]"
     );
 }
+
+#[test]
+fn intersection_estimated_size_grows_with_updates() {
+    let mut intersection = TupleIntersection::new(SumPolicy);
+    assert_eq!(intersection.estimated_size(), 72);
+
+    let sketch = tuple_sketch_with_range(0, 1000);
+    intersection.update(&sketch).unwrap();
+    assert_eq!(intersection.estimated_size(), 32840);
+}

datasketches/tests/tuple_test/union.rs

@@ -151,3 +151,13 @@ fn estimation_bounds_cover_the_true_union() {
         "expected 75000 in [{lower}, {upper}]"
     );
 }
+
+#[test]
+fn union_estimated_size_grows_with_updates() {
+    let mut union = default_union_builder().build();
+    assert_eq!(union.estimated_size(), 2120);
+
+    let sketch = tuple_sketch_with_range(0, 1000);
+    union.update(&sketch).unwrap();
+    assert_eq!(union.estimated_size(), 131144);
+}