From 52ce670869d91fc55251771b6e4790ece3d1ae73 Mon Sep 17 00:00:00 2001 From: zenotme Date: Tue, 7 Jul 2026 01:25:28 +0800 Subject: [PATCH 1/3] feat(theta): add theta union --- datasketches/src/theta/mod.rs | 3 + datasketches/src/theta/union.rs | 233 +++++++++ datasketches/tests/theta_union_test.rs | 684 +++++++++++++++++++++++++ 3 files changed, 920 insertions(+) create mode 100644 datasketches/src/theta/union.rs create mode 100644 datasketches/tests/theta_union_test.rs diff --git a/datasketches/src/theta/mod.rs b/datasketches/src/theta/mod.rs index 03b5e2a..d1deca8 100644 --- a/datasketches/src/theta/mod.rs +++ b/datasketches/src/theta/mod.rs @@ -44,12 +44,15 @@ mod hash_table; mod intersection; mod serialization; mod sketch; +mod union; 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; +pub use self::union::ThetaUnion; +pub use self::union::ThetaUnionBuilder; /// Maximum theta value (signed max for compatibility with Java) const MAX_THETA: u64 = i64::MAX as u64; diff --git a/datasketches/src/theta/union.rs b/datasketches/src/theta/union.rs new file mode 100644 index 0000000..e7c5278 --- /dev/null +++ b/datasketches/src/theta/union.rs @@ -0,0 +1,233 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::common::ResizeFactor; +use crate::error::Error; +use crate::hash::DEFAULT_UPDATE_SEED; +use crate::theta::CompactThetaSketch; +use crate::theta::DEFAULT_LG_K; +use crate::theta::MAX_LG_K; +use crate::theta::MAX_THETA; +use crate::theta::MIN_LG_K; +use crate::theta::ThetaSketchView; +use crate::theta::hash_table::ThetaHashTable; + +/// Stateful union operator for Theta sketches. +#[derive(Debug)] +pub struct ThetaUnion { + table: ThetaHashTable, + union_theta: u64, +} + +impl ThetaUnion { + /// Create a new builder for ThetaUnion + /// + /// # Examples + /// + /// ``` + /// # use datasketches::theta::ThetaUnion; + /// let _union = ThetaUnion::builder().lg_k(12).build(); + /// ``` + pub fn builder() -> ThetaUnionBuilder { + ThetaUnionBuilder::default() + } + + /// Update this union with a given sketch. + pub fn update(&mut self, sketch: &S) -> Result<(), Error> { + if sketch.is_empty() { + return Ok(()); + } + + if self.table.seed_hash() != sketch.seed_hash() { + return Err(Error::invalid_argument(format!( + "incompatible seed hash: expected {}, got {}", + self.table.seed_hash(), + sketch.seed_hash(), + ))); + } + + self.table.set_empty(false); + self.union_theta = self.union_theta.min(sketch.theta64()); + + for hash in sketch.iter() { + if hash < self.union_theta && hash < self.table.theta() { + self.table.try_insert_hash(hash); + } else if sketch.is_ordered() { + break; + } + } + self.union_theta = self.union_theta.min(self.table.theta()); + + Ok(()) + } + + /// Return this union in compact form. + pub fn result(&self) -> CompactThetaSketch { + self.result_with_ordered(true) + } + + /// Return this union in compact form. + /// + /// If `ordered` is true, retained hash values are sorted in ascending order. + pub fn result_with_ordered(&self, ordered: bool) -> CompactThetaSketch { + let empty = self.table.is_empty(); + if empty { + return CompactThetaSketch::from_parts( + Vec::new(), + self.union_theta, + self.table.seed_hash(), + true, + true, + ); + } + + let mut theta = self.union_theta.min(self.table.theta()); + let mut entries = if self.union_theta >= self.table.theta() { + self.table.iter().collect::>() + } else { + self.table + .iter() + .filter(|&hash| hash < theta) + .collect::>() + }; + + let nominal_num = 1usize << self.table.lg_nom_size(); + if entries.len() > nominal_num { + let (_, kth, _) = entries.select_nth_unstable(nominal_num); + theta = *kth; + entries.truncate(nominal_num); + } + + let ordered = ordered || (entries.len() == 1 && theta == MAX_THETA); + + if ordered && entries.len() > 1 { + entries.sort_unstable(); + } + + CompactThetaSketch::from_parts(entries, theta, self.table.seed_hash(), ordered, false) + } + + /// Reset the union to empty state. + pub fn reset(&mut self) { + self.table.reset(); + self.union_theta = self.table.theta(); + } +} + +/// Builder for [`ThetaUnion`]. +#[derive(Debug, Clone)] +pub struct ThetaUnionBuilder { + lg_k: u8, + resize_factor: ResizeFactor, + sampling_probability: f32, + seed: u64, +} + +impl Default for ThetaUnionBuilder { + fn default() -> Self { + Self { + lg_k: DEFAULT_LG_K, + resize_factor: ResizeFactor::X8, + sampling_probability: 1.0, + seed: DEFAULT_UPDATE_SEED, + } + } +} + +impl ThetaUnionBuilder { + /// Set lg_k (log2 of nominal size k). + /// + /// # Panics + /// + /// If lg_k is not in range [5, 26] + /// + /// # Examples + /// + /// ``` + /// # use datasketches::theta::ThetaUnion; + /// let _union = ThetaUnion::builder().lg_k(12).build(); + /// ``` + pub fn lg_k(mut self, lg_k: u8) -> Self { + assert!( + (MIN_LG_K..=MAX_LG_K).contains(&lg_k), + "lg_k must be in [{MIN_LG_K}, {MAX_LG_K}], got {lg_k}" + ); + self.lg_k = lg_k; + self + } + + /// Set resize factor. + pub fn resize_factor(mut self, resize_factor: ResizeFactor) -> Self { + self.resize_factor = resize_factor; + self + } + + /// Set sampling probability p. + /// + /// # Panics + /// + /// Panics if p is not in range `(0.0, 1.0]` + /// + /// # Examples + /// + /// ``` + /// # use datasketches::theta::ThetaUnion; + /// let _union = ThetaUnion::builder().sampling_probability(0.5).build(); + /// ``` + pub fn sampling_probability(mut self, p: f32) -> Self { + assert!( + (0.0..=1.0).contains(&p) && p > 0.0, + "sampling_probability must be in (0.0, 1.0], got {p}" + ); + self.sampling_probability = p; + self + } + + /// Set hash seed. + /// + /// # Examples + /// + /// ``` + /// # use datasketches::theta::ThetaUnion; + /// let _union = ThetaUnion::builder().seed(7).build(); + /// ``` + pub fn seed(mut self, seed: u64) -> Self { + self.seed = seed; + self + } + + /// Build the ThetaUnion. + /// + /// # Examples + /// + /// ``` + /// # use datasketches::theta::ThetaUnion; + /// let _union = ThetaUnion::builder().lg_k(10).build(); + /// ``` + pub fn build(self) -> ThetaUnion { + let table = ThetaHashTable::new( + self.lg_k, + self.resize_factor, + self.sampling_probability, + self.seed, + ); + ThetaUnion { + union_theta: table.theta(), + table, + } + } +} diff --git a/datasketches/tests/theta_union_test.rs b/datasketches/tests/theta_union_test.rs new file mode 100644 index 0000000..5eeacec --- /dev/null +++ b/datasketches/tests/theta_union_test.rs @@ -0,0 +1,684 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#![cfg(feature = "theta")] + +use datasketches::theta::CompactThetaSketch; +use datasketches::theta::ThetaSketch; +use datasketches::theta::ThetaUnion; + +fn sketch_with_range(lg_k: u8, start: i64, count: i64) -> ThetaSketch { + let mut sketch = ThetaSketch::builder().lg_k(lg_k).build(); + for value in start..start + count { + sketch.update(value); + } + sketch +} + +fn assert_estimate_close(sketch: &CompactThetaSketch, expected: f64, tolerance: f64) { + assert!( + (sketch.estimate() - expected).abs() <= tolerance, + "estimate={}, expected={}, tolerance={}, theta={}, retained={}", + sketch.estimate(), + expected, + tolerance, + sketch.theta(), + sketch.num_retained() + ); +} + +#[test] +fn test_empty_union() { + let sketch = ThetaSketch::builder().build(); + let mut union = ThetaUnion::builder().build(); + let result = union.result(); + assert_eq!(result.num_retained(), 0); + assert!(result.is_empty()); + assert!(!result.is_estimation_mode()); + + union.update(&sketch).unwrap(); + let result = union.result(); + assert_eq!(result.num_retained(), 0); + assert!(result.is_empty()); + assert!(!result.is_estimation_mode()); +} + +#[test] +fn test_non_empty_no_retained_keys() { + let mut sketch = ThetaSketch::builder().sampling_probability(0.001).build(); + sketch.update(1u64); + + let mut union = ThetaUnion::builder().build(); + union.update(&sketch).unwrap(); + let result = union.result(); + assert_eq!(result.num_retained(), 0); + assert!(!result.is_empty()); + assert!(result.is_estimation_mode()); + assert!((result.theta() - 0.001).abs() < 1e-10); +} + +#[test] +fn test_exact_mode_half_overlap() { + let mut sketch1 = ThetaSketch::builder().build(); + for value in 0i64..1000i64 { + sketch1.update(value); + } + + let mut sketch2 = ThetaSketch::builder().build(); + for value in 500i64..1500i64 { + sketch2.update(value); + } + + let mut union = ThetaUnion::builder().build(); + union.update(&sketch1).unwrap(); + union.update(&sketch2).unwrap(); + let result = union.result(); + assert!(!result.is_empty()); + assert!(!result.is_estimation_mode()); + assert_eq!(result.estimate(), 1500.0); + + union.reset(); + let reset = union.result(); + assert_eq!(reset.num_retained(), 0); + assert!(reset.is_empty()); + assert!(!reset.is_estimation_mode()); +} + +#[test] +fn test_exact_mode_half_overlap_compact() { + let mut sketch1 = ThetaSketch::builder().build(); + for value in 0i64..1000i64 { + sketch1.update(value); + } + let compact1 = CompactThetaSketch::deserialize(&sketch1.compact(true).serialize()).unwrap(); + + let mut sketch2 = ThetaSketch::builder().build(); + for value in 500i64..1500i64 { + sketch2.update(value); + } + let compact2 = CompactThetaSketch::deserialize(&sketch2.compact(true).serialize()).unwrap(); + + let mut union = ThetaUnion::builder().build(); + union.update(&compact1).unwrap(); + union.update(&compact2).unwrap(); + let result = union.result(); + assert!(!result.is_empty()); + assert!(!result.is_estimation_mode()); + assert_eq!(result.estimate(), 1500.0); +} + +#[test] +fn test_estimation_mode_half_overlap() { + let mut sketch1 = ThetaSketch::builder().build(); + for value in 0i64..10000i64 { + sketch1.update(value); + } + + let mut sketch2 = ThetaSketch::builder().build(); + for value in 5000i64..15000i64 { + sketch2.update(value); + } + + let mut union = ThetaUnion::builder().build(); + union.update(&sketch1).unwrap(); + union.update(&sketch2).unwrap(); + let result = union.result(); + assert!(!result.is_empty()); + assert!(result.is_estimation_mode()); + assert!( + (result.estimate() - 15000.0).abs() <= 15000.0 * 0.01, + "estimate={}, theta={}, retained={}", + result.estimate(), + result.theta(), + result.num_retained() + ); +} + +#[test] +fn test_seed_mismatch() { + let mut sketch = ThetaSketch::builder().build(); + sketch.update(1u64); + + let mut union = ThetaUnion::builder().seed(123).build(); + assert!(union.update(&sketch).is_err()); +} + +#[test] +fn test_larger_k() { + let mut sketch1 = ThetaSketch::builder().lg_k(14).build(); + for value in 0i64..16384i64 { + sketch1.update(value); + } + + let mut sketch2 = ThetaSketch::builder().lg_k(14).build(); + for value in 0i64..26384i64 { + sketch2.update(value); + } + + let mut sketch3 = ThetaSketch::builder().lg_k(14).build(); + for value in 0i64..86384i64 { + sketch3.update(value); + } + + let mut union1 = ThetaUnion::builder().lg_k(16).build(); + union1.update(&sketch2).unwrap(); + union1.update(&sketch1).unwrap(); + union1.update(&sketch3).unwrap(); + let result1 = union1.result(); + assert_eq!(result1.estimate(), sketch3.estimate()); + + let mut union2 = ThetaUnion::builder().lg_k(16).build(); + union2.update(&sketch1).unwrap(); + union2.update(&sketch3).unwrap(); + union2.update(&sketch2).unwrap(); + let result2 = union2.result(); + assert_eq!(result2.estimate(), sketch3.estimate()); +} + +#[test] +fn test_exact_union_no_overlap() { + let lg_k = 9; + let k = 1i64 << lg_k; + let sketch1 = sketch_with_range(lg_k, 0, k / 2); + let sketch2 = sketch_with_range(lg_k, k / 2, k / 2); + + let mut union = ThetaUnion::builder().lg_k(lg_k).build(); + union.update(&sketch1).unwrap(); + union.update(&sketch2).unwrap(); + + let result = union.result(); + assert!(!result.is_empty()); + assert!(!result.is_estimation_mode()); + assert_eq!(result.estimate(), k as f64); +} + +#[test] +fn test_estimation_union_no_overlap() { + let lg_k = 12; + let k = 1i64 << lg_k; + let sketch1 = sketch_with_range(lg_k, 0, 2 * k); + let sketch2 = sketch_with_range(lg_k, 2 * k, 2 * k); + + let mut union = ThetaUnion::builder().lg_k(lg_k).build(); + union.update(&sketch1).unwrap(); + union.update(&sketch2).unwrap(); + + assert_estimate_close(&union.result(), (4 * k) as f64, 0.05 * (4 * k) as f64); +} + +#[test] +fn test_exact_union_with_overlap() { + let lg_k = 9; + let k = 1i64 << lg_k; + let sketch1 = sketch_with_range(lg_k, 0, k / 2); + let sketch2 = sketch_with_range(lg_k, 0, k); + + let mut union = ThetaUnion::builder().lg_k(lg_k).build(); + union.update(&sketch1).unwrap(); + union.update(&sketch2).unwrap(); + + let result = union.result(); + assert!(!result.is_empty()); + assert!(!result.is_estimation_mode()); + assert_eq!(result.estimate(), k as f64); +} + +#[test] +fn test_ordered_and_unordered_compact_inputs() { + let lg_k = 12; + let k = 1i64 << lg_k; + let sketch1 = sketch_with_range(lg_k, 0, 2 * k); + let sketch2 = sketch_with_range(lg_k + 1, 2 * k, 2 * k); + let compact_ordered = sketch2.compact(true); + let compact_unordered = sketch2.compact(false); + + let mut ordered_union = ThetaUnion::builder().lg_k(lg_k).build(); + ordered_union.update(&sketch1).unwrap(); + ordered_union.update(&compact_ordered).unwrap(); + + let mut unordered_union = ThetaUnion::builder().lg_k(lg_k).build(); + unordered_union.update(&sketch1).unwrap(); + unordered_union.update(&compact_unordered).unwrap(); + + assert_eq!( + ordered_union.result().estimate(), + unordered_union.result().estimate() + ); + assert_estimate_close( + &ordered_union.result(), + (4 * k) as f64, + 0.05 * (4 * k) as f64, + ); +} + +#[test] +fn test_result_ordering_forms_have_same_estimate() { + let sketch1 = sketch_with_range(12, 0, 8192); + let sketch2 = sketch_with_range(12, 8192, 1024); + + let mut union = ThetaUnion::builder().lg_k(12).build(); + union.update(&sketch1).unwrap(); + union.update(&sketch2).unwrap(); + + let unordered = union.result_with_ordered(false); + let ordered = union.result_with_ordered(true); + + assert!(!unordered.is_ordered()); + assert!(ordered.is_ordered()); + assert_eq!(unordered.estimate(), ordered.estimate()); +} + +#[test] +fn test_multi_union() { + let lg_k = 13; + let ranges = [ + (0, 100_000), + (100_000, 26_797), + (126_797, 26_797), + (153_594, 26_797), + ]; + let mut union = ThetaUnion::builder().lg_k(lg_k).build(); + + for (start, count) in ranges { + let sketch = sketch_with_range(lg_k, start, count); + union.update(&sketch).unwrap(); + } + + assert_estimate_close(&union.result(), 180_391.0, 180_391.0 * 0.02); +} + +#[test] +fn test_result_does_not_reset_union() { + let lg_k = 9; + let k = 1i64 << lg_k; + let compact1 = sketch_with_range(lg_k, 0, k).compact(true); + let compact2 = sketch_with_range(lg_k, k, k).compact(true); + + let mut union = ThetaUnion::builder().lg_k(lg_k).build(); + union.update(&compact1).unwrap(); + union.update(&compact2).unwrap(); + let first = union.result(); + let second = union.result(); + + assert_eq!(first.estimate(), second.estimate()); + assert!(!second.is_empty()); +} + +#[test] +fn test_union_full_overlap() { + let lg_k = 9; + let k = 1i64 << lg_k; + let compact1 = sketch_with_range(lg_k, 0, k).compact(true); + let compact2 = sketch_with_range(lg_k, 0, k).compact(true); + + let mut union = ThetaUnion::builder().lg_k(lg_k).build(); + union.update(&compact1).unwrap(); + union.update(&compact2).unwrap(); + let result = union.result(); + + assert_eq!(result.estimate(), k as f64); +} + +#[test] +fn test_ordered_input_early_stop_matches_unordered_input() { + let lg_k = 10; + let k = 1i64 << lg_k; + let mut value = 0i64; + + for _ in 0..10 { + let sketch1 = sketch_with_range(lg_k, value, 4 * k); + let sketch2 = sketch_with_range(lg_k, value + 2 * k, 4 * k); + value += 6 * k; + let ordered1 = sketch1.compact(true); + let ordered2 = sketch2.compact(true); + let unordered1 = sketch1.compact(false); + let unordered2 = sketch2.compact(false); + + let mut ordered_union = ThetaUnion::builder().lg_k(lg_k + 1).build(); + ordered_union.update(&ordered1).unwrap(); + ordered_union.update(&ordered2).unwrap(); + + let mut unordered_union = ThetaUnion::builder().lg_k(lg_k + 1).build(); + unordered_union.update(&unordered1).unwrap(); + unordered_union.update(&unordered2).unwrap(); + + assert_eq!( + ordered_union.result().estimate(), + unordered_union.result().estimate() + ); + } +} + +#[test] +fn test_union_cutback_to_k() { + let lg_k = 10; + let k = 1i64 << lg_k; + let compact1 = sketch_with_range(lg_k, 0, 3 * k).compact(true); + let compact2 = sketch_with_range(lg_k, 6 * k, 3 * k).compact(true); + + let mut union = ThetaUnion::builder().lg_k(lg_k).build(); + union.update(&compact1).unwrap(); + union.update(&compact2).unwrap(); + let result = union.result(); + + assert_estimate_close(&result, (6 * k) as f64, (6 * k) as f64 * 0.06); + assert!(result.num_retained() <= k as usize); +} + +#[test] +fn test_union_empty_valid_rules() { + let empty1 = ThetaSketch::builder().build().compact(true); + let empty2 = ThetaSketch::builder().build().compact(true); + let mut one = ThetaSketch::builder().build(); + one.update(1i64); + let one = one.compact(true); + + let mut empty_union = ThetaUnion::builder().lg_k(5).build(); + empty_union.update(&empty1).unwrap(); + empty_union.update(&empty2).unwrap(); + assert!(empty_union.result().is_empty()); + + let mut left_non_empty_union = ThetaUnion::builder().lg_k(5).build(); + left_non_empty_union.update(&one).unwrap(); + left_non_empty_union.update(&empty2).unwrap(); + assert!(!left_non_empty_union.result().is_empty()); + + let mut right_non_empty_union = ThetaUnion::builder().lg_k(5).build(); + right_non_empty_union.update(&empty1).unwrap(); + right_non_empty_union.update(&one).unwrap(); + assert!(!right_non_empty_union.result().is_empty()); +} + +#[test] +fn test_trim_to_k() { + let hi_sketch = sketch_with_range(10, 0, 3749); + let lo_sketch = sketch_with_range(9, 10_000, 1783); + + let mut union = ThetaUnion::builder().lg_k(10).build(); + union.update(&hi_sketch).unwrap(); + union.update(&lo_sketch).unwrap(); + let result = union.result(); + + assert_eq!(result.num_retained(), 1024); +} + +#[test] +fn test_builder_lg_k() { + let sketch = sketch_with_range(10, 0, 1000); + let mut union = ThetaUnion::builder().lg_k(10).build(); + union.update(&sketch).unwrap(); + + assert_eq!(union.result().estimate(), 1000.0); +} + +#[derive(Clone, Copy, Debug)] +enum CornerSketchState { + Empty, + Exact, + Estimation, + Degenerate, +} + +fn corner_sketch(state: CornerSketchState, p: f32, value: i64) -> ThetaSketch { + let builder = ThetaSketch::builder().lg_k(5); + let mut sketch = match state { + CornerSketchState::Empty | CornerSketchState::Exact => builder.build(), + CornerSketchState::Estimation | CornerSketchState::Degenerate => { + builder.sampling_probability(p).build() + } + }; + if !matches!(state, CornerSketchState::Empty) { + sketch.update(value); + } + sketch +} + +#[test] +fn test_corner_case_union_states() { + const GT_MIDP_VALUE: i64 = 3; + const MIDP: f32 = 0.5; + const GT_LOWP_VALUE: i64 = 6; + const LOWP: f32 = 0.1; + const LT_LOWP_VALUE: i64 = 4; + + let cases = [ + ( + CornerSketchState::Empty, + 1.0, + 0, + CornerSketchState::Empty, + 1.0, + 0, + 1.0, + 0, + true, + ), + ( + CornerSketchState::Empty, + 1.0, + 0, + CornerSketchState::Exact, + 1.0, + GT_MIDP_VALUE, + 1.0, + 1, + false, + ), + ( + CornerSketchState::Empty, + 1.0, + 0, + CornerSketchState::Degenerate, + LOWP, + GT_LOWP_VALUE, + LOWP as f64, + 0, + false, + ), + ( + CornerSketchState::Empty, + 1.0, + 0, + CornerSketchState::Estimation, + LOWP, + LT_LOWP_VALUE, + LOWP as f64, + 1, + false, + ), + ( + CornerSketchState::Exact, + 1.0, + GT_MIDP_VALUE, + CornerSketchState::Empty, + 1.0, + 0, + 1.0, + 1, + false, + ), + ( + CornerSketchState::Exact, + 1.0, + GT_MIDP_VALUE, + CornerSketchState::Exact, + 1.0, + GT_MIDP_VALUE, + 1.0, + 1, + false, + ), + ( + CornerSketchState::Exact, + 1.0, + LT_LOWP_VALUE, + CornerSketchState::Degenerate, + LOWP, + GT_LOWP_VALUE, + LOWP as f64, + 1, + false, + ), + ( + CornerSketchState::Exact, + 1.0, + LT_LOWP_VALUE, + CornerSketchState::Estimation, + LOWP, + LT_LOWP_VALUE, + LOWP as f64, + 1, + false, + ), + ( + CornerSketchState::Estimation, + LOWP, + LT_LOWP_VALUE, + CornerSketchState::Empty, + 1.0, + 0, + LOWP as f64, + 1, + false, + ), + ( + CornerSketchState::Estimation, + LOWP, + LT_LOWP_VALUE, + CornerSketchState::Exact, + 1.0, + LT_LOWP_VALUE, + LOWP as f64, + 1, + false, + ), + ( + CornerSketchState::Estimation, + MIDP, + LT_LOWP_VALUE, + CornerSketchState::Degenerate, + LOWP, + GT_LOWP_VALUE, + LOWP as f64, + 1, + false, + ), + ( + CornerSketchState::Estimation, + MIDP, + LT_LOWP_VALUE, + CornerSketchState::Estimation, + LOWP, + LT_LOWP_VALUE, + LOWP as f64, + 1, + false, + ), + ( + CornerSketchState::Degenerate, + LOWP, + GT_LOWP_VALUE, + CornerSketchState::Empty, + 1.0, + 0, + LOWP as f64, + 0, + false, + ), + ( + CornerSketchState::Degenerate, + LOWP, + GT_LOWP_VALUE, + CornerSketchState::Exact, + 1.0, + LT_LOWP_VALUE, + LOWP as f64, + 1, + false, + ), + ( + CornerSketchState::Degenerate, + MIDP, + GT_MIDP_VALUE, + CornerSketchState::Degenerate, + LOWP, + GT_LOWP_VALUE, + LOWP as f64, + 0, + false, + ), + ( + CornerSketchState::Degenerate, + MIDP, + GT_MIDP_VALUE, + CornerSketchState::Estimation, + LOWP, + LT_LOWP_VALUE, + LOWP as f64, + 1, + false, + ), + ]; + + for ( + state_a, + p_a, + value_a, + state_b, + p_b, + value_b, + expected_theta, + expected_count, + expected_empty, + ) in cases + { + let sketch_a = corner_sketch(state_a, p_a, value_a); + let sketch_b = corner_sketch(state_b, p_b, value_b); + + let mut union = ThetaUnion::builder().build(); + union.update(&sketch_a).unwrap(); + union.update(&sketch_b).unwrap(); + let result = union.result(); + + assert!( + (result.theta() - expected_theta).abs() < 1e-6, + "state_a={state_a:?}, state_b={state_b:?}, theta={}", + result.theta() + ); + assert_eq!( + result.num_retained(), + expected_count, + "state_a={state_a:?}, state_b={state_b:?}" + ); + assert_eq!( + result.is_empty(), + expected_empty, + "state_a={state_a:?}, state_b={state_b:?}" + ); + + let compact_a = sketch_a.compact(true); + let compact_b = sketch_b.compact(true); + let mut union = ThetaUnion::builder().build(); + union.update(&compact_a).unwrap(); + union.update(&compact_b).unwrap(); + let compact_result = union.result(); + + assert!((compact_result.theta() - expected_theta).abs() < 1e-6); + assert_eq!(compact_result.num_retained(), expected_count); + assert_eq!(compact_result.is_empty(), expected_empty); + } +} From 4c167cb2a95fef34035ad391ecdc37683351f9c6 Mon Sep 17 00:00:00 2001 From: zenotme Date: Fri, 10 Jul 2026 23:39:27 +0800 Subject: [PATCH 2/3] refactor(theta): extract shared theta primitives --- datasketches/src/lib.rs | 3 + datasketches/src/theta/hash_table.rs | 18 +- datasketches/src/theta/intersection.rs | 6 +- datasketches/src/theta/mod.rs | 25 +- datasketches/src/theta/sketch.rs | 102 +++---- datasketches/src/theta/union.rs | 106 +++----- .../hash_table.rs} | 14 +- datasketches/src/theta_common/mod.rs | 41 +++ datasketches/src/theta_common/sketch_view.rs | 47 ++++ datasketches/src/theta_common/union.rs | 251 ++++++++++++++++++ 10 files changed, 447 insertions(+), 166 deletions(-) rename datasketches/src/{theta/raw_hash_table.rs => theta_common/hash_table.rs} (97%) create mode 100644 datasketches/src/theta_common/mod.rs create mode 100644 datasketches/src/theta_common/sketch_view.rs create mode 100644 datasketches/src/theta_common/union.rs diff --git a/datasketches/src/lib.rs b/datasketches/src/lib.rs index 65f0ddd..083f064 100644 --- a/datasketches/src/lib.rs +++ b/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; diff --git a/datasketches/src/theta/hash_table.rs b/datasketches/src/theta/hash_table.rs index 4fbd5e6..ffac00b 100644 --- a/datasketches/src/theta/hash_table.rs +++ b/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; +/// 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 { diff --git a/datasketches/src/theta/intersection.rs b/datasketches/src/theta/intersection.rs index 85d2eb7..cbc68bf 100644 --- a/datasketches/src/theta/intersection.rs +++ b/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 { diff --git a/datasketches/src/theta/mod.rs b/datasketches/src/theta/mod.rs index 7c25440..1df43fb 100644 --- a/datasketches/src/theta/mod.rs +++ b/datasketches/src/theta/mod.rs @@ -42,11 +42,17 @@ mod bit_pack; mod hash_table; mod intersection; -mod raw_hash_table; mod serialization; mod sketch; mod union; +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; + +pub use self::hash_table::ThetaEntry; pub use self::intersection::ThetaIntersection; pub use self::sketch::CompactThetaSketch; pub use self::sketch::ThetaSketch; @@ -54,20 +60,3 @@ pub use self::sketch::ThetaSketchBuilder; pub use self::sketch::ThetaSketchView; pub use self::union::ThetaUnion; pub use self::union::ThetaUnionBuilder; - -/// 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; diff --git a/datasketches/src/theta/sketch.rs b/datasketches/src/theta/sketch.rs index 6ee9a23..a998f14 100644 --- a/datasketches/src/theta/sketch.rs +++ b/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 {} - /// Returns theta as `u64`. - fn theta64(&self) -> u64; +impl> 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 + '_; +impl RawThetaSketchView 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 + '_ { + 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 + '_ { - self.table.iter() + pub fn iter(&self) -> impl Iterator + '_ { + 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 = self.iter().collect(); + let mut entries: Vec = 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 + '_ { - 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 + '_ { - self.entries.iter().copied() + /// Return iterator over retained entries. + pub fn iter(&self) -> impl Iterator + '_ { + 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 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 + '_ { - CompactThetaSketch::iter(self) + fn is_ordered(&self) -> bool { + CompactThetaSketch::is_ordered(self) + } + + fn iter(&self) -> impl Iterator + '_ { + 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 { - let mut entries: Vec = sketch.iter().collect(); + let mut entries: Vec = sketch.iter().map(|entry| entry.hash()).collect(); entries.sort_unstable(); entries } fn sorted_compact_entries(sketch: &CompactThetaSketch) -> Vec { - let mut entries: Vec = sketch.iter().collect(); + let mut entries: Vec = sketch.iter().map(|entry| entry.hash()).collect(); entries.sort_unstable(); entries } diff --git a/datasketches/src/theta/union.rs b/datasketches/src/theta/union.rs index e7c5278..cf73ade 100644 --- a/datasketches/src/theta/union.rs +++ b/datasketches/src/theta/union.rs @@ -21,16 +21,23 @@ use crate::hash::DEFAULT_UPDATE_SEED; use crate::theta::CompactThetaSketch; use crate::theta::DEFAULT_LG_K; use crate::theta::MAX_LG_K; -use crate::theta::MAX_THETA; use crate::theta::MIN_LG_K; use crate::theta::ThetaSketchView; -use crate::theta::hash_table::ThetaHashTable; +use crate::theta::hash_table::ThetaEntry; +use crate::theta_common::union::RawThetaUnion; +use crate::theta_common::union::RawThetaUnionPolicy; /// Stateful union operator for Theta sketches. #[derive(Debug)] pub struct ThetaUnion { - table: ThetaHashTable, - union_theta: u64, + raw: RawThetaUnion, +} + +#[derive(Debug)] +struct NoopUnionPolicy; + +impl RawThetaUnionPolicy for NoopUnionPolicy { + fn merge(&self, _existing: &mut ThetaEntry, _incoming: ThetaEntry) {} } impl ThetaUnion { @@ -48,31 +55,7 @@ impl ThetaUnion { /// Update this union with a given sketch. pub fn update(&mut self, sketch: &S) -> Result<(), Error> { - if sketch.is_empty() { - return Ok(()); - } - - if self.table.seed_hash() != sketch.seed_hash() { - return Err(Error::invalid_argument(format!( - "incompatible seed hash: expected {}, got {}", - self.table.seed_hash(), - sketch.seed_hash(), - ))); - } - - self.table.set_empty(false); - self.union_theta = self.union_theta.min(sketch.theta64()); - - for hash in sketch.iter() { - if hash < self.union_theta && hash < self.table.theta() { - self.table.try_insert_hash(hash); - } else if sketch.is_ordered() { - break; - } - } - self.union_theta = self.union_theta.min(self.table.theta()); - - Ok(()) + self.raw.update(sketch) } /// Return this union in compact form. @@ -84,47 +67,23 @@ impl ThetaUnion { /// /// If `ordered` is true, retained hash values are sorted in ascending order. pub fn result_with_ordered(&self, ordered: bool) -> CompactThetaSketch { - let empty = self.table.is_empty(); - if empty { - return CompactThetaSketch::from_parts( - Vec::new(), - self.union_theta, - self.table.seed_hash(), - true, - true, - ); - } - - let mut theta = self.union_theta.min(self.table.theta()); - let mut entries = if self.union_theta >= self.table.theta() { - self.table.iter().collect::>() - } else { - self.table - .iter() - .filter(|&hash| hash < theta) - .collect::>() - }; - - let nominal_num = 1usize << self.table.lg_nom_size(); - if entries.len() > nominal_num { - let (_, kth, _) = entries.select_nth_unstable(nominal_num); - theta = *kth; - entries.truncate(nominal_num); - } - - let ordered = ordered || (entries.len() == 1 && theta == MAX_THETA); - - if ordered && entries.len() > 1 { - entries.sort_unstable(); - } - - CompactThetaSketch::from_parts(entries, theta, self.table.seed_hash(), ordered, false) + let result = self.raw.result(ordered); + CompactThetaSketch::from_parts( + result + .entries + .into_iter() + .map(|entry| entry.hash()) + .collect(), + result.theta, + result.seed_hash, + result.ordered, + result.empty, + ) } /// Reset the union to empty state. pub fn reset(&mut self) { - self.table.reset(); - self.union_theta = self.table.theta(); + self.raw.reset(); } } @@ -219,15 +178,14 @@ impl ThetaUnionBuilder { /// let _union = ThetaUnion::builder().lg_k(10).build(); /// ``` pub fn build(self) -> ThetaUnion { - let table = ThetaHashTable::new( - self.lg_k, - self.resize_factor, - self.sampling_probability, - self.seed, - ); ThetaUnion { - union_theta: table.theta(), - table, + raw: RawThetaUnion::new( + self.lg_k, + self.resize_factor, + self.sampling_probability, + self.seed, + NoopUnionPolicy, + ), } } } diff --git a/datasketches/src/theta/raw_hash_table.rs b/datasketches/src/theta_common/hash_table.rs similarity index 97% rename from datasketches/src/theta/raw_hash_table.rs rename to datasketches/src/theta_common/hash_table.rs index f171a36..b7b8fe5 100644 --- a/datasketches/src/theta/raw_hash_table.rs +++ b/datasketches/src/theta_common/hash_table.rs @@ -17,16 +17,18 @@ use std::hash::Hash; +use super::HASH_TABLE_REBUILD_THRESHOLD; +use super::HASH_TABLE_RESIZE_THRESHOLD; +use super::MAX_THETA; +use super::MIN_LG_K; +use super::STRIDE_MASK; use crate::common::ResizeFactor; use crate::hash::MurmurHash3X64128; use crate::hash::compute_seed_hash; -use crate::theta::HASH_TABLE_REBUILD_THRESHOLD; -use crate::theta::HASH_TABLE_RESIZE_THRESHOLD; -use crate::theta::MAX_THETA; -use crate::theta::MIN_LG_K; -use crate::theta::STRIDE_MASK; -pub(crate) trait RawHashTableEntry { +/// An entry retained by a Theta-family hash table. +pub trait RawHashTableEntry { + /// Return the hash used as this entry's key. fn hash(&self) -> u64; } diff --git a/datasketches/src/theta_common/mod.rs b/datasketches/src/theta_common/mod.rs new file mode 100644 index 0000000..c912694 --- /dev/null +++ b/datasketches/src/theta_common/mod.rs @@ -0,0 +1,41 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +pub(crate) mod hash_table; +pub(crate) mod sketch_view; +pub(crate) mod union; + +// Public because public view APIs expose this trait and its entry bound, e.g. `ThetaSketchView: RawThetaSketchView`. +pub use self::hash_table::RawHashTableEntry; +pub use self::sketch_view::RawThetaSketchView; + +/// 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; diff --git a/datasketches/src/theta_common/sketch_view.rs b/datasketches/src/theta_common/sketch_view.rs new file mode 100644 index 0000000..69ffd50 --- /dev/null +++ b/datasketches/src/theta_common/sketch_view.rs @@ -0,0 +1,47 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::theta_common::hash_table::RawHashTableEntry; + +pub(crate) mod private { + pub(crate) trait Sealed {} +} + +/// Read-only input accepted by a raw Theta union. +/// +/// This trait carries complete retained entries, so tuple unions can use the same state machine +/// while merging their per-key summaries. +#[allow(private_bounds)] +pub trait RawThetaSketchView: private::Sealed { + /// Return the 16-bit seed hash. + fn seed_hash(&self) -> u16; + + /// Return theta as a `u64` threshold. + fn theta64(&self) -> u64; + + /// Return whether this sketch has not received any updates. + fn is_empty(&self) -> bool; + + /// Return whether retained entries are ordered by ascending hash. + fn is_ordered(&self) -> bool; + + /// Return an iterator over retained entries. + fn iter(&self) -> impl Iterator + '_; + + /// Return the number of retained entries. + fn num_retained(&self) -> usize; +} diff --git a/datasketches/src/theta_common/union.rs b/datasketches/src/theta_common/union.rs new file mode 100644 index 0000000..25d51fd --- /dev/null +++ b/datasketches/src/theta_common/union.rs @@ -0,0 +1,251 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::common::ResizeFactor; +use crate::error::Error; +use crate::theta_common::MAX_THETA; +use crate::theta_common::hash_table::RawHashTable; +use crate::theta_common::hash_table::RawHashTableEntry; +use crate::theta_common::sketch_view::RawThetaSketchView; + +/// Merges an incoming entry into an existing entry with the same hash. +pub(crate) trait RawThetaUnionPolicy { + fn merge(&self, existing: &mut E, incoming: E); +} + +/// Generic state machine shared by Theta and Tuple unions. +/// +/// `E` is the retained entry type. Ordinary Theta entries only contain a hash, while tuple +/// entries also carry a summary. `P` defines how equal-hash entries are combined. +#[derive(Debug)] +pub(crate) struct RawThetaUnion { + table: RawHashTable, + policy: P, + union_theta: u64, +} + +/// Raw compact-union state from which a sketch family creates its compact result type. +#[derive(Debug)] +pub(crate) struct RawThetaUnionResult { + pub(crate) entries: Vec, + pub(crate) theta: u64, + pub(crate) seed_hash: u16, + pub(crate) ordered: bool, + pub(crate) empty: bool, +} + +impl RawThetaUnion +where + E: RawHashTableEntry, +{ + pub(crate) fn new( + lg_k: u8, + resize_factor: ResizeFactor, + sampling_probability: f32, + seed: u64, + policy: P, + ) -> Self { + let table = RawHashTable::new(lg_k, resize_factor, sampling_probability, seed); + Self { + union_theta: table.theta(), + table, + policy, + } + } + + /// Incorporate a sketch into the union. + pub(crate) fn update(&mut self, sketch: &S) -> Result<(), Error> + where + S: RawThetaSketchView, + P: RawThetaUnionPolicy, + { + if sketch.is_empty() { + return Ok(()); + } + + if self.table.seed_hash() != sketch.seed_hash() { + return Err(Error::invalid_argument(format!( + "incompatible seed hash: expected {}, got {}", + self.table.seed_hash(), + sketch.seed_hash(), + ))); + } + + self.table.set_empty(false); + self.union_theta = self.union_theta.min(sketch.theta64()); + + for entry in sketch.iter() { + let hash = entry.hash(); + if hash < self.union_theta && hash < self.table.theta() { + self.table.upsert_entry(hash, |existing| match existing { + Some(existing) => { + self.policy.merge(existing, entry); + None + } + None => Some(entry), + }); + } else if sketch.is_ordered() { + break; + } + } + self.union_theta = self.union_theta.min(self.table.theta()); + + Ok(()) + } + + /// Return the current compact-union state. + pub(crate) fn result(&self, ordered: bool) -> RawThetaUnionResult + where + E: Clone, + { + if self.table.is_empty() { + return RawThetaUnionResult { + entries: Vec::new(), + theta: self.union_theta, + seed_hash: self.table.seed_hash(), + ordered: true, + empty: true, + }; + } + + let mut theta = self.union_theta.min(self.table.theta()); + let mut entries = if self.union_theta >= self.table.theta() { + self.table.iter_entries().cloned().collect::>() + } else { + self.table + .iter_entries() + .filter(|entry| entry.hash() < theta) + .cloned() + .collect::>() + }; + + let nominal_num = 1usize << self.table.lg_nom_size(); + if entries.len() > nominal_num { + let (_, kth, _) = entries.select_nth_unstable_by_key(nominal_num, |entry| entry.hash()); + theta = kth.hash(); + entries.truncate(nominal_num); + } + + let ordered = ordered || (entries.len() == 1 && theta == MAX_THETA); + if ordered { + entries.sort_unstable_by_key(RawHashTableEntry::hash); + } + + RawThetaUnionResult { + entries, + theta, + seed_hash: self.table.seed_hash(), + ordered, + empty: false, + } + } + + /// Reset the union to its initial state. + pub(crate) fn reset(&mut self) { + self.table.reset(); + self.union_theta = self.table.theta(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::hash::DEFAULT_UPDATE_SEED; + use crate::theta_common::sketch_view::RawThetaSketchView; + + #[derive(Clone, Debug, Eq, PartialEq)] + struct TestEntry { + hash: u64, + summary: u64, + } + + impl RawHashTableEntry for TestEntry { + fn hash(&self) -> u64 { + self.hash + } + } + + struct TestSketch { + entries: Vec, + } + + impl crate::theta_common::sketch_view::private::Sealed for TestSketch {} + + impl RawThetaSketchView for TestSketch { + fn seed_hash(&self) -> u16 { + crate::hash::compute_seed_hash(DEFAULT_UPDATE_SEED) + } + + fn theta64(&self) -> u64 { + MAX_THETA + } + + fn is_empty(&self) -> bool { + false + } + + fn is_ordered(&self) -> bool { + false + } + + fn iter(&self) -> impl Iterator + '_ { + self.entries.iter().cloned() + } + + fn num_retained(&self) -> usize { + self.entries.len() + } + } + + struct SumPolicy; + + impl RawThetaUnionPolicy for SumPolicy { + fn merge(&self, existing: &mut TestEntry, incoming: TestEntry) { + existing.summary += incoming.summary; + } + } + + #[test] + fn merges_equal_hash_entries_with_policy() { + let mut union = + RawThetaUnion::new(5, ResizeFactor::X1, 1.0, DEFAULT_UPDATE_SEED, SumPolicy); + union + .update(&TestSketch { + entries: vec![TestEntry { + hash: 1, + summary: 2, + }], + }) + .unwrap(); + union + .update(&TestSketch { + entries: vec![TestEntry { + hash: 1, + summary: 3, + }], + }) + .unwrap(); + + assert_eq!( + union.result(true).entries, + vec![TestEntry { + hash: 1, + summary: 5, + }] + ); + } +} From 9e22934ea4a8728f17e5b335006ee12a413bc4c4 Mon Sep 17 00:00:00 2001 From: zenotme Date: Sat, 11 Jul 2026 01:03:00 +0800 Subject: [PATCH 3/3] refactor(theta): rename `result`, `result_with_ordered` to `to_sketch` --- datasketches/src/theta/intersection.rs | 13 +--- datasketches/src/theta/mod.rs | 11 ++-- datasketches/src/theta/union.rs | 11 +--- datasketches/src/theta_common/mod.rs | 3 +- datasketches/tests/theta_intersection_test.rs | 42 ++++++------ datasketches/tests/theta_union_test.rs | 66 ++++++++++--------- 6 files changed, 67 insertions(+), 79 deletions(-) diff --git a/datasketches/src/theta/intersection.rs b/datasketches/src/theta/intersection.rs index cbc68bf..46b928c 100644 --- a/datasketches/src/theta/intersection.rs +++ b/datasketches/src/theta/intersection.rs @@ -202,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 = self.table.iter().collect(); if ordered { diff --git a/datasketches/src/theta/mod.rs b/datasketches/src/theta/mod.rs index 1df43fb..59d7b7d 100644 --- a/datasketches/src/theta/mod.rs +++ b/datasketches/src/theta/mod.rs @@ -46,12 +46,6 @@ mod serialization; mod sketch; mod union; -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; - pub use self::hash_table::ThetaEntry; pub use self::intersection::ThetaIntersection; pub use self::sketch::CompactThetaSketch; @@ -60,3 +54,8 @@ pub use self::sketch::ThetaSketchBuilder; pub use self::sketch::ThetaSketchView; 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; diff --git a/datasketches/src/theta/union.rs b/datasketches/src/theta/union.rs index cf73ade..891d06e 100644 --- a/datasketches/src/theta/union.rs +++ b/datasketches/src/theta/union.rs @@ -58,15 +58,8 @@ impl ThetaUnion { self.raw.update(sketch) } - /// Return this union in compact form. - pub fn result(&self) -> CompactThetaSketch { - self.result_with_ordered(true) - } - - /// Return this union in compact form. - /// - /// If `ordered` is true, retained hash values are sorted in ascending order. - pub fn result_with_ordered(&self, ordered: bool) -> CompactThetaSketch { + /// Return this union as a compact sketch. + pub fn to_sketch(&self, ordered: bool) -> CompactThetaSketch { let result = self.raw.result(ordered); CompactThetaSketch::from_parts( result diff --git a/datasketches/src/theta_common/mod.rs b/datasketches/src/theta_common/mod.rs index c912694..96fcfdb 100644 --- a/datasketches/src/theta_common/mod.rs +++ b/datasketches/src/theta_common/mod.rs @@ -19,7 +19,8 @@ pub(crate) mod hash_table; pub(crate) mod sketch_view; pub(crate) mod union; -// Public because public view APIs expose this trait and its entry bound, e.g. `ThetaSketchView: RawThetaSketchView`. +// Public because public view APIs expose this trait and its entry bound, e.g. `ThetaSketchView: +// RawThetaSketchView`. pub use self::hash_table::RawHashTableEntry; pub use self::sketch_view::RawThetaSketchView; diff --git a/datasketches/tests/theta_intersection_test.rs b/datasketches/tests/theta_intersection_test.rs index 9b98062..bb662b6 100644 --- a/datasketches/tests/theta_intersection_test.rs +++ b/datasketches/tests/theta_intersection_test.rs @@ -38,14 +38,14 @@ fn test_has_result_state_machine() { assert!(!i.has_result()); i.update(&a).unwrap(); assert!(i.has_result()); - assert!(i.result().estimate() >= 1.0); + assert!(i.to_sketch(true).estimate() >= 1.0); } #[test] fn test_result_before_update_panics() { let i = ThetaIntersection::new(123); let result = std::panic::catch_unwind(|| { - let _ = i.result(); + let _ = i.to_sketch(true); }); assert!(result.is_err()); } @@ -64,7 +64,7 @@ fn test_update_accepts_compact_sketch() { i.update(&a.compact(true)).unwrap(); i.update(&b).unwrap(); - let r = i.result(); + let r = i.to_sketch(true); assert!(r.estimate() == 1.0); assert!(r.is_ordered()); @@ -75,7 +75,7 @@ fn test_update_accepts_compact_sketch() { i.update(&c.compact(false)).unwrap(); - let r = i.result_with_ordered(false); + let r = i.to_sketch(false); assert!(r.estimate() == 0.0); assert!(!r.is_ordered()); } @@ -87,7 +87,7 @@ fn test_seed_mismatch_behaviour_for_empty_sketch() { i.update(&empty_other_seed).unwrap(); assert!(i.has_result()); - let r = i.result(); + let r = i.to_sketch(true); assert!(r.is_empty()); } @@ -111,12 +111,12 @@ fn test_terminal_empty_state_ignores_future_updates() { i.update(&empty).unwrap(); i.update(&non_empty).unwrap(); - let r = i.result(); + let r = i.to_sketch(true); assert!(r.is_empty()); } #[test] -fn test_result_with_ordered_false_is_not_ordered() { +fn test_to_sketch_unordered_is_not_ordered() { let mut a = ThetaSketch::builder().build(); for i in 0..64 { a.update(i); @@ -124,7 +124,7 @@ fn test_result_with_ordered_false_is_not_ordered() { let mut i = ThetaIntersection::new_with_default_seed(); i.update(&a).unwrap(); - let r = i.result_with_ordered(false); + let r = i.to_sketch(false); assert!(!r.is_ordered()); } @@ -134,14 +134,14 @@ fn test_empty_update_twice() { let mut i = ThetaIntersection::new_with_default_seed(); i.update(&empty).unwrap(); - let r1 = i.result(); + let r1 = i.to_sketch(true); assert_eq!(r1.num_retained(), 0); assert!(r1.is_empty()); assert!(!r1.is_estimation_mode()); assert_eq!(r1.estimate(), 0.0); i.update(&empty).unwrap(); - let r2 = i.result(); + let r2 = i.to_sketch(true); assert_eq!(r2.num_retained(), 0); assert!(r2.is_empty()); assert!(!r2.is_estimation_mode()); @@ -155,7 +155,7 @@ fn test_non_empty_no_retained_keys() { let mut i = ThetaIntersection::new_with_default_seed(); i.update(&s).unwrap(); - let r1 = i.result(); + let r1 = i.to_sketch(true); assert_eq!(r1.num_retained(), 0); assert!(!r1.is_empty()); assert!(r1.is_estimation_mode()); @@ -163,7 +163,7 @@ fn test_non_empty_no_retained_keys() { assert_eq!(r1.estimate(), 0.0); i.update(&s).unwrap(); - let r2 = i.result(); + let r2 = i.to_sketch(true); assert_eq!(r2.num_retained(), 0); assert!(!r2.is_empty()); assert!(r2.is_estimation_mode()); @@ -179,7 +179,7 @@ fn test_exact_half_overlap_unordered() { let mut i = ThetaIntersection::new_with_default_seed(); i.update(&s1).unwrap(); i.update(&s2).unwrap(); - let r = i.result(); + let r = i.to_sketch(true); assert!(!r.is_empty()); assert!(!r.is_estimation_mode()); @@ -194,7 +194,7 @@ fn test_exact_half_overlap_ordered() { let mut i = ThetaIntersection::new_with_default_seed(); i.update(&s1.compact(true)).unwrap(); i.update(&s2.compact(true)).unwrap(); - let r = i.result(); + let r = i.to_sketch(true); assert!(!r.is_empty()); assert!(!r.is_estimation_mode()); @@ -209,7 +209,7 @@ fn test_exact_disjoint_unordered() { let mut i = ThetaIntersection::new_with_default_seed(); i.update(&s1).unwrap(); i.update(&s2).unwrap(); - let r = i.result(); + let r = i.to_sketch(true); assert!(r.is_empty()); assert!(!r.is_estimation_mode()); @@ -224,7 +224,7 @@ fn test_exact_disjoint_ordered() { let mut i = ThetaIntersection::new_with_default_seed(); i.update(&s1.compact(true)).unwrap(); i.update(&s2.compact(true)).unwrap(); - let r = i.result(); + let r = i.to_sketch(true); assert!(r.is_empty()); assert!(!r.is_estimation_mode()); @@ -239,7 +239,7 @@ fn test_estimation_half_overlap_unordered() { let mut i = ThetaIntersection::new_with_default_seed(); i.update(&s1).unwrap(); i.update(&s2).unwrap(); - let r = i.result(); + let r = i.to_sketch(true); assert!(!r.is_empty()); assert!(r.is_estimation_mode()); @@ -254,7 +254,7 @@ fn test_estimation_half_overlap_ordered() { let mut i = ThetaIntersection::new_with_default_seed(); i.update(&s1.compact(true)).unwrap(); i.update(&s2.compact(true)).unwrap(); - let r = i.result(); + let r = i.to_sketch(true); assert!(!r.is_empty()); assert!(r.is_estimation_mode()); @@ -271,7 +271,7 @@ fn test_estimation_half_overlap_ordered_deserialized_compact() { let mut i = ThetaIntersection::new_with_default_seed(); i.update(&c1).unwrap(); i.update(&c2).unwrap(); - let r = i.result(); + let r = i.to_sketch(true); assert!(!r.is_empty()); assert!(r.is_estimation_mode()); @@ -286,7 +286,7 @@ fn test_estimation_disjoint_unordered() { let mut i = ThetaIntersection::new_with_default_seed(); i.update(&s1).unwrap(); i.update(&s2).unwrap(); - let r = i.result(); + let r = i.to_sketch(true); assert!(!r.is_empty()); assert!(r.is_estimation_mode()); @@ -301,7 +301,7 @@ fn test_estimation_disjoint_ordered() { let mut i = ThetaIntersection::new_with_default_seed(); i.update(&s1.compact(true)).unwrap(); i.update(&s2.compact(true)).unwrap(); - let r = i.result(); + let r = i.to_sketch(true); assert!(!r.is_empty()); assert!(r.is_estimation_mode()); diff --git a/datasketches/tests/theta_union_test.rs b/datasketches/tests/theta_union_test.rs index 5eeacec..e6e9160 100644 --- a/datasketches/tests/theta_union_test.rs +++ b/datasketches/tests/theta_union_test.rs @@ -45,13 +45,13 @@ fn assert_estimate_close(sketch: &CompactThetaSketch, expected: f64, tolerance: fn test_empty_union() { let sketch = ThetaSketch::builder().build(); let mut union = ThetaUnion::builder().build(); - let result = union.result(); + let result = union.to_sketch(true); assert_eq!(result.num_retained(), 0); assert!(result.is_empty()); assert!(!result.is_estimation_mode()); union.update(&sketch).unwrap(); - let result = union.result(); + let result = union.to_sketch(true); assert_eq!(result.num_retained(), 0); assert!(result.is_empty()); assert!(!result.is_estimation_mode()); @@ -64,7 +64,7 @@ fn test_non_empty_no_retained_keys() { let mut union = ThetaUnion::builder().build(); union.update(&sketch).unwrap(); - let result = union.result(); + let result = union.to_sketch(true); assert_eq!(result.num_retained(), 0); assert!(!result.is_empty()); assert!(result.is_estimation_mode()); @@ -86,13 +86,13 @@ fn test_exact_mode_half_overlap() { let mut union = ThetaUnion::builder().build(); union.update(&sketch1).unwrap(); union.update(&sketch2).unwrap(); - let result = union.result(); + let result = union.to_sketch(true); assert!(!result.is_empty()); assert!(!result.is_estimation_mode()); assert_eq!(result.estimate(), 1500.0); union.reset(); - let reset = union.result(); + let reset = union.to_sketch(true); assert_eq!(reset.num_retained(), 0); assert!(reset.is_empty()); assert!(!reset.is_estimation_mode()); @@ -115,7 +115,7 @@ fn test_exact_mode_half_overlap_compact() { let mut union = ThetaUnion::builder().build(); union.update(&compact1).unwrap(); union.update(&compact2).unwrap(); - let result = union.result(); + let result = union.to_sketch(true); assert!(!result.is_empty()); assert!(!result.is_estimation_mode()); assert_eq!(result.estimate(), 1500.0); @@ -136,7 +136,7 @@ fn test_estimation_mode_half_overlap() { let mut union = ThetaUnion::builder().build(); union.update(&sketch1).unwrap(); union.update(&sketch2).unwrap(); - let result = union.result(); + let result = union.to_sketch(true); assert!(!result.is_empty()); assert!(result.is_estimation_mode()); assert!( @@ -178,14 +178,14 @@ fn test_larger_k() { union1.update(&sketch2).unwrap(); union1.update(&sketch1).unwrap(); union1.update(&sketch3).unwrap(); - let result1 = union1.result(); + let result1 = union1.to_sketch(true); assert_eq!(result1.estimate(), sketch3.estimate()); let mut union2 = ThetaUnion::builder().lg_k(16).build(); union2.update(&sketch1).unwrap(); union2.update(&sketch3).unwrap(); union2.update(&sketch2).unwrap(); - let result2 = union2.result(); + let result2 = union2.to_sketch(true); assert_eq!(result2.estimate(), sketch3.estimate()); } @@ -200,7 +200,7 @@ fn test_exact_union_no_overlap() { union.update(&sketch1).unwrap(); union.update(&sketch2).unwrap(); - let result = union.result(); + let result = union.to_sketch(true); assert!(!result.is_empty()); assert!(!result.is_estimation_mode()); assert_eq!(result.estimate(), k as f64); @@ -217,7 +217,11 @@ fn test_estimation_union_no_overlap() { union.update(&sketch1).unwrap(); union.update(&sketch2).unwrap(); - assert_estimate_close(&union.result(), (4 * k) as f64, 0.05 * (4 * k) as f64); + assert_estimate_close( + &union.to_sketch(true), + (4 * k) as f64, + 0.05 * (4 * k) as f64, + ); } #[test] @@ -231,7 +235,7 @@ fn test_exact_union_with_overlap() { union.update(&sketch1).unwrap(); union.update(&sketch2).unwrap(); - let result = union.result(); + let result = union.to_sketch(true); assert!(!result.is_empty()); assert!(!result.is_estimation_mode()); assert_eq!(result.estimate(), k as f64); @@ -255,11 +259,11 @@ fn test_ordered_and_unordered_compact_inputs() { unordered_union.update(&compact_unordered).unwrap(); assert_eq!( - ordered_union.result().estimate(), - unordered_union.result().estimate() + ordered_union.to_sketch(true).estimate(), + unordered_union.to_sketch(true).estimate() ); assert_estimate_close( - &ordered_union.result(), + &ordered_union.to_sketch(true), (4 * k) as f64, 0.05 * (4 * k) as f64, ); @@ -274,8 +278,8 @@ fn test_result_ordering_forms_have_same_estimate() { union.update(&sketch1).unwrap(); union.update(&sketch2).unwrap(); - let unordered = union.result_with_ordered(false); - let ordered = union.result_with_ordered(true); + let unordered = union.to_sketch(false); + let ordered = union.to_sketch(true); assert!(!unordered.is_ordered()); assert!(ordered.is_ordered()); @@ -298,7 +302,7 @@ fn test_multi_union() { union.update(&sketch).unwrap(); } - assert_estimate_close(&union.result(), 180_391.0, 180_391.0 * 0.02); + assert_estimate_close(&union.to_sketch(true), 180_391.0, 180_391.0 * 0.02); } #[test] @@ -311,8 +315,8 @@ fn test_result_does_not_reset_union() { let mut union = ThetaUnion::builder().lg_k(lg_k).build(); union.update(&compact1).unwrap(); union.update(&compact2).unwrap(); - let first = union.result(); - let second = union.result(); + let first = union.to_sketch(true); + let second = union.to_sketch(true); assert_eq!(first.estimate(), second.estimate()); assert!(!second.is_empty()); @@ -328,7 +332,7 @@ fn test_union_full_overlap() { let mut union = ThetaUnion::builder().lg_k(lg_k).build(); union.update(&compact1).unwrap(); union.update(&compact2).unwrap(); - let result = union.result(); + let result = union.to_sketch(true); assert_eq!(result.estimate(), k as f64); } @@ -357,8 +361,8 @@ fn test_ordered_input_early_stop_matches_unordered_input() { unordered_union.update(&unordered2).unwrap(); assert_eq!( - ordered_union.result().estimate(), - unordered_union.result().estimate() + ordered_union.to_sketch(true).estimate(), + unordered_union.to_sketch(true).estimate() ); } } @@ -373,7 +377,7 @@ fn test_union_cutback_to_k() { let mut union = ThetaUnion::builder().lg_k(lg_k).build(); union.update(&compact1).unwrap(); union.update(&compact2).unwrap(); - let result = union.result(); + let result = union.to_sketch(true); assert_estimate_close(&result, (6 * k) as f64, (6 * k) as f64 * 0.06); assert!(result.num_retained() <= k as usize); @@ -390,17 +394,17 @@ fn test_union_empty_valid_rules() { let mut empty_union = ThetaUnion::builder().lg_k(5).build(); empty_union.update(&empty1).unwrap(); empty_union.update(&empty2).unwrap(); - assert!(empty_union.result().is_empty()); + assert!(empty_union.to_sketch(true).is_empty()); let mut left_non_empty_union = ThetaUnion::builder().lg_k(5).build(); left_non_empty_union.update(&one).unwrap(); left_non_empty_union.update(&empty2).unwrap(); - assert!(!left_non_empty_union.result().is_empty()); + assert!(!left_non_empty_union.to_sketch(true).is_empty()); let mut right_non_empty_union = ThetaUnion::builder().lg_k(5).build(); right_non_empty_union.update(&empty1).unwrap(); right_non_empty_union.update(&one).unwrap(); - assert!(!right_non_empty_union.result().is_empty()); + assert!(!right_non_empty_union.to_sketch(true).is_empty()); } #[test] @@ -411,7 +415,7 @@ fn test_trim_to_k() { let mut union = ThetaUnion::builder().lg_k(10).build(); union.update(&hi_sketch).unwrap(); union.update(&lo_sketch).unwrap(); - let result = union.result(); + let result = union.to_sketch(true); assert_eq!(result.num_retained(), 1024); } @@ -422,7 +426,7 @@ fn test_builder_lg_k() { let mut union = ThetaUnion::builder().lg_k(10).build(); union.update(&sketch).unwrap(); - assert_eq!(union.result().estimate(), 1000.0); + assert_eq!(union.to_sketch(true).estimate(), 1000.0); } #[derive(Clone, Copy, Debug)] @@ -652,7 +656,7 @@ fn test_corner_case_union_states() { let mut union = ThetaUnion::builder().build(); union.update(&sketch_a).unwrap(); union.update(&sketch_b).unwrap(); - let result = union.result(); + let result = union.to_sketch(true); assert!( (result.theta() - expected_theta).abs() < 1e-6, @@ -675,7 +679,7 @@ fn test_corner_case_union_states() { let mut union = ThetaUnion::builder().build(); union.update(&compact_a).unwrap(); union.update(&compact_b).unwrap(); - let compact_result = union.result(); + let compact_result = union.to_sketch(true); assert!((compact_result.theta() - expected_theta).abs() < 1e-6); assert_eq!(compact_result.num_retained(), expected_count);