diff --git a/CHANGELOG.md b/CHANGELOG.md index 997607e..ce8552a 100644 --- a/CHANGELOG.md +++ b/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` to update from `&str` without allocating on existing-key hits. Frequency queries also accept borrowed key forms matching `Borrow`. * `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 diff --git a/datasketches/src/cpc/union.rs b/datasketches/src/cpc/union.rs index 9dc2a6c..5cf74d7 100644 --- a/datasketches/src/cpc/union.rs +++ b/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::(). + let heap_size = match &self.state { + UnionState::Accumulator(sketch) => sketch.estimated_size() - size_of::(), + UnionState::BitMatrix(matrix) => matrix.capacity() * size_of::(), + }; + size_of::() + heap_size + } } // testing methods diff --git a/datasketches/src/hll/union.rs b/datasketches/src/hll/union.rs index 4de24d9..87da53e 100644 --- a/datasketches/src/hll/union.rs +++ b/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::(). + size_of::() - size_of::() + self.gadget.estimated_size() + } } /// Convert a coupon mode (List or Set) to Hll8 target type diff --git a/datasketches/src/thetafamily/common/intersection.rs b/datasketches/src/thetafamily/common/intersection.rs index 3707c85..61cd74a 100644 --- a/datasketches/src/thetafamily/common/intersection.rs +++ b/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 where diff --git a/datasketches/src/thetafamily/common/union.rs b/datasketches/src/thetafamily/common/union.rs index 642de04..c22771f 100644 --- a/datasketches/src/thetafamily/common/union.rs +++ b/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)] diff --git a/datasketches/src/thetafamily/theta/intersection.rs b/datasketches/src/thetafamily/theta/intersection.rs index e518bc7..9afd6ba 100644 --- a/datasketches/src/thetafamily/theta/intersection.rs +++ b/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.state.estimated_size() + } + /// Returns the intersection result as a compact theta sketch. /// /// # Panics diff --git a/datasketches/src/thetafamily/theta/union.rs b/datasketches/src/thetafamily/theta/union.rs index d521612..3d58877 100644 --- a/datasketches/src/thetafamily/theta/union.rs +++ b/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.state.estimated_size() + } } /// Builder for [`ThetaUnion`]. diff --git a/datasketches/src/thetafamily/tuple/intersection.rs b/datasketches/src/thetafamily/tuple/intersection.rs index 58524d4..5dfaaf7 100644 --- a/datasketches/src/thetafamily/tuple/intersection.rs +++ b/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.state.estimated_size() + } + /// Returns the intersection result as a compact Tuple sketch. /// /// If `ordered` is true, retained entries are sorted ascending by hash. diff --git a/datasketches/src/thetafamily/tuple/union.rs b/datasketches/src/thetafamily/tuple/union.rs index c64bbcd..4568ed8 100644 --- a/datasketches/src/thetafamily/tuple/union.rs +++ b/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.state.estimated_size() + } } /// Builder for [`TupleUnion`]. diff --git a/datasketches/tests/cpc_test/union.rs b/datasketches/tests/cpc_test/union.rs index aaa5fa3..1d83da9 100644 --- a/datasketches/tests/cpc_test/union.rs +++ b/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); +} diff --git a/datasketches/tests/hll_test/union.rs b/datasketches/tests/hll_test/union.rs index 6d6d5d3..ef67b39 100644 --- a/datasketches/tests/hll_test/union.rs +++ b/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); +} diff --git a/datasketches/tests/theta_test/intersection.rs b/datasketches/tests/theta_test/intersection.rs index 21d32c5..65f9d70 100644 --- a/datasketches/tests/theta_test/intersection.rs +++ b/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); +} diff --git a/datasketches/tests/theta_test/union.rs b/datasketches/tests/theta_test/union.rs index 56fc64d..2a6f1cc 100644 --- a/datasketches/tests/theta_test/union.rs +++ b/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); +} diff --git a/datasketches/tests/tuple_test/intersection.rs b/datasketches/tests/tuple_test/intersection.rs index e339055..ea6e5af 100644 --- a/datasketches/tests/tuple_test/intersection.rs +++ b/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); +} diff --git a/datasketches/tests/tuple_test/union.rs b/datasketches/tests/tuple_test/union.rs index bb4af8b..23e75aa 100644 --- a/datasketches/tests/tuple_test/union.rs +++ b/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); +}