From 0b5d3c3c8a67bf615e0a23d17462db05a53084bb Mon Sep 17 00:00:00 2001 From: jaideeppyne Date: Mon, 31 Aug 2026 17:03:35 +0530 Subject: [PATCH 1/4] fix(theta,tuple): report MAX_THETA for an empty sampled sketch An empty sketch built with a sampling probability below 1.0 reported the sampling theta rather than MAX_THETA, so theta() returned p instead of 1.0 and is_estimation_mode() returned true. Java and C++ mask theta to MAX_THETA while the sketch is empty. The stale theta also reached ThetaANotB and TupleANotB results computed from an empty input, which serialized to a three-preamble-long image whose theta deserialization discarded, so those results did not survive a serialization round trip. --- CHANGELOG.md | 1 + datasketches/src/thetafamily/theta/sketch.rs | 13 ++++-- datasketches/src/thetafamily/tuple/sketch.rs | 24 +++++++++-- tests-integration/src/lib.rs | 3 ++ tests-integration/tests/theta_test/a_not_b.rs | 28 +++++++++++++ tests-integration/tests/theta_test/sketch.rs | 42 ++++++++++++++++--- tests-integration/tests/tuple_test/a_not_b.rs | 31 ++++++++++++++ tests-integration/tests/tuple_test/sketch.rs | 37 +++++++++++++++- 8 files changed, 166 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3f1f1bcf..b322863b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -48,6 +48,7 @@ All significant changes to this project will be documented in this file. * Malformed CPC images now return `InvalidData` instead of panicking. * Seeded deserializers now return `InvalidData` rather than panicking when the caller supplies a seed whose hash is the reserved zero value. * Fix T-Digest interpolation and tail calculations that could produce non-monotonic or out-of-range quantiles and invalid rank, CDF, or PMF values. +* An empty `ThetaSketch` or `TupleSketch` built with a sampling probability below `1.0` now reports `theta` of `1.0`, `theta64` of `MAX_THETA`, and `is_estimation_mode` of `false`, matching Java and C++. The sampling theta previously reported by such a sketch also reached `ThetaANotB` and `TupleANotB` results computed from an empty input, which serialized to an image whose theta deserialization then discarded, so those results did not survive a serialization round trip. ## v0.4.0 (2026-08-18) diff --git a/datasketches/src/thetafamily/theta/sketch.rs b/datasketches/src/thetafamily/theta/sketch.rs index 28e2c12d..85272fd1 100644 --- a/datasketches/src/thetafamily/theta/sketch.rs +++ b/datasketches/src/thetafamily/theta/sketch.rs @@ -253,12 +253,19 @@ impl ThetaSketch { /// Returns theta as a fraction in `[0.0, 1.0]`. pub fn theta(&self) -> f64 { - self.table.theta() as f64 / MAX_THETA as f64 + self.theta64() as f64 / MAX_THETA as f64 } /// Returns theta as a `u64`. + /// + /// An empty sketch reports `MAX_THETA` even when it was built with a sampling probability + /// below `1.0`, matching the other DataSketches implementations. pub fn theta64(&self) -> u64 { - self.table.theta() + if self.is_empty() { + MAX_THETA + } else { + self.table.theta() + } } /// Returns the 16-bit seed hash. @@ -273,7 +280,7 @@ impl ThetaSketch { /// Returns `true` if the sketch is in estimation mode. pub fn is_estimation_mode(&self) -> bool { - self.table.theta() < MAX_THETA + self.theta64() < MAX_THETA } /// Returns the number of retained entries. diff --git a/datasketches/src/thetafamily/tuple/sketch.rs b/datasketches/src/thetafamily/tuple/sketch.rs index 22e924f7..beec47e7 100644 --- a/datasketches/src/thetafamily/tuple/sketch.rs +++ b/datasketches/src/thetafamily/tuple/sketch.rs @@ -134,9 +134,18 @@ impl<'a, S> TupleSketchView<'a, S> { } /// Returns theta as a `u64` threshold. + /// + /// An empty sketch reports `MAX_THETA` even when it was built with a sampling probability + /// below `1.0`, matching the other DataSketches implementations. pub fn theta64(&self) -> u64 { match self.0 { - TupleSketchViewState::Mutable(table) => table.theta(), + TupleSketchViewState::Mutable(table) => { + if table.is_empty() { + MAX_THETA + } else { + table.theta() + } + } TupleSketchViewState::Compact(sketch) => sketch.theta64(), } } @@ -301,12 +310,19 @@ where /// Returns theta as a fraction in `[0.0, 1.0]`. pub fn theta(&self) -> f64 { - self.table.theta() as f64 / MAX_THETA as f64 + self.theta64() as f64 / MAX_THETA as f64 } /// Returns theta as a `u64`. + /// + /// An empty sketch reports `MAX_THETA` even when it was built with a sampling probability + /// below `1.0`, matching the other DataSketches implementations. pub fn theta64(&self) -> u64 { - self.table.theta() + if self.is_empty() { + MAX_THETA + } else { + self.table.theta() + } } /// Returns the 16-bit seed hash. @@ -321,7 +337,7 @@ where /// Returns `true` if the sketch is in estimation mode. pub fn is_estimation_mode(&self) -> bool { - self.table.theta() < MAX_THETA + self.theta64() < MAX_THETA } /// Returns the number of retained entries. diff --git a/tests-integration/src/lib.rs b/tests-integration/src/lib.rs index 2407ef3b..bf99dcc5 100644 --- a/tests-integration/src/lib.rs +++ b/tests-integration/src/lib.rs @@ -19,3 +19,6 @@ /// A seed whose 16-bit seed hash is the reserved zero value. pub const ZERO_HASH_SEED: u64 = 50_541; + +/// The Theta-family theta value that represents exact (non-sampled) mode. +pub const MAX_THETA: u64 = i64::MAX as u64; diff --git a/tests-integration/tests/theta_test/a_not_b.rs b/tests-integration/tests/theta_test/a_not_b.rs index 9b17c006..5e62c1c9 100644 --- a/tests-integration/tests/theta_test/a_not_b.rs +++ b/tests-integration/tests/theta_test/a_not_b.rs @@ -29,6 +29,7 @@ use googletest::prelude::anything; use googletest::prelude::err; use googletest::prelude::lt; use googletest::prelude::near; +use tests_integration::MAX_THETA; fn sketch_with_range(start: u64, count: u64) -> ThetaSketch { let mut sketch = ThetaSketchBuilder::default().build().unwrap(); @@ -257,3 +258,30 @@ fn test_estimation_disjoint_returns_a() { assert!(r.is_estimation_mode()); assert_that!(r.estimate(), near(10000.0, 10000.0 * 0.02)); } + +#[test] +fn test_empty_sampled_inputs_produce_a_serializable_empty_result() { + for probability in [1.0, 0.5, 0.1, 0.001] { + let a = ThetaSketchBuilder::default() + .lg_k(12) + .sampling_probability(probability) + .build() + .unwrap(); + let b = ThetaSketchBuilder::default() + .lg_k(12) + .sampling_probability(probability) + .build() + .unwrap(); + + let r = ThetaANotB::default().compute(&a, &b, true).unwrap(); + + assert!(r.is_empty()); + assert_eq!(r.theta64(), MAX_THETA); + assert!(!r.is_estimation_mode()); + + let bytes = r.serialize(); + let restored = CompactThetaSketch::deserialize(&bytes).unwrap(); + assert_eq!(restored.serialize(), bytes); + assert_eq!(restored.theta64(), r.theta64()); + } +} diff --git a/tests-integration/tests/theta_test/sketch.rs b/tests-integration/tests/theta_test/sketch.rs index 3e86c936..cf82065b 100644 --- a/tests-integration/tests/theta_test/sketch.rs +++ b/tests-integration/tests/theta_test/sketch.rs @@ -25,6 +25,7 @@ use googletest::prelude::gt; use googletest::prelude::le; use googletest::prelude::lt; use googletest::prelude::near; +use tests_integration::MAX_THETA; use tests_integration::ZERO_HASH_SEED; #[test] @@ -291,23 +292,54 @@ fn test_bounds_all_num_std_devs() { } #[test] -fn test_bounds_empty_estimation_mode() { - // Create a sketch with sampling probability < 1.0 to force estimation mode +fn test_bounds_empty_with_sampling() { let sketch = ThetaSketchBuilder::default() .lg_k(12) .sampling_probability(0.1) .build() .unwrap(); - // The sketch is empty but theta < 1.0, so it's in estimation mode - // However, when empty, both bounds should return 0.0 per Java implementation assert!(sketch.is_empty()); - assert!(sketch.is_estimation_mode()); + assert!(!sketch.is_estimation_mode()); assert_eq!(sketch.estimate(), 0.0); assert_eq!(sketch.lower_bound(NumStdDev::One), 0.0); assert_eq!(sketch.upper_bound(NumStdDev::One), 0.0); } +#[test] +fn test_empty_sketch_reports_max_theta_for_every_sampling_probability() { + for probability in [1.0, 0.9, 0.5, 0.1, 0.01, 0.001] { + let sketch = ThetaSketchBuilder::default() + .lg_k(12) + .sampling_probability(probability) + .build() + .unwrap(); + + assert!(sketch.is_empty()); + assert_eq!(sketch.theta64(), MAX_THETA); + assert_eq!(sketch.theta(), 1.0); + assert!(!sketch.is_estimation_mode()); + assert_eq!(sketch.theta64(), sketch.compact(true).theta64()); + } +} + +#[test] +fn test_sampling_theta_applies_once_the_sketch_is_non_empty() { + let mut sketch = ThetaSketchBuilder::default() + .lg_k(12) + .sampling_probability(0.5) + .build() + .unwrap(); + + assert_eq!(sketch.theta64(), MAX_THETA); + + sketch.update(1u64); + + assert!(!sketch.is_empty()); + assert!(sketch.is_estimation_mode()); + assert_that!(sketch.theta64(), lt(MAX_THETA)); +} + #[test] fn test_compact_preserves_logical_non_empty_after_screened_update() { let screened_value = (0u64..) diff --git a/tests-integration/tests/tuple_test/a_not_b.rs b/tests-integration/tests/tuple_test/a_not_b.rs index 397d9576..eb9e9840 100644 --- a/tests-integration/tests/tuple_test/a_not_b.rs +++ b/tests-integration/tests/tuple_test/a_not_b.rs @@ -23,6 +23,7 @@ use googletest::assert_that; use googletest::prelude::all; use googletest::prelude::ge; use googletest::prelude::le; +use tests_integration::MAX_THETA; use crate::default_tuple_sketch_builder; use crate::tuple_sketch_with_range; @@ -185,3 +186,33 @@ fn estimation_bounds_cover_the_true_difference() { assert!(result.is_estimation_mode()); assert_that!(25_000.0, all!(ge(lower), le(upper))); } + +#[test] +fn empty_sampled_inputs_produce_a_serializable_empty_result() { + for probability in [1.0, 0.5, 0.1, 0.001] { + let a = default_tuple_sketch_builder() + .lg_k(12) + .sampling_probability(probability) + .build() + .unwrap(); + let b = default_tuple_sketch_builder() + .lg_k(12) + .sampling_probability(probability) + .build() + .unwrap(); + + assert_eq!(a.theta64(), MAX_THETA); + assert!(!a.is_estimation_mode()); + + let result = TupleANotB::default().compute(&a, &b, true).unwrap(); + + assert!(result.is_empty()); + assert_eq!(result.theta64(), MAX_THETA); + assert!(!result.is_estimation_mode()); + + let bytes = result.serialize(); + let restored = CompactTupleSketch::::deserialize(&bytes).unwrap(); + assert_eq!(restored.serialize(), bytes); + assert_eq!(restored.theta64(), result.theta64()); + } +} diff --git a/tests-integration/tests/tuple_test/sketch.rs b/tests-integration/tests/tuple_test/sketch.rs index a4d858f2..30bd74cc 100644 --- a/tests-integration/tests/tuple_test/sketch.rs +++ b/tests-integration/tests/tuple_test/sketch.rs @@ -29,6 +29,7 @@ use googletest::assert_that; use googletest::prelude::gt; use googletest::prelude::le; use googletest::prelude::lt; +use tests_integration::MAX_THETA; use tests_integration::ZERO_HASH_SEED; use crate::default_tuple_sketch_builder; @@ -183,12 +184,46 @@ fn empty_sampled_sketch_has_zero_bounds() { .unwrap(); assert!(sketch.is_empty()); - assert!(sketch.is_estimation_mode()); + assert!(!sketch.is_estimation_mode()); assert_eq!(sketch.estimate(), 0.0); assert_eq!(sketch.lower_bound(NumStdDev::Three), 0.0); assert_eq!(sketch.upper_bound(NumStdDev::Three), 0.0); } +#[test] +fn empty_sketch_reports_max_theta_for_every_sampling_probability() { + for probability in [1.0, 0.9, 0.5, 0.1, 0.01, 0.001] { + let sketch = default_tuple_sketch_builder() + .lg_k(12) + .sampling_probability(probability) + .build() + .unwrap(); + + assert!(sketch.is_empty()); + assert_eq!(sketch.theta64(), MAX_THETA); + assert_eq!(sketch.theta(), 1.0); + assert!(!sketch.is_estimation_mode()); + assert_eq!(sketch.theta64(), sketch.compact(true).theta64()); + } +} + +#[test] +fn sampling_theta_applies_once_the_sketch_is_non_empty() { + let mut sketch = default_tuple_sketch_builder() + .lg_k(12) + .sampling_probability(0.5) + .build() + .unwrap(); + + assert_eq!(sketch.theta64(), MAX_THETA); + + sketch.update(1u64, 1u64); + + assert!(!sketch.is_empty()); + assert!(sketch.is_estimation_mode()); + assert_that!(sketch.theta64(), lt(MAX_THETA)); +} + fn sorted_entries<'a>(entries: impl Iterator>) -> Vec<(u64, u64)> { let mut entries: Vec<_> = entries .map(|entry| (entry.hash(), *entry.summary())) From f175ce6a272e9b104a9481a3b738f2860b10f0d1 Mon Sep 17 00:00:00 2001 From: tison Date: Mon, 31 Aug 2026 22:58:39 +0800 Subject: [PATCH 2/4] refactor(theta,tuple): model canonical sketch states --- CHANGELOG.md | 2 +- .../src/thetafamily/common/a_not_b.rs | 114 +++--- .../src/thetafamily/common/hash_table.rs | 166 ++++---- .../src/thetafamily/common/intersection.rs | 152 ++++---- .../thetafamily/common/jaccard_similarity.rs | 128 +++---- datasketches/src/thetafamily/common/mod.rs | 13 +- .../src/thetafamily/common/sketch_state.rs | 270 +++++++++++++ datasketches/src/thetafamily/common/union.rs | 96 +++-- datasketches/src/thetafamily/theta/a_not_b.rs | 15 +- .../src/thetafamily/theta/intersection.rs | 22 +- datasketches/src/thetafamily/theta/sketch.rs | 355 +++++++++--------- datasketches/src/thetafamily/theta/union.rs | 17 +- datasketches/src/thetafamily/tuple/a_not_b.rs | 10 +- .../src/thetafamily/tuple/intersection.rs | 14 +- datasketches/src/thetafamily/tuple/sketch.rs | 202 +++++----- datasketches/src/thetafamily/tuple/union.rs | 9 +- .../tests/theta_test/intersection.rs | 3 +- tests-integration/tests/theta_test/union.rs | 35 ++ tests-integration/tests/tuple_test/union.rs | 36 ++ 19 files changed, 957 insertions(+), 702 deletions(-) create mode 100644 datasketches/src/thetafamily/common/sketch_state.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index b322863b..7baaad94 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -48,7 +48,7 @@ All significant changes to this project will be documented in this file. * Malformed CPC images now return `InvalidData` instead of panicking. * Seeded deserializers now return `InvalidData` rather than panicking when the caller supplies a seed whose hash is the reserved zero value. * Fix T-Digest interpolation and tail calculations that could produce non-monotonic or out-of-range quantiles and invalid rank, CDF, or PMF values. -* An empty `ThetaSketch` or `TupleSketch` built with a sampling probability below `1.0` now reports `theta` of `1.0`, `theta64` of `MAX_THETA`, and `is_estimation_mode` of `false`, matching Java and C++. The sampling theta previously reported by such a sketch also reached `ThetaANotB` and `TupleANotB` results computed from an empty input, which serialized to an image whose theta deserialization then discarded, so those results did not survive a serialization round trip. +* Empty Theta-family and Tuple-family sketches now consistently report `theta` of `1.0`, `theta64` of `MAX_THETA`, and `is_estimation_mode` of `false`, including update sketches and unions configured with a sampling probability below `1.0`. Empty compact results now use the canonical ordered representation, so set-operation results survive serialization round trips without changing their reported state. ## v0.4.0 (2026-08-18) diff --git a/datasketches/src/thetafamily/common/a_not_b.rs b/datasketches/src/thetafamily/common/a_not_b.rs index 643d7dfe..2f3bd0bd 100644 --- a/datasketches/src/thetafamily/common/a_not_b.rs +++ b/datasketches/src/thetafamily/common/a_not_b.rs @@ -23,9 +23,8 @@ use crate::hash::check_seed_hash; use crate::thetacommon::EntrySketch; use crate::thetacommon::KeySketch; use crate::thetacommon::SketchEntry; -use crate::thetacommon::SketchScalars; -use crate::thetacommon::constants::MAX_THETA; -use crate::thetacommon::hash_table::CompactSketchParts; +use crate::thetacommon::sketch_state::CompactSketchState; +use crate::thetacommon::sketch_state::ThetaThreshold; /// Computes `a and not b` for Theta-family sketch views. /// @@ -38,54 +37,51 @@ pub fn compute( a: A, b: B, ordered: bool, -) -> Result, Error> +) -> Result, Error> where A: EntrySketch, B: KeySketch, { - let SketchScalars { - seed_hash: a_seed_hash, - theta: a_theta, - empty: a_empty, - ordered: a_ordered, - .. - } = a.scalars(); + let a_metadata = a.metadata(); // If A is empty the result is an (empty) copy of A. As with the union and intersection, an // empty input carries no keys, so its seed is not validated. - if a_empty { - return Ok(parts_from_sketch(a, ordered)); + if a_metadata.is_empty() { + return Ok(compact_state_from_sketch(a, ordered)); } // A is non-empty, so its seed must be compatible. - check_seed_hash(seed_hash, a_seed_hash, "A", ErrorKind::InvalidArgument)?; + check_seed_hash( + seed_hash, + a_metadata.seed_hash(), + "A", + ErrorKind::InvalidArgument, + )?; - let SketchScalars { - seed_hash: b_seed_hash, - theta: b_theta, - empty: b_empty, - ordered: b_ordered, - num_retained: b_num_retained, - } = b.scalars(); + let b_metadata = b.metadata(); // An empty B subtracts nothing, so the result is simply a copy of A. This also covers the // "A is non-empty but has no retained keys" state: B's seed and theta must not influence // the result. - if b_empty { - return Ok(parts_from_sketch(a, ordered)); + if b_metadata.is_empty() { + return Ok(compact_state_from_sketch(a, ordered)); } // B is non-empty, so its seed must be compatible. - check_seed_hash(seed_hash, b_seed_hash, "B", ErrorKind::InvalidArgument)?; + check_seed_hash( + seed_hash, + b_metadata.seed_hash(), + "B", + ErrorKind::InvalidArgument, + )?; - let theta = a_theta.min(b_theta); - // A is non-empty here; the result only becomes empty if everything is subtracted in exact - // mode (handled below). - let mut is_empty = false; + let theta = a_metadata.theta().min(b_metadata.theta()); - let entries: Vec = if b_num_retained == 0 { - a.entries().filter(|entry| entry.hash() < theta).collect() - } else if a_ordered && b_ordered { + let entries: Vec = if b_metadata.num_retained() == 0 { + a.entries() + .filter(|entry| entry.hash() < theta.get()) + .collect() + } else if a_metadata.is_ordered() && b_metadata.is_ordered() { // Both inputs are sorted ascending by hash: merge-scan without a hash set. Only // B hashes below theta can exclude an A entry (A entries are all < theta), so // unexamined B entries at or above theta are harmless. @@ -93,7 +89,7 @@ where let mut entries = vec![]; for entry in a.entries() { let hash = entry.hash(); - if hash >= theta { + if hash >= theta.get() { break; } while let Some(&b_hash) = b_hashes.peek() { @@ -109,11 +105,11 @@ where } entries } else { - let mut b_keys: HashSet = HashSet::with_capacity(b_num_retained); + let mut b_keys: HashSet = HashSet::with_capacity(b_metadata.num_retained()); for hash in b.hashes() { - if hash < theta { + if hash < theta.get() { b_keys.insert(hash); - } else if b_ordered { + } else if b_metadata.is_ordered() { break; } } @@ -121,57 +117,51 @@ where let mut entries = vec![]; for entry in a.entries() { let hash = entry.hash(); - if hash < theta { + if hash < theta.get() { if !b_keys.contains(&hash) { entries.push(entry); } - } else if a_ordered { + } else if a_metadata.is_ordered() { break; } } entries }; - if entries.is_empty() && theta == MAX_THETA { - is_empty = true; + if entries.is_empty() && theta == ThetaThreshold::MAX { + return Ok(CompactSketchState::empty(seed_hash)); } - let out_ordered = ordered || a_ordered; let mut entries = entries; - if ordered && !a_ordered && entries.len() > 1 { + if ordered && !a_metadata.is_ordered() && entries.len() > 1 { entries.sort_unstable_by_key(SketchEntry::hash); } + let out_ordered = + ordered || a_metadata.is_ordered() || (entries.len() == 1 && theta == ThetaThreshold::MAX); - Ok(CompactSketchParts { + Ok(CompactSketchState::non_empty( entries, theta, seed_hash, - ordered: out_ordered, - empty: is_empty, - }) + out_ordered, + )) } -fn parts_from_sketch(sketch: S, ordered: bool) -> CompactSketchParts +fn compact_state_from_sketch(sketch: S, ordered: bool) -> CompactSketchState where S: EntrySketch, { - let SketchScalars { - seed_hash, - theta, - empty, - ordered: input_ordered, - .. - } = sketch.scalars(); + let metadata = sketch.metadata(); + if metadata.is_empty() { + return CompactSketchState::empty(metadata.seed_hash()); + } + let mut entries: Vec = sketch.entries().collect(); - let out_ordered = ordered || input_ordered; - if ordered && !input_ordered && entries.len() > 1 { + if ordered && !metadata.is_ordered() && entries.len() > 1 { entries.sort_unstable_by_key(SketchEntry::hash); } - CompactSketchParts { - entries, - theta, - seed_hash, - ordered: out_ordered, - empty, - } + let theta = metadata.theta(); + let out_ordered = + ordered || metadata.is_ordered() || (entries.len() == 1 && theta == ThetaThreshold::MAX); + CompactSketchState::non_empty(entries, theta, metadata.seed_hash(), out_ordered) } diff --git a/datasketches/src/thetafamily/common/hash_table.rs b/datasketches/src/thetafamily/common/hash_table.rs index 8400f69c..bc71acb1 100644 --- a/datasketches/src/thetafamily/common/hash_table.rs +++ b/datasketches/src/thetafamily/common/hash_table.rs @@ -27,19 +27,11 @@ use crate::thetacommon::SketchEntry; use crate::thetacommon::constants::HASH_TABLE_REBUILD_THRESHOLD; use crate::thetacommon::constants::HASH_TABLE_RESIZE_THRESHOLD; use crate::thetacommon::constants::MAX_LG_K; -use crate::thetacommon::constants::MAX_THETA; use crate::thetacommon::constants::MIN_LG_K; use crate::thetacommon::constants::STRIDE_MASK; - -/// Compact-sketch state from which a sketch family creates its compact result type. -#[derive(Debug)] -pub struct CompactSketchParts { - pub entries: Vec, - pub theta: u64, - pub seed_hash: u16, - pub ordered: bool, - pub empty: bool, -} +use crate::thetacommon::sketch_state::CompactSketchState; +use crate::thetacommon::sketch_state::ThetaSketchState; +use crate::thetacommon::sketch_state::ThetaThreshold; pub struct SketchHashTableIter<'a, E>(slice::Iter<'a, Option>); @@ -58,7 +50,7 @@ impl<'a, E> Iterator for SketchHashTableIter<'a, E> { /// Generic hash-table mechanics shared by Theta and Tuple sketches. /// /// 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. +/// theta screening, probing, resizing, rebuilding, and trimming. /// /// 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. @@ -75,15 +67,10 @@ pub struct SketchHashTable { seed: u64, seed_hash: u16, - // Logical emptiness of the source set. - // - // * `false` if any update has been attempted (even if screened by theta) - // * `true` if no updates have been attempted. - // - // This can be false even when `num_retained` is 0. - is_empty: bool, - - theta: u64, + // The operational threshold used to screen future updates. This is intentionally independent + // of the sketch's externally visible empty state: a never-updated sketch built with p < 1.0 + // must retain this sampling threshold even though its public theta is MAX_THETA. + retention_theta: ThetaThreshold, entries: Vec>, @@ -120,32 +107,48 @@ where let seed_hash = compute_seed_hash(seed, ErrorKind::InvalidArgument)?; let lg_max_size = lg_nom_size + 1; let lg_cur_size = starting_sub_multiple(lg_max_size, MIN_LG_K, resize_factor.lg_value()); - Ok(Self::from_raw_parts( + Ok(Self::allocate_empty( lg_cur_size, lg_nom_size, resize_factor, sampling_probability, - starting_theta_from_sampling_probability(sampling_probability), + starting_retention_theta(sampling_probability), seed, seed_hash, - true, )) } - /// Constructs a table from raw internal state. + /// Creates a table used internally by a Theta-family set operation. /// /// # Panics /// /// Panics if `lg_cur_size > lg_nom_size + 1`. (`lg_nom_size + 1 == lg_max_size`) - pub fn from_raw_parts( + pub fn for_set_operation( + lg_cur_size: u8, + lg_nom_size: u8, + retention_theta: ThetaThreshold, + seed: u64, + seed_hash: u16, + ) -> Self { + Self::allocate_empty( + lg_cur_size, + lg_nom_size, + ResizeFactor::X1, + 1.0, + retention_theta, + seed, + seed_hash, + ) + } + + fn allocate_empty( lg_cur_size: u8, lg_nom_size: u8, resize_factor: ResizeFactor, sampling_probability: f32, - theta: u64, + retention_theta: ThetaThreshold, seed: u64, seed_hash: u16, - is_empty: bool, ) -> Self { let lg_max_size = lg_nom_size + 1; assert!( @@ -162,8 +165,7 @@ where sampling_probability, seed, seed_hash, - is_empty, - theta, + retention_theta, entries, num_retained: 0, } @@ -195,9 +197,7 @@ where where F: FnOnce(Option<&mut E>) -> Option, { - self.is_empty = false; - - if hash == 0 || hash >= self.theta { + if hash == 0 || hash >= self.retention_theta.get() { return false; } @@ -260,9 +260,9 @@ where } } - /// Reset the table to empty state. + /// Restores the table's initial capacity and retention threshold and removes all entries. pub fn reset(&mut self) { - let init_theta = starting_theta_from_sampling_probability(self.sampling_probability); + let initial_retention_theta = starting_retention_theta(self.sampling_probability); let init_lg_cur = starting_sub_multiple( self.lg_nom_size + 1, MIN_LG_K, @@ -273,8 +273,7 @@ where self.entries.clear(); self.entries.resize_with(size, || None); self.num_retained = 0; - self.theta = init_theta; - self.is_empty = true; + self.retention_theta = initial_retention_theta; self.lg_cur_size = init_lg_cur; } @@ -283,14 +282,9 @@ where self.num_retained } - /// Get theta. - pub fn theta(&self) -> u64 { - self.theta - } - - /// Check logical emptiness of the source set. - pub fn is_empty(&self) -> bool { - self.is_empty + /// Returns the operational theta used to screen retained entries. + pub fn retention_theta(&self) -> ThetaThreshold { + self.retention_theta } /// Get iterator over retained entries. @@ -298,30 +292,30 @@ where SketchHashTableIter(self.entries.iter()) } - /// Returns the retained entries and theta as compact-sketch parts. - /// - /// An empty table reports `MAX_THETA` rather than its current theta, matching Java's - /// `correctThetaOnCompact()` behavior for never-updated sketches initialized with p < 1.0. - /// Empty and single-entry exact-mode results are always marked ordered (Java/C++ - /// compatibility). - pub fn to_compact_parts(&self, ordered: bool) -> CompactSketchParts + /// Creates canonical compact-sketch state from this table and its owning sketch state. + pub fn to_compact_sketch_state( + &self, + theta_sketch_state: ThetaSketchState, + ordered: bool, + ) -> CompactSketchState where E: Clone, { - let mut entries: Vec = self.iter_entries().cloned().collect(); - let empty = self.is_empty(); - let theta = if empty { MAX_THETA } else { self.theta() }; - let is_single = entries.len() == 1 && theta == MAX_THETA; - let ordered = ordered || empty || is_single; - if ordered && entries.len() > 1 { - entries.sort_unstable_by_key(SketchEntry::hash); - } - CompactSketchParts { - entries, - theta, - seed_hash: self.seed_hash(), - ordered, - empty, + match theta_sketch_state { + ThetaSketchState::Empty => { + debug_assert_eq!(self.num_retained, 0); + CompactSketchState::empty(self.seed_hash) + } + ThetaSketchState::NonEmpty { theta } => { + debug_assert_eq!(theta, self.retention_theta); + let mut retained_entries: Vec = self.iter_entries().cloned().collect(); + let ordered = + ordered || (retained_entries.len() == 1 && theta == ThetaThreshold::MAX); + if ordered && retained_entries.len() > 1 { + retained_entries.sort_unstable_by_key(SketchEntry::hash); + } + CompactSketchState::non_empty(retained_entries, theta, self.seed_hash, ordered) + } } } @@ -335,23 +329,14 @@ where self.seed_hash } - /// Set empty flag. - pub fn set_empty(&mut self, is_empty: bool) { - self.is_empty = is_empty; - } - /// Get the seed used by this table. pub fn seed(&self) -> u64 { self.seed } - /// Sets theta value. - pub fn set_theta(&mut self, theta: u64) { - assert!( - (1..=MAX_THETA).contains(&theta), - "theta must be in [1, {MAX_THETA}], got {theta}" - ); - self.theta = theta; + /// Sets the operational theta used to screen retained entries. + pub fn set_retention_theta(&mut self, retention_theta: ThetaThreshold) { + self.retention_theta = retention_theta; } /// Returns minimal lg_size where rebuild-capacity can hold `count`. @@ -435,7 +420,7 @@ where let (_lesser, kth, _greater) = retained.select_nth_unstable_by_key(k, |e| e.hash()); kth.hash() }; - self.theta = kth_hash; + self.retention_theta = ThetaThreshold::new(kth_hash); retained.truncate(k); let size = 1 << self.lg_cur_size; @@ -478,11 +463,26 @@ pub fn starting_sub_multiple(lg_target: u8, lg_min: u8, lg_resize_factor: u8) -> } } -/// Compute initial theta for hash table based on sampling probability. -pub fn starting_theta_from_sampling_probability(sampling_probability: f32) -> u64 { +/// Computes the initial operational theta from a sampling probability. +pub fn starting_retention_theta(sampling_probability: f32) -> ThetaThreshold { if sampling_probability < 1.0 { - (MAX_THETA as f64 * sampling_probability as f64) as u64 + let scaled_theta = (ThetaThreshold::MAX.get() as f64 * sampling_probability as f64) as u64; + // Threshold one and zero screen the same set of usable hashes because hash zero is + // reserved. Keep the state valid when a positive f32 probability rounds below one. + ThetaThreshold::new(scaled_theta.max(1)) } else { - MAX_THETA + ThetaThreshold::MAX + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn smallest_positive_probability_has_a_valid_retention_theta() { + let theta = starting_retention_theta(f32::from_bits(1)); + + assert_eq!(theta.get(), 1); } } diff --git a/datasketches/src/thetafamily/common/intersection.rs b/datasketches/src/thetafamily/common/intersection.rs index 45d47822..a49ed92f 100644 --- a/datasketches/src/thetafamily/common/intersection.rs +++ b/datasketches/src/thetafamily/common/intersection.rs @@ -15,18 +15,17 @@ // specific language governing permissions and limitations // under the License. -use crate::common::ResizeFactor; use crate::error::Error; use crate::error::ErrorKind; use crate::hash::check_seed_hash; use crate::hash::compute_seed_hash; use crate::thetacommon::EntrySketch; use crate::thetacommon::SketchEntry; -use crate::thetacommon::SketchScalars; use crate::thetacommon::constants::HASH_TABLE_REBUILD_THRESHOLD; -use crate::thetacommon::constants::MAX_THETA; -use crate::thetacommon::hash_table::CompactSketchParts; use crate::thetacommon::hash_table::SketchHashTable; +use crate::thetacommon::sketch_state::CompactSketchState; +use crate::thetacommon::sketch_state::ThetaSketchState; +use crate::thetacommon::sketch_state::ThetaThreshold; /// Merges an incoming entry into an existing entry with the same hash. /// @@ -44,7 +43,15 @@ pub trait IntersectionMergePolicy { pub struct IntersectionState { table: SketchHashTable, policy: P, - has_result: bool, + result_state: IntersectionResultState, +} + +/// State of the value currently represented by an intersection operator. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum IntersectionResultState { + Uninitialized, + Empty, + NonEmpty, } impl IntersectionState @@ -55,17 +62,8 @@ where pub fn new(seed: u64, policy: P) -> Result { let seed_hash = compute_seed_hash(seed, ErrorKind::InvalidArgument)?; Ok(Self { - has_result: false, - table: SketchHashTable::from_raw_parts( - 0, - 0, - ResizeFactor::X1, - 1.0, - MAX_THETA, - seed, - seed_hash, - false, - ), + result_state: IntersectionResultState::Uninitialized, + table: SketchHashTable::for_set_operation(0, 0, ThetaThreshold::MAX, seed, seed_hash), policy, }) } @@ -80,76 +78,64 @@ where E: Clone, P: IntersectionMergePolicy, { - let SketchScalars { - seed_hash, - theta, - empty, - ordered, - num_retained, - } = sketch.scalars(); - let new_default_table = |table: &SketchHashTable| { - SketchHashTable::from_raw_parts( + let metadata = sketch.metadata(); + let table_without_entries = |table: &SketchHashTable, retention_theta| { + SketchHashTable::for_set_operation( 0, 0, - ResizeFactor::X1, - 1.0, - table.theta(), + retention_theta, table.seed(), table.seed_hash(), - table.is_empty(), ) }; - if self.table.is_empty() { + if self.result_state == IntersectionResultState::Empty { return Ok(()); } - if empty { - self.table.set_empty(true); - } else { - check_seed_hash( - self.table.seed_hash(), - seed_hash, - "intersection update", - ErrorKind::InvalidArgument, - )?; + if metadata.is_empty() { + self.result_state = IntersectionResultState::Empty; + self.table = table_without_entries(&self.table, ThetaThreshold::MAX); + return Ok(()); } - self.table.set_theta(if self.table.is_empty() { - MAX_THETA - } else { - self.table.theta().min(theta) - }); + check_seed_hash( + self.table.seed_hash(), + metadata.seed_hash(), + "intersection update", + ErrorKind::InvalidArgument, + )?; + + let result_theta = self.table.retention_theta().min(metadata.theta()); + self.table.set_retention_theta(result_theta); - if self.has_result && self.table.num_retained() == 0 { + if self.result_state == IntersectionResultState::NonEmpty && self.table.num_retained() == 0 + { return Ok(()); } - if num_retained == 0 { - self.has_result = true; - self.table = new_default_table(&self.table); + if metadata.num_retained() == 0 { + self.result_state = IntersectionResultState::NonEmpty; + self.table = table_without_entries(&self.table, result_theta); return Ok(()); } // first update, copy incoming entries - if !self.has_result { - self.has_result = true; + if self.result_state == IntersectionResultState::Uninitialized { + self.result_state = IntersectionResultState::NonEmpty; let lg_size = SketchHashTable::::lg_size_from_count_for_rebuild( - num_retained, + metadata.num_retained(), HASH_TABLE_REBUILD_THRESHOLD, ); - // num_retained >= 1 here (the zero case returned early above), so lg_size >= 1 and - // lg_size - 1 below cannot underflow. + // The retained count is at least one here, so lg_size >= 1 and lg_size - 1 below + // cannot underflow. debug_assert!(lg_size >= 1); - self.table = SketchHashTable::from_raw_parts( + self.table = SketchHashTable::for_set_operation( lg_size, lg_size - 1, - ResizeFactor::X1, - 1.0, - self.table.theta(), + result_theta, self.table.seed(), self.table.seed_hash(), - self.table.is_empty(), ); for entry in sketch.entries() { let hash = entry.hash(); @@ -163,18 +149,18 @@ where } } // Safety check. - if self.table.num_retained() != num_retained { + if self.table.num_retained() != metadata.num_retained() { return Err(Error::invalid_argument( "num entries mismatch, possibly corrupted input sketch", )); } } else { - let max_matches = self.table.num_retained().min(num_retained); + let max_matches = self.table.num_retained().min(metadata.num_retained()); let mut matched_entries = Vec::with_capacity(max_matches); let mut count = 0; for entry in sketch.entries() { let hash = entry.hash(); - if hash < self.table.theta() { + if hash < self.table.retention_theta().get() { if let Some(existing) = self.table.entry(hash) { if matched_entries.len() == max_matches { return Err(Error::invalid_argument( @@ -185,25 +171,25 @@ where self.policy.merge(&mut merged, entry); matched_entries.push(merged); } - } else if ordered { + } else if metadata.is_ordered() { break; // early stop for ordered sketches } count += 1; } // Safety check. - if count > num_retained { + if count > metadata.num_retained() { return Err(Error::invalid_argument( "more keys than expected, possibly corrupted input sketch", )); - } else if !ordered && count < num_retained { + } else if !metadata.is_ordered() && count < metadata.num_retained() { return Err(Error::invalid_argument( "fewer keys than expected, possibly corrupted input sketch", )); } if matched_entries.is_empty() { - self.table = new_default_table(&self.table); - if self.table.theta() == MAX_THETA { - self.table.set_empty(true); + self.table = table_without_entries(&self.table, result_theta); + if result_theta == ThetaThreshold::MAX { + self.result_state = IntersectionResultState::Empty; } } else { let lg_size = SketchHashTable::::lg_size_from_count_for_rebuild( @@ -213,15 +199,12 @@ where // matched_entries is non-empty here (the empty case is handled above), so // lg_size >= 1 and lg_size - 1 below cannot underflow. debug_assert!(lg_size >= 1); - self.table = SketchHashTable::from_raw_parts( + self.table = SketchHashTable::for_set_operation( lg_size, lg_size - 1, - ResizeFactor::X1, - 1.0, - self.table.theta(), + result_theta, self.table.seed(), self.table.seed_hash(), - self.table.is_empty(), ); for entry in matched_entries { let hash = entry.hash(); @@ -241,7 +224,7 @@ where /// Returns whether this operator has received at least one update. pub fn has_result(&self) -> bool { - self.has_result + self.result_state != IntersectionResultState::Uninitialized } /// Returns the estimated size of the heap allocations in bytes. @@ -249,21 +232,20 @@ where self.table.estimated_size() } - /// Returns the current intersection state as compact-sketch parts. - pub fn to_compact_parts(&self, ordered: bool) -> CompactSketchParts + /// Returns the current intersection as canonical compact-sketch state. + pub fn to_compact_sketch_state(&self, ordered: bool) -> Option> where E: Clone, { - let mut entries: Vec = self.table.iter_entries().cloned().collect(); - if ordered { - entries.sort_unstable_by_key(SketchEntry::hash); - } - CompactSketchParts { - entries, - theta: self.table.theta(), - seed_hash: self.table.seed_hash(), - ordered, - empty: self.table.is_empty(), + match self.result_state { + IntersectionResultState::Uninitialized => None, + IntersectionResultState::Empty => { + Some(CompactSketchState::empty(self.table.seed_hash())) + } + IntersectionResultState::NonEmpty => Some(self.table.to_compact_sketch_state( + ThetaSketchState::non_empty(self.table.retention_theta()), + ordered, + )), } } } diff --git a/datasketches/src/thetafamily/common/jaccard_similarity.rs b/datasketches/src/thetafamily/common/jaccard_similarity.rs index a761cbd5..80c7ed24 100644 --- a/datasketches/src/thetafamily/common/jaccard_similarity.rs +++ b/datasketches/src/thetafamily/common/jaccard_similarity.rs @@ -23,14 +23,15 @@ use crate::hash::compute_seed_hash; use crate::thetacommon::EntrySketch; use crate::thetacommon::KeySketch; use crate::thetacommon::SketchEntry; -use crate::thetacommon::SketchScalars; +use crate::thetacommon::ThetaSketchMetadata; use crate::thetacommon::binomial_bounds; use crate::thetacommon::constants::MAX_LG_K; use crate::thetacommon::constants::MAX_THETA; use crate::thetacommon::constants::MIN_LG_K; -use crate::thetacommon::hash_table::CompactSketchParts; use crate::thetacommon::intersection::IntersectionMergePolicy; use crate::thetacommon::intersection::IntersectionState; +use crate::thetacommon::sketch_state::CompactSketchState; +use crate::thetacommon::sketch_state::ThetaThreshold; use crate::thetacommon::union::UnionMergePolicy; use crate::thetacommon::union::UnionState; @@ -70,7 +71,11 @@ impl JaccardSimilarity { } } - fn ratio_bounds(union_count: u64, intersection_count: u64, theta: u64) -> Result { + fn ratio_bounds( + union_count: u64, + intersection_count: u64, + theta: ThetaThreshold, + ) -> Result { if intersection_count > union_count { return Err(Error::invalid_argument(format!( "intersection count cannot exceed union count: {intersection_count} > {union_count}" @@ -84,13 +89,8 @@ impl JaccardSimilarity { }); } - let sampling_probability = theta as f64 / MAX_THETA as f64; - if sampling_probability <= 0.0 || sampling_probability > 1.0 { - return Err(Error::invalid_argument(format!( - "theta must produce a probability in (0.0, 1.0], got {sampling_probability}" - ))); - } - if sampling_probability == 1.0 { + let sampling_probability = theta.get() as f64 / MAX_THETA as f64; + if theta == ThetaThreshold::MAX { return Ok(Self::exact(intersection_count as f64 / union_count as f64)); } @@ -140,8 +140,8 @@ impl KeySketch for KeyEntries where S: KeySketch, { - fn scalars(self) -> SketchScalars { - self.0.scalars() + fn metadata(self) -> ThetaSketchMetadata { + self.0.metadata() } fn hashes(self) -> impl Iterator { @@ -165,27 +165,17 @@ where A: KeySketch, B: KeySketch, { - let SketchScalars { - theta: a_theta, - empty: a_empty, - num_retained: a_num_retained, - .. - } = sketch_a.scalars(); - let SketchScalars { - theta: b_theta, - empty: b_empty, - num_retained: b_num_retained, - .. - } = sketch_b.scalars(); - if a_empty && b_empty { + let a_metadata = sketch_a.metadata(); + let b_metadata = sketch_b.metadata(); + if a_metadata.is_empty() && b_metadata.is_empty() { return Ok(JaccardSimilarity::exact(1.0)); } - if a_empty || b_empty { + if a_metadata.is_empty() || b_metadata.is_empty() { return Ok(JaccardSimilarity::exact(0.0)); } - let sketch_a_state = (a_num_retained, a_theta); - let sketch_b_state = (b_num_retained, b_theta); + let sketch_a_state = (a_metadata.num_retained(), a_metadata.theta()); + let sketch_b_state = (b_metadata.num_retained(), b_metadata.theta()); let union = compute_union(seed, sketch_a, sketch_b)?; if identical_sets(sketch_a_state, sketch_b_state, &union) { return Ok(JaccardSimilarity::exact(1.0)); @@ -194,17 +184,20 @@ where let mut intersection = IntersectionState::new(seed, NoopMergePolicy)?; intersection.update(KeyEntries(sketch_a))?; intersection.update(KeyEntries(sketch_b))?; - let intersection = intersection.to_compact_parts(false); + let intersection = intersection + .to_compact_sketch_state(false) + .expect("two intersection updates must produce a result"); + let union_theta = union.theta_sketch_state().theta(); let intersection_count = intersection - .entries + .retained_entries() .iter() - .filter(|entry| entry.hash < union.theta) + .filter(|entry| entry.hash < union_theta.get()) .count(); JaccardSimilarity::ratio_bounds( - union.entries.len() as u64, + union.retained_entries().len() as u64, intersection_count as u64, - union.theta, + union_theta, ) } @@ -213,27 +206,17 @@ where A: KeySketch, B: KeySketch, { - let SketchScalars { - theta: a_theta, - empty: a_empty, - num_retained: a_num_retained, - .. - } = sketch_a.scalars(); - let SketchScalars { - theta: b_theta, - empty: b_empty, - num_retained: b_num_retained, - .. - } = sketch_b.scalars(); - if a_empty && b_empty { + let a_metadata = sketch_a.metadata(); + let b_metadata = sketch_b.metadata(); + if a_metadata.is_empty() && b_metadata.is_empty() { return Ok(true); } - if a_empty || b_empty { + if a_metadata.is_empty() || b_metadata.is_empty() { return Ok(false); } - let sketch_a_state = (a_num_retained, a_theta); - let sketch_b_state = (b_num_retained, b_theta); + let sketch_a_state = (a_metadata.num_retained(), a_metadata.theta()); + let sketch_b_state = (b_metadata.num_retained(), b_metadata.theta()); let union = compute_union(seed, sketch_a, sketch_b)?; Ok(identical_sets(sketch_a_state, sketch_b_state, &union)) } @@ -242,27 +225,29 @@ fn compute_union( seed: u64, sketch_a: A, sketch_b: B, -) -> Result, Error> +) -> Result, Error> where A: KeySketch, B: KeySketch, { - let SketchScalars { - seed_hash: a_seed_hash, - num_retained: a_num_retained, - .. - } = sketch_a.scalars(); - let SketchScalars { - seed_hash: b_seed_hash, - num_retained: b_num_retained, - .. - } = sketch_b.scalars(); + let a_metadata = sketch_a.metadata(); + let b_metadata = sketch_b.metadata(); let seed_hash = compute_seed_hash(seed, ErrorKind::InvalidArgument)?; - check_seed_hash(seed_hash, a_seed_hash, "A", ErrorKind::InvalidData)?; - check_seed_hash(seed_hash, b_seed_hash, "B", ErrorKind::InvalidData)?; + check_seed_hash( + seed_hash, + a_metadata.seed_hash(), + "A", + ErrorKind::InvalidData, + )?; + check_seed_hash( + seed_hash, + b_metadata.seed_hash(), + "B", + ErrorKind::InvalidData, + )?; let mut union = UnionState::new( - union_lg_k(a_num_retained, b_num_retained), + union_lg_k(a_metadata.num_retained(), b_metadata.num_retained()), ResizeFactor::X8, 1.0, seed, @@ -270,7 +255,7 @@ where )?; union.update(KeyEntries(sketch_a))?; union.update(KeyEntries(sketch_b))?; - Ok(union.to_compact_parts(false)) + Ok(union.to_compact_sketch_state(false)) } /// Returns whether both sketches have the same retained keys and theta. @@ -278,14 +263,15 @@ where /// When the union retains no additional keys and preserves both input theta values, each input /// contains exactly the same retained key set represented by the union. fn identical_sets( - sketch_a: (usize, u64), - sketch_b: (usize, u64), - union: &CompactSketchParts, + sketch_a: (usize, ThetaThreshold), + sketch_b: (usize, ThetaThreshold), + union: &CompactSketchState, ) -> bool { - union.entries.len() == sketch_a.0 - && union.entries.len() == sketch_b.0 - && union.theta == sketch_a.1 - && union.theta == sketch_b.1 + let union_state = union.theta_sketch_state(); + union.retained_entries().len() == sketch_a.0 + && union.retained_entries().len() == sketch_b.0 + && union_state.theta() == sketch_a.1 + && union_state.theta() == sketch_b.1 } fn sampling_adjuster(sampling_probability: f64) -> f64 { diff --git a/datasketches/src/thetafamily/common/mod.rs b/datasketches/src/thetafamily/common/mod.rs index 59f5dae8..7827ae6b 100644 --- a/datasketches/src/thetafamily/common/mod.rs +++ b/datasketches/src/thetafamily/common/mod.rs @@ -23,6 +23,7 @@ pub(super) mod constants; pub(super) mod hash_table; pub(super) mod intersection; pub(super) mod jaccard_similarity; +pub(super) mod sketch_state; pub(super) mod union; pub use self::jaccard_similarity::JaccardSimilarity; @@ -33,7 +34,7 @@ pub(super) trait SketchEntry { } pub(super) trait KeySketch: Copy { - fn scalars(self) -> SketchScalars; + fn metadata(self) -> ThetaSketchMetadata; fn hashes(self) -> impl Iterator; } @@ -44,12 +45,4 @@ pub(super) trait EntrySketch: KeySketch { fn entries(self) -> impl Iterator; } -/// Scalar sketch values inspected by Theta-family set operations. -#[derive(Clone, Copy, Debug)] -pub(super) struct SketchScalars { - pub seed_hash: u16, - pub theta: u64, - pub empty: bool, - pub ordered: bool, - pub num_retained: usize, -} +pub(super) use self::sketch_state::ThetaSketchMetadata; diff --git a/datasketches/src/thetafamily/common/sketch_state.rs b/datasketches/src/thetafamily/common/sketch_state.rs new file mode 100644 index 00000000..bdc93d71 --- /dev/null +++ b/datasketches/src/thetafamily/common/sketch_state.rs @@ -0,0 +1,270 @@ +// 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 std::num::NonZeroU64; + +use crate::thetacommon::constants::MAX_THETA; + +/// A validated Theta-family retention threshold. +/// +/// Hash zero is reserved and hashes are made non-negative by dropping their high bit, so every +/// usable threshold is in `1..=MAX_THETA`. +#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)] +#[repr(transparent)] +pub struct ThetaThreshold(NonZeroU64); + +impl ThetaThreshold { + pub const MAX: Self = Self(NonZeroU64::new(MAX_THETA).unwrap()); + + /// Creates a threshold if `value` is in the valid Theta-family range. + pub fn try_new(value: u64) -> Option { + let value = NonZeroU64::new(value)?; + (value.get() <= MAX_THETA).then_some(Self(value)) + } + + /// Creates a threshold known by the caller to be valid. + /// + /// # Panics + /// + /// Panics if `value` is outside `1..=MAX_THETA`. + pub fn new(value: u64) -> Self { + Self::try_new(value) + .unwrap_or_else(|| panic!("theta must be in [1, {MAX_THETA}], got {value}")) + } + + pub fn get(self) -> u64 { + self.0.get() + } + + pub fn is_estimation_mode(self) -> bool { + self < Self::MAX + } +} + +/// Canonical state exposed by a Theta-family sketch. +/// +/// `NonEmpty` means that the sketch has observed or represents non-empty input. It may still +/// retain zero entries when every observed hash was screened out by theta. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ThetaSketchState { + Empty, + NonEmpty { theta: ThetaThreshold }, +} + +impl ThetaSketchState { + pub fn non_empty(theta: ThetaThreshold) -> Self { + Self::NonEmpty { theta } + } + + pub fn theta(self) -> ThetaThreshold { + match self { + Self::Empty => ThetaThreshold::MAX, + Self::NonEmpty { theta } => theta, + } + } + + pub fn is_empty(self) -> bool { + matches!(self, Self::Empty) + } + + pub fn is_estimation_mode(self) -> bool { + matches!(self, Self::NonEmpty { theta } if theta.is_estimation_mode()) + } +} + +/// Whether an update sketch has observed an update call since construction or reset. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum UpdateSketchState { + NeverUpdated, + Updated, +} + +impl UpdateSketchState { + pub fn theta_sketch_state(self, retention_theta: ThetaThreshold) -> ThetaSketchState { + match self { + Self::NeverUpdated => ThetaSketchState::Empty, + Self::Updated => ThetaSketchState::non_empty(retention_theta), + } + } +} + +/// Observable metadata consumed by shared Theta-family set operations. +/// +/// Empty sketches cannot carry a retained count, theta, or ordering claim in this representation. +#[derive(Clone, Copy, Debug)] +pub enum ThetaSketchMetadata { + Empty { + seed_hash: u16, + }, + NonEmpty { + seed_hash: u16, + theta: ThetaThreshold, + ordered: bool, + num_retained: usize, + }, +} + +impl ThetaSketchMetadata { + pub fn from_theta_sketch_state( + seed_hash: u16, + theta_sketch_state: ThetaSketchState, + ordered: bool, + num_retained: usize, + ) -> Self { + match theta_sketch_state { + ThetaSketchState::Empty => { + debug_assert_eq!(num_retained, 0); + Self::Empty { seed_hash } + } + ThetaSketchState::NonEmpty { theta } => Self::NonEmpty { + seed_hash, + theta, + ordered, + num_retained, + }, + } + } + + pub fn seed_hash(self) -> u16 { + match self { + Self::Empty { seed_hash } | Self::NonEmpty { seed_hash, .. } => seed_hash, + } + } + + pub fn theta_sketch_state(self) -> ThetaSketchState { + match self { + Self::Empty { .. } => ThetaSketchState::Empty, + Self::NonEmpty { theta, .. } => ThetaSketchState::non_empty(theta), + } + } + + pub fn is_empty(self) -> bool { + matches!(self, Self::Empty { .. }) + } + + pub fn theta(self) -> ThetaThreshold { + self.theta_sketch_state().theta() + } + + pub fn is_ordered(self) -> bool { + match self { + Self::Empty { .. } => true, + Self::NonEmpty { ordered, .. } => ordered, + } + } + + pub fn num_retained(self) -> usize { + match self { + Self::Empty { .. } => 0, + Self::NonEmpty { num_retained, .. } => num_retained, + } + } +} + +/// Canonical in-memory state for a compact Theta-family sketch. +/// +/// The empty variant deliberately has neither retained entries nor theta. This makes the only +/// representable empty state use `MAX_THETA`, report exact mode, and serialize through the +/// canonical empty-image path. +#[derive(Clone, Debug)] +pub enum CompactSketchState { + Empty { + seed_hash: u16, + }, + NonEmpty { + retained_entries: Vec, + theta: ThetaThreshold, + seed_hash: u16, + ordered: bool, + }, +} + +impl CompactSketchState { + pub fn empty(seed_hash: u16) -> Self { + Self::Empty { seed_hash } + } + + pub fn non_empty( + retained_entries: Vec, + theta: ThetaThreshold, + seed_hash: u16, + ordered: bool, + ) -> Self { + Self::NonEmpty { + retained_entries, + theta, + seed_hash, + ordered, + } + } + + pub fn theta_sketch_state(&self) -> ThetaSketchState { + match self { + Self::Empty { .. } => ThetaSketchState::Empty, + Self::NonEmpty { theta, .. } => ThetaSketchState::non_empty(*theta), + } + } + + pub fn seed_hash(&self) -> u16 { + match self { + Self::Empty { seed_hash } | Self::NonEmpty { seed_hash, .. } => *seed_hash, + } + } + + pub fn retained_entries(&self) -> &[E] { + match self { + Self::Empty { .. } => &[], + Self::NonEmpty { + retained_entries, .. + } => retained_entries, + } + } + + pub fn retained_entries_capacity(&self) -> usize { + match self { + Self::Empty { .. } => 0, + Self::NonEmpty { + retained_entries, .. + } => retained_entries.capacity(), + } + } + + pub fn is_ordered(&self) -> bool { + match self { + Self::Empty { .. } => true, + Self::NonEmpty { ordered, .. } => *ordered, + } + } + + #[cfg(feature = "theta")] + pub fn map_retained_entries(self, mut f: impl FnMut(E) -> T) -> CompactSketchState { + match self { + Self::Empty { seed_hash } => CompactSketchState::empty(seed_hash), + Self::NonEmpty { + retained_entries, + theta, + seed_hash, + ordered, + } => CompactSketchState::non_empty( + retained_entries.into_iter().map(&mut f).collect(), + theta, + seed_hash, + ordered, + ), + } + } +} diff --git a/datasketches/src/thetafamily/common/union.rs b/datasketches/src/thetafamily/common/union.rs index 95af17bc..43201a04 100644 --- a/datasketches/src/thetafamily/common/union.rs +++ b/datasketches/src/thetafamily/common/union.rs @@ -21,10 +21,9 @@ use crate::error::ErrorKind; use crate::hash::check_seed_hash; use crate::thetacommon::EntrySketch; use crate::thetacommon::SketchEntry; -use crate::thetacommon::SketchScalars; -use crate::thetacommon::constants::MAX_THETA; -use crate::thetacommon::hash_table::CompactSketchParts; use crate::thetacommon::hash_table::SketchHashTable; +use crate::thetacommon::sketch_state::CompactSketchState; +use crate::thetacommon::sketch_state::ThetaThreshold; /// Merges an incoming entry into an existing entry with the same hash. pub trait UnionMergePolicy { @@ -39,7 +38,14 @@ pub trait UnionMergePolicy { pub struct UnionState { table: SketchHashTable, policy: P, - union_theta: u64, + result_state: UnionResultState, +} + +/// State of the value currently represented by a union operator. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum UnionResultState { + Empty, + NonEmpty { theta: ThetaThreshold }, } impl UnionState @@ -55,7 +61,7 @@ where ) -> Result { let table = SketchHashTable::new(lg_k, resize_factor, sampling_probability, seed)?; Ok(Self { - union_theta: table.theta(), + result_state: UnionResultState::Empty, table, policy, }) @@ -67,30 +73,30 @@ where S: EntrySketch, P: UnionMergePolicy, { - let SketchScalars { - seed_hash, - theta, - empty, - ordered, - .. - } = sketch.scalars(); - if empty { + let metadata = sketch.metadata(); + if metadata.is_empty() { return Ok(()); } check_seed_hash( self.table.seed_hash(), - seed_hash, + metadata.seed_hash(), "union update", ErrorKind::InvalidArgument, )?; - self.table.set_empty(false); - self.union_theta = self.union_theta.min(theta); + let current_theta = match self.result_state { + UnionResultState::Empty => self.table.retention_theta(), + UnionResultState::NonEmpty { theta } => theta, + }; + let result_theta = current_theta.min(metadata.theta()); + self.result_state = UnionResultState::NonEmpty { + theta: result_theta, + }; for entry in sketch.entries() { let hash = entry.hash(); - if hash < self.union_theta && hash < self.table.theta() { + if hash < result_theta.get() && hash < self.table.retention_theta().get() { self.table.upsert_entry(hash, |existing| match existing { Some(existing) => { self.policy.merge(existing, entry); @@ -98,68 +104,60 @@ where } None => Some(entry), }); - } else if ordered { + } else if metadata.is_ordered() { break; } } - self.union_theta = self.union_theta.min(self.table.theta()); + self.result_state = UnionResultState::NonEmpty { + theta: result_theta.min(self.table.retention_theta()), + }; Ok(()) } - /// Return the current compact-union state as compact-sketch parts. - pub fn to_compact_parts(&self, ordered: bool) -> CompactSketchParts + /// Returns the union as canonical compact-sketch state. + pub fn to_compact_sketch_state(&self, ordered: bool) -> CompactSketchState where E: Clone, { - let seed_hash = self.table.seed_hash(); - - if self.table.is_empty() { - return CompactSketchParts { - entries: vec![], - theta: self.union_theta, - seed_hash, - ordered: true, - empty: true, - }; - } + let result_theta = match self.result_state { + UnionResultState::Empty => { + return CompactSketchState::empty(self.table.seed_hash()); + } + UnionResultState::NonEmpty { theta } => theta, + }; - let mut theta = self.union_theta.min(self.table.theta()); - let mut entries = if self.union_theta >= self.table.theta() { + let mut theta = result_theta.min(self.table.retention_theta()); + let mut retained_entries = if result_theta >= self.table.retention_theta() { self.table.iter_entries().cloned().collect::>() } else { self.table .iter_entries() - .filter(|entry| entry.hash() < theta) + .filter(|entry| entry.hash() < theta.get()) .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); + if retained_entries.len() > nominal_num { + let (_, kth, _) = + retained_entries.select_nth_unstable_by_key(nominal_num, |entry| entry.hash()); + theta = ThetaThreshold::new(kth.hash()); + retained_entries.truncate(nominal_num); } - let ordered = ordered || (entries.len() == 1 && theta == MAX_THETA); + let ordered = ordered || (retained_entries.len() == 1 && theta == ThetaThreshold::MAX); if ordered { - entries.sort_unstable_by_key(SketchEntry::hash); + retained_entries.sort_unstable_by_key(SketchEntry::hash); } - CompactSketchParts { - entries, - theta, - seed_hash, - ordered, - empty: false, - } + CompactSketchState::non_empty(retained_entries, theta, self.table.seed_hash(), ordered) } /// Reset the union to its initial state. pub fn reset(&mut self) { self.table.reset(); - self.union_theta = self.table.theta(); + self.result_state = UnionResultState::Empty; } /// Returns the estimated size of the heap allocations in bytes. diff --git a/datasketches/src/thetafamily/theta/a_not_b.rs b/datasketches/src/thetafamily/theta/a_not_b.rs index 8f3b43eb..7f89b86f 100644 --- a/datasketches/src/thetafamily/theta/a_not_b.rs +++ b/datasketches/src/thetafamily/theta/a_not_b.rs @@ -91,17 +91,8 @@ impl ThetaANotB { ) -> Result { let a = a.into(); let b = b.into(); - let parts = a_not_b::compute(self.seed_hash, a, b, ordered)?; - Ok(CompactThetaSketch::from_parts( - parts - .entries - .into_iter() - .map(|entry| entry.hash()) - .collect(), - parts.theta, - parts.seed_hash, - parts.ordered, - parts.empty, - )) + let compact_state = a_not_b::compute(self.seed_hash, a, b, ordered)? + .map_retained_entries(|entry| entry.hash()); + Ok(CompactThetaSketch::from_compact_state(compact_state)) } } diff --git a/datasketches/src/thetafamily/theta/intersection.rs b/datasketches/src/thetafamily/theta/intersection.rs index 3f324cdc..8f59f9ad 100644 --- a/datasketches/src/thetafamily/theta/intersection.rs +++ b/datasketches/src/thetafamily/theta/intersection.rs @@ -85,20 +85,12 @@ impl ThetaIntersection { /// /// If `ordered` is `true`, retained hashes are sorted in ascending order. pub fn to_sketch(&self, ordered: bool) -> Option { - if !self.state.has_result() { - return None; - } - let parts = self.state.to_compact_parts(ordered); - Some(CompactThetaSketch::from_parts( - parts - .entries - .into_iter() - .map(|entry| entry.hash()) - .collect(), - parts.theta, - parts.seed_hash, - parts.ordered, - parts.empty, - )) + self.state + .to_compact_sketch_state(ordered) + .map(|compact_state| { + CompactThetaSketch::from_compact_state( + compact_state.map_retained_entries(|entry| entry.hash()), + ) + }) } } diff --git a/datasketches/src/thetafamily/theta/sketch.rs b/datasketches/src/thetafamily/theta/sketch.rs index 85272fd1..192ba8ee 100644 --- a/datasketches/src/thetafamily/theta/sketch.rs +++ b/datasketches/src/thetafamily/theta/sketch.rs @@ -48,7 +48,7 @@ use crate::theta::serialization::V2_PREAMBLE_ESTIMATE; use crate::theta::serialization::V2_PREAMBLE_PRECISE; use crate::thetacommon::EntrySketch; use crate::thetacommon::KeySketch; -use crate::thetacommon::SketchScalars; +use crate::thetacommon::ThetaSketchMetadata; use crate::thetacommon::binomial_bounds; use crate::thetacommon::constants::DEFAULT_LG_K; use crate::thetacommon::constants::FLAGS_IS_COMPACT; @@ -57,6 +57,10 @@ use crate::thetacommon::constants::FLAGS_IS_ORDERED; use crate::thetacommon::constants::FLAGS_IS_READ_ONLY; use crate::thetacommon::constants::MAX_THETA; use crate::thetacommon::hash_table::SketchHashTableIter; +use crate::thetacommon::sketch_state::CompactSketchState; +use crate::thetacommon::sketch_state::ThetaSketchState; +use crate::thetacommon::sketch_state::ThetaThreshold; +use crate::thetacommon::sketch_state::UpdateSketchState; /// Read-only view for Theta sketches. /// @@ -117,18 +121,12 @@ impl<'a> ThetaSketchView<'a> { /// Returns theta as a `u64` threshold. pub fn theta64(&self) -> u64 { - match self.0 { - ThetaSketchViewState::Mutable(sketch) => sketch.theta64(), - ThetaSketchViewState::Compact(sketch) => sketch.theta64(), - } + self.theta_sketch_state().theta().get() } - /// Returns whether the viewed sketch has not received any updates. + /// Returns `true` if the viewed sketch is empty. pub fn is_empty(&self) -> bool { - match self.0 { - ThetaSketchViewState::Mutable(sketch) => sketch.is_empty(), - ThetaSketchViewState::Compact(sketch) => sketch.is_empty(), - } + self.theta_sketch_state().is_empty() } /// Returns whether retained entries are ordered by ascending hash. @@ -146,7 +144,7 @@ impl<'a> ThetaSketchView<'a> { ThetaSketchIter::Mutable(sketch.table.iter_entries()) } ThetaSketchViewState::Compact(sketch) => { - ThetaSketchIter::Compact(sketch.entries.iter()) + ThetaSketchIter::Compact(sketch.compact_state.retained_entries().iter()) } } } @@ -158,17 +156,23 @@ impl<'a> ThetaSketchView<'a> { ThetaSketchViewState::Compact(sketch) => sketch.num_retained(), } } + + fn theta_sketch_state(&self) -> ThetaSketchState { + match self.0 { + ThetaSketchViewState::Mutable(sketch) => sketch.theta_sketch_state(), + ThetaSketchViewState::Compact(sketch) => sketch.theta_sketch_state(), + } + } } impl KeySketch for ThetaSketchView<'_> { - fn scalars(self) -> SketchScalars { - SketchScalars { - seed_hash: self.seed_hash(), - theta: self.theta64(), - empty: self.is_empty(), - ordered: self.is_ordered(), - num_retained: self.num_retained(), - } + fn metadata(self) -> ThetaSketchMetadata { + ThetaSketchMetadata::from_theta_sketch_state( + self.seed_hash(), + self.theta_sketch_state(), + self.is_ordered(), + self.num_retained(), + ) } fn hashes(self) -> impl Iterator { @@ -200,6 +204,7 @@ impl<'a> From<&'a CompactThetaSketch> for ThetaSketchView<'a> { #[derive(Debug)] pub struct ThetaSketch { table: ThetaHashTable, + update_state: UpdateSketchState, } impl ThetaSketch { @@ -228,6 +233,7 @@ impl ThetaSketch { /// assert!(sketch.estimate() >= 1.0); /// ``` pub fn update(&mut self, value: T) { + self.update_state = UpdateSketchState::Updated; self.table.try_insert(value); } @@ -247,7 +253,7 @@ impl ThetaSketch { return 0.0; } let num_retained = self.table.num_retained() as f64; - let theta = self.table.theta() as f64 / MAX_THETA as f64; + let theta = self.theta_sketch_state().theta().get() as f64 / MAX_THETA as f64; num_retained / theta } @@ -261,11 +267,7 @@ impl ThetaSketch { /// An empty sketch reports `MAX_THETA` even when it was built with a sampling probability /// below `1.0`, matching the other DataSketches implementations. pub fn theta64(&self) -> u64 { - if self.is_empty() { - MAX_THETA - } else { - self.table.theta() - } + self.theta_sketch_state().theta().get() } /// Returns the 16-bit seed hash. @@ -275,12 +277,12 @@ impl ThetaSketch { /// Returns `true` if the sketch is empty. pub fn is_empty(&self) -> bool { - self.table.is_empty() + self.theta_sketch_state().is_empty() } /// Returns `true` if the sketch is in estimation mode. pub fn is_estimation_mode(&self) -> bool { - self.theta64() < MAX_THETA + self.theta_sketch_state().is_estimation_mode() } /// Returns the number of retained entries. @@ -301,6 +303,7 @@ impl ThetaSketch { /// Resets the sketch to its empty state. pub fn reset(&mut self) { self.table.reset(); + self.update_state = UpdateSketchState::NeverUpdated; } /// Returns an iterator over retained entries. @@ -334,18 +337,11 @@ impl ThetaSketch { /// assert_eq!(compact.num_retained(), 1); /// ``` pub fn compact(&self, ordered: bool) -> CompactThetaSketch { - let parts = self.table.to_compact_parts(ordered); - CompactThetaSketch::from_parts( - parts - .entries - .into_iter() - .map(|entry| entry.hash()) - .collect(), - parts.theta, - parts.seed_hash, - parts.ordered, - parts.empty, - ) + let compact_state = self + .table + .to_compact_sketch_state(self.theta_sketch_state(), ordered) + .map_retained_entries(|entry| entry.hash()); + CompactThetaSketch::from_compact_state(compact_state) } /// Returns the approximate lower error bound for the specified number of standard deviations. @@ -425,6 +421,11 @@ impl ThetaSketch { pub fn estimated_size(&self) -> usize { size_of::() + self.table.estimated_size() } + + fn theta_sketch_state(&self) -> ThetaSketchState { + self.update_state + .theta_sketch_state(self.table.retention_theta()) + } } /// Compact (immutable) theta sketch. @@ -433,28 +434,12 @@ impl ThetaSketch { /// plus theta and a 16-bit seed hash. It can be ordered (sorted ascending) or unordered. #[derive(Clone, Debug)] pub struct CompactThetaSketch { - entries: Vec, - theta: u64, - seed_hash: u16, - ordered: bool, - empty: bool, + compact_state: CompactSketchState, } impl CompactThetaSketch { - pub(super) fn from_parts( - entries: Vec, - theta: u64, - seed_hash: u16, - ordered: bool, - empty: bool, - ) -> Self { - Self { - entries, - theta, - seed_hash, - ordered, - empty, - } + pub(super) fn from_compact_state(compact_state: CompactSketchState) -> Self { + Self { compact_state } } /// Returns a read-only view accepted by Theta set operations. @@ -468,51 +453,59 @@ impl CompactThetaSketch { return 0.0; } let num_retained = self.num_retained() as f64; - if self.theta == MAX_THETA { + if self.theta64() == MAX_THETA { return num_retained; } - let theta = self.theta as f64 / MAX_THETA as f64; + let theta = self.theta(); num_retained / theta } /// Returns theta as a fraction in `[0.0, 1.0]`. pub fn theta(&self) -> f64 { - self.theta as f64 / MAX_THETA as f64 + self.theta64() as f64 / MAX_THETA as f64 } /// Returns theta as a `u64`. pub fn theta64(&self) -> u64 { - self.theta + self.theta_sketch_state().theta().get() } /// Returns `true` if this sketch is empty. pub fn is_empty(&self) -> bool { - self.empty + self.theta_sketch_state().is_empty() } /// Returns `true` if this sketch is in estimation mode. pub fn is_estimation_mode(&self) -> bool { - self.theta < MAX_THETA + self.theta_sketch_state().is_estimation_mode() } /// Returns the number of retained entries. pub fn num_retained(&self) -> usize { - self.entries.len() + self.retained_hashes().len() } /// Returns `true` if retained entries are ordered (sorted ascending). pub fn is_ordered(&self) -> bool { - self.ordered + self.compact_state.is_ordered() } /// Returns the 16-bit seed hash. pub fn seed_hash(&self) -> u16 { - self.seed_hash + self.compact_state.seed_hash() } /// Returns an iterator over retained entries. pub fn iter(&self) -> impl Iterator + '_ { - self.entries.iter().copied().map(ThetaEntry::new) + self.retained_hashes().iter().copied().map(ThetaEntry::new) + } + + fn retained_hashes(&self) -> &[u64] { + self.compact_state.retained_entries() + } + + fn theta_sketch_state(&self) -> ThetaSketchState { + self.compact_state.theta_sketch_state() } /// Returns the approximate lower error bound for the specified number of standard deviations. @@ -543,7 +536,7 @@ impl CompactThetaSketch { if self.is_estimation_mode() { 2 } else { 1 } } else if self.is_estimation_mode() { 3 - } else if self.is_empty() || self.entries.len() == 1 { + } else if self.is_empty() || self.num_retained() == 1 { 1 } else { 2 @@ -563,14 +556,15 @@ impl CompactThetaSketch { } fn is_suitable_for_compression(&self) -> bool { - self.ordered - && !self.entries.is_empty() - && (self.entries.len() != 1 || self.is_estimation_mode()) + self.is_ordered() + && self.num_retained() != 0 + && (self.num_retained() != 1 || self.is_estimation_mode()) } /// Serializes this sketch into the uncompressed compact theta format. pub fn serialize(&self) -> Vec { - let mut bytes = SketchBytes::with_capacity(64 + self.entries.len() * 8); + let retained_hashes = self.retained_hashes(); + let mut bytes = SketchBytes::with_capacity(64 + retained_hashes.len() * 8); let pre_longs = self.preamble_longs(false); bytes.write_u8(pre_longs); @@ -589,28 +583,29 @@ impl CompactThetaSketch { } bytes.write_u8(flags); - bytes.write_u16_le(self.seed_hash); + bytes.write_u16_le(self.seed_hash()); if pre_longs > 1 { - bytes.write_u32_le(self.entries.len() as u32); + bytes.write_u32_le(retained_hashes.len() as u32); bytes.write_u32_be(0); // not used by compact sketches; match Java/C++ } if self.is_estimation_mode() { bytes.write_u64_le(self.theta64()); } - for hash in self.entries.iter() { + for hash in retained_hashes { bytes.write_u64_le(*hash); } bytes.into_bytes() } fn serialize_v4(&self) -> Vec { + let retained_hashes = self.retained_hashes(); let pre_longs = self.preamble_longs(true); - let entry_bits = Self::compute_entry_bits(&self.entries); - let num_entries_bytes = Self::num_entries_bytes(self.entries.len()); + let entry_bits = Self::compute_entry_bits(retained_hashes); + let num_entries_bytes = Self::num_entries_bytes(retained_hashes.len()); // Pre-size exactly: preamble longs (8 bytes each) + num_entries_bytes + packed bits. - let compressed_bits = entry_bits as usize * self.entries.len(); + let compressed_bits = entry_bits as usize * retained_hashes.len(); let compressed_bytes = compressed_bits.div_ceil(8); let out_bytes = (pre_longs as usize * 8) + (num_entries_bytes as usize) + compressed_bytes; let mut bytes = SketchBytes::with_capacity(out_bytes); @@ -627,12 +622,12 @@ impl CompactThetaSketch { flags |= FLAGS_IS_ORDERED; bytes.write_u8(flags); - bytes.write_u16_le(self.seed_hash); + bytes.write_u16_le(self.seed_hash()); if self.is_estimation_mode() { - bytes.write_u64_le(self.theta); + bytes.write_u64_le(self.theta64()); } - let mut n = self.entries.len() as u32; + let mut n = retained_hashes.len() as u32; for _ in 0..num_entries_bytes { bytes.write_u8((n & 0xff) as u8); n >>= 8; @@ -642,10 +637,10 @@ impl CompactThetaSketch { let mut previous = 0u64; let mut i = 0usize; let mut block = vec![0u8; entry_bits as usize]; - while i + BLOCK_WIDTH <= self.entries.len() { + while i + BLOCK_WIDTH <= retained_hashes.len() { let mut deltas = [0u64; BLOCK_WIDTH]; for j in 0..BLOCK_WIDTH { - let entry = self.entries[i + j]; + let entry = retained_hashes[i + j]; deltas[j] = entry - previous; previous = entry; } @@ -656,12 +651,12 @@ impl CompactThetaSketch { } // pack extra deltas if fewer than 8 of them left - if i < self.entries.len() { + if i < retained_hashes.len() { let mut block = vec![0u8; entry_bits as usize]; let mut packer = BitPacker::new(&mut block); - while i < self.entries.len() { - let delta = self.entries[i] - previous; - previous = self.entries[i]; + while i < retained_hashes.len() { + let delta = retained_hashes[i] - previous; + previous = retained_hashes[i]; packer.pack_value(delta, entry_bits); i += 1; } @@ -733,7 +728,7 @@ impl CompactThetaSketch { fn read_entries( cursor: &mut SketchSlice<'_>, num_entries: usize, - theta: u64, + theta: ThetaThreshold, ) -> Result, Error> { let required_bytes = num_entries .checked_mul(size_of::()) @@ -747,7 +742,7 @@ impl CompactThetaSketch { let mut entries = Vec::with_capacity(num_entries); for _ in 0..num_entries { let hash = cursor.read_u64_le().map_err(insufficient_data("entries"))?; - if hash == 0 || hash >= theta { + if hash == 0 || hash >= theta.get() { return Err(Error::deserial("corrupted: invalid retained hash value")); } entries.push(hash); @@ -755,6 +750,14 @@ impl CompactThetaSketch { Ok(entries) } + fn deserialize_theta(value: u64) -> Result { + ThetaThreshold::try_new(value).ok_or_else(|| { + Error::deserial(format!( + "corrupted: theta must be in [1, {MAX_THETA}], got {value}" + )) + }) + } + fn deserialize_v1(mut cursor: SketchSlice<'_>, expected_seed_hash: u16) -> Result { let seed_hash = expected_seed_hash; cursor.read_u8().map_err(insufficient_data(""))?; @@ -767,30 +770,23 @@ impl CompactThetaSketch { cursor .read_u32_le() .map_err(insufficient_data(""))?; - let theta = cursor - .read_u64_le() - .map_err(insufficient_data("theta_long"))?; + let theta = Self::deserialize_theta( + cursor + .read_u64_le() + .map_err(insufficient_data("theta_long"))?, + )?; - let empty = num_entries == 0 && theta == MAX_THETA; - if empty { - return Ok(Self { - entries: vec![], - theta, + if num_entries == 0 && theta == ThetaThreshold::MAX { + return Ok(Self::from_compact_state(CompactSketchState::empty( seed_hash, - ordered: true, - empty: true, - }); + ))); } let entries = Self::read_entries(&mut cursor, num_entries, theta)?; - Ok(Self { - entries, - theta, - seed_hash, - ordered: true, - empty: false, - }) + Ok(Self::from_compact_state(CompactSketchState::non_empty( + entries, theta, seed_hash, true, + ))) } fn deserialize_v2( @@ -813,13 +809,9 @@ impl CompactThetaSketch { )?; match pre_longs { - V2_PREAMBLE_EMPTY => Ok(Self { - entries: vec![], - theta: MAX_THETA, + V2_PREAMBLE_EMPTY => Ok(Self::from_compact_state(CompactSketchState::empty( seed_hash, - ordered: true, - empty: true, - }), + ))), V2_PREAMBLE_PRECISE => { let num_entries = cursor .read_u32_le() @@ -828,14 +820,18 @@ impl CompactThetaSketch { cursor .read_u32_le() .map_err(insufficient_data(""))?; - let entries = Self::read_entries(&mut cursor, num_entries, MAX_THETA)?; - Ok(Self { + let entries = Self::read_entries(&mut cursor, num_entries, ThetaThreshold::MAX)?; + if num_entries == 0 { + return Ok(Self::from_compact_state(CompactSketchState::empty( + seed_hash, + ))); + } + Ok(Self::from_compact_state(CompactSketchState::non_empty( entries, - theta: MAX_THETA, + ThetaThreshold::MAX, seed_hash, - ordered: true, - empty: num_entries == 0, - }) + true, + ))) } V2_PREAMBLE_ESTIMATE => { let num_entries = cursor @@ -845,18 +841,20 @@ impl CompactThetaSketch { cursor .read_u32_le() .map_err(insufficient_data(""))?; - let theta = cursor - .read_u64_le() - .map_err(insufficient_data("theta_long"))?; - let empty = (num_entries == 0) && (theta == MAX_THETA); + let theta = Self::deserialize_theta( + cursor + .read_u64_le() + .map_err(insufficient_data("theta_long"))?, + )?; let entries = Self::read_entries(&mut cursor, num_entries, theta)?; - Ok(Self { - entries, - theta, - seed_hash, - ordered: true, - empty, - }) + if num_entries == 0 && theta == ThetaThreshold::MAX { + return Ok(Self::from_compact_state(CompactSketchState::empty( + seed_hash, + ))); + } + Ok(Self::from_compact_state(CompactSketchState::non_empty( + entries, theta, seed_hash, true, + ))) } _ => Err(Error::invalid_preamble_longs(&[1, 2, 3], pre_longs)), } @@ -876,41 +874,42 @@ impl CompactThetaSketch { .map_err(insufficient_data("seed_hash"))?; let empty = (flags & FLAGS_IS_EMPTY) != 0; - let mut theta = MAX_THETA; - let num_entries; - let mut entries = vec![]; - if !empty { - check_seed_hash( - expected_seed_hash, + if empty { + return Ok(Self::from_compact_state(CompactSketchState::empty( seed_hash, - "deserialized CompactThetaSketch v3", - ErrorKind::InvalidData, - )?; - if pre_longs == 1 { - num_entries = 1; - } else { - num_entries = cursor - .read_u32_le() - .map_err(insufficient_data("num_entries"))?; - cursor - .read_u32_le() - .map_err(insufficient_data(""))?; - if pre_longs > 2 { - theta = cursor + ))); + } + + check_seed_hash( + expected_seed_hash, + seed_hash, + "deserialized CompactThetaSketch v3", + ErrorKind::InvalidData, + )?; + let mut theta = ThetaThreshold::MAX; + let num_entries = if pre_longs == 1 { + 1 + } else { + let num_entries = cursor + .read_u32_le() + .map_err(insufficient_data("num_entries"))?; + cursor + .read_u32_le() + .map_err(insufficient_data(""))?; + if pre_longs > 2 { + theta = Self::deserialize_theta( + cursor .read_u64_le() - .map_err(insufficient_data("theta_long"))?; - } + .map_err(insufficient_data("theta_long"))?, + )?; } - entries = Self::read_entries(&mut cursor, num_entries as usize, theta)?; - } + num_entries + }; + let entries = Self::read_entries(&mut cursor, num_entries as usize, theta)?; let ordered = (flags & FLAGS_IS_ORDERED) != 0; - Ok(Self { - entries, - theta, - seed_hash, - ordered, - empty, - }) + Ok(Self::from_compact_state(CompactSketchState::non_empty( + entries, theta, seed_hash, ordered, + ))) } fn deserialize_v4( @@ -939,11 +938,13 @@ impl CompactThetaSketch { )?; } let theta = if pre_longs > 1 { - cursor - .read_u64_le() - .map_err(insufficient_data("theta_long"))? + Self::deserialize_theta( + cursor + .read_u64_le() + .map_err(insufficient_data("theta_long"))?, + )? } else { - MAX_THETA + ThetaThreshold::MAX }; // unpack num_entries @@ -1006,25 +1007,24 @@ impl CompactThetaSketch { .checked_add(previous) .ok_or_else(|| Error::deserial("Theta entry delta overflows"))?; previous = *e; - if *e == 0 || *e >= theta { + if *e == 0 || *e >= theta.get() { return Err(Error::deserial("corrupted: invalid retained hash value")); } } let ordered = (flags & FLAGS_IS_ORDERED) != 0; - Ok(Self { - entries, - theta, - seed_hash, - ordered, - empty, - }) + let compact_state = if empty { + CompactSketchState::empty(seed_hash) + } else { + CompactSketchState::non_empty(entries, theta, seed_hash, ordered) + }; + Ok(Self::from_compact_state(compact_state)) } /// Returns the estimated size of the sketch in bytes. pub fn estimated_size(&self) -> usize { - size_of::() + self.entries.capacity() * size_of::() + size_of::() + self.compact_state.retained_entries_capacity() * size_of::() } } @@ -1128,7 +1128,10 @@ impl ThetaSketchBuilder { self.seed, )?; - Ok(ThetaSketch { table }) + Ok(ThetaSketch { + table, + update_state: UpdateSketchState::NeverUpdated, + }) } } diff --git a/datasketches/src/thetafamily/theta/union.rs b/datasketches/src/thetafamily/theta/union.rs index 72805fc4..12dff158 100644 --- a/datasketches/src/thetafamily/theta/union.rs +++ b/datasketches/src/thetafamily/theta/union.rs @@ -47,18 +47,11 @@ impl ThetaUnion { /// Returns this union as a compact sketch. pub fn to_sketch(&self, ordered: bool) -> CompactThetaSketch { - let parts = self.state.to_compact_parts(ordered); - CompactThetaSketch::from_parts( - parts - .entries - .into_iter() - .map(|entry| entry.hash()) - .collect(), - parts.theta, - parts.seed_hash, - parts.ordered, - parts.empty, - ) + let compact_state = self + .state + .to_compact_sketch_state(ordered) + .map_retained_entries(|entry| entry.hash()); + CompactThetaSketch::from_compact_state(compact_state) } /// Resets the union to its empty state. diff --git a/datasketches/src/thetafamily/tuple/a_not_b.rs b/datasketches/src/thetafamily/tuple/a_not_b.rs index 55a7cad5..58c3568b 100644 --- a/datasketches/src/thetafamily/tuple/a_not_b.rs +++ b/datasketches/src/thetafamily/tuple/a_not_b.rs @@ -100,13 +100,7 @@ impl TupleANotB { { let a = a.into(); let b = b.into(); - let parts = a_not_b::compute(self.seed_hash, a, b, ordered)?; - Ok(CompactTupleSketch::from_parts( - parts.entries, - parts.theta, - parts.seed_hash, - parts.ordered, - parts.empty, - )) + let compact_state = a_not_b::compute(self.seed_hash, a, b, ordered)?; + Ok(CompactTupleSketch::from_compact_state(compact_state)) } } diff --git a/datasketches/src/thetafamily/tuple/intersection.rs b/datasketches/src/thetafamily/tuple/intersection.rs index a38ca05b..38dfe9ed 100644 --- a/datasketches/src/thetafamily/tuple/intersection.rs +++ b/datasketches/src/thetafamily/tuple/intersection.rs @@ -153,16 +153,8 @@ where where P::Summary: Clone, { - if !self.state.has_result() { - return None; - } - let parts = self.state.to_compact_parts(ordered); - Some(CompactTupleSketch::from_parts( - parts.entries, - parts.theta, - parts.seed_hash, - parts.ordered, - parts.empty, - )) + self.state + .to_compact_sketch_state(ordered) + .map(CompactTupleSketch::from_compact_state) } } diff --git a/datasketches/src/thetafamily/tuple/sketch.rs b/datasketches/src/thetafamily/tuple/sketch.rs index beec47e7..8e7f20e4 100644 --- a/datasketches/src/thetafamily/tuple/sketch.rs +++ b/datasketches/src/thetafamily/tuple/sketch.rs @@ -39,7 +39,7 @@ use crate::hash::check_seed_hash; use crate::hash::compute_seed_hash; use crate::thetacommon::EntrySketch; use crate::thetacommon::KeySketch; -use crate::thetacommon::SketchScalars; +use crate::thetacommon::ThetaSketchMetadata; use crate::thetacommon::binomial_bounds; use crate::thetacommon::constants::DEFAULT_LG_K; use crate::thetacommon::constants::FLAGS_IS_COMPACT; @@ -48,6 +48,10 @@ use crate::thetacommon::constants::FLAGS_IS_ORDERED; use crate::thetacommon::constants::FLAGS_IS_READ_ONLY; use crate::thetacommon::constants::MAX_THETA; use crate::thetacommon::hash_table::SketchHashTableIter; +use crate::thetacommon::sketch_state::CompactSketchState; +use crate::thetacommon::sketch_state::ThetaSketchState; +use crate::thetacommon::sketch_state::ThetaThreshold; +use crate::thetacommon::sketch_state::UpdateSketchState; use crate::tuple::hash_table::TupleEntry; use crate::tuple::hash_table::TupleHashTable; use crate::tuple::policy::SummaryPolicy; @@ -81,7 +85,10 @@ pub struct TupleSketchView<'a, S>(TupleSketchViewState<'a, S>); #[derive(Debug)] enum TupleSketchViewState<'a, S> { - Mutable(&'a TupleHashTable), + Mutable { + table: &'a TupleHashTable, + update_state: UpdateSketchState, + }, Compact(&'a CompactTupleSketch), } @@ -128,40 +135,25 @@ impl<'a, S> TupleSketchView<'a, S> { /// Returns the 16-bit seed hash. pub fn seed_hash(&self) -> u16 { match self.0 { - TupleSketchViewState::Mutable(table) => table.seed_hash(), + TupleSketchViewState::Mutable { table, .. } => table.seed_hash(), TupleSketchViewState::Compact(sketch) => sketch.seed_hash(), } } /// Returns theta as a `u64` threshold. - /// - /// An empty sketch reports `MAX_THETA` even when it was built with a sampling probability - /// below `1.0`, matching the other DataSketches implementations. pub fn theta64(&self) -> u64 { - match self.0 { - TupleSketchViewState::Mutable(table) => { - if table.is_empty() { - MAX_THETA - } else { - table.theta() - } - } - TupleSketchViewState::Compact(sketch) => sketch.theta64(), - } + self.theta_sketch_state().theta().get() } - /// Returns whether the viewed sketch has not received any updates. + /// Returns `true` if the viewed sketch is empty. pub fn is_empty(&self) -> bool { - match self.0 { - TupleSketchViewState::Mutable(table) => table.is_empty(), - TupleSketchViewState::Compact(sketch) => sketch.is_empty(), - } + self.theta_sketch_state().is_empty() } /// Returns whether retained entries are ordered by ascending hash. pub fn is_ordered(&self) -> bool { match self.0 { - TupleSketchViewState::Mutable(_) => false, + TupleSketchViewState::Mutable { .. } => false, TupleSketchViewState::Compact(sketch) => sketch.is_ordered(), } } @@ -169,9 +161,11 @@ impl<'a, S> TupleSketchView<'a, S> { /// Returns an iterator over retained entries. pub fn iter(self) -> impl Iterator> + 'a { match self.0 { - TupleSketchViewState::Mutable(table) => TupleSketchIter::Mutable(table.iter_entries()), + TupleSketchViewState::Mutable { table, .. } => { + TupleSketchIter::Mutable(table.iter_entries()) + } TupleSketchViewState::Compact(sketch) => { - TupleSketchIter::Compact(sketch.entries.iter()) + TupleSketchIter::Compact(sketch.compact_state.retained_entries().iter()) } } } @@ -179,21 +173,30 @@ impl<'a, S> TupleSketchView<'a, S> { /// Returns the number of retained entries. pub fn num_retained(&self) -> usize { match self.0 { - TupleSketchViewState::Mutable(table) => table.num_retained(), + TupleSketchViewState::Mutable { table, .. } => table.num_retained(), TupleSketchViewState::Compact(sketch) => sketch.num_retained(), } } + + fn theta_sketch_state(&self) -> ThetaSketchState { + match self.0 { + TupleSketchViewState::Mutable { + table, + update_state, + } => update_state.theta_sketch_state(table.retention_theta()), + TupleSketchViewState::Compact(sketch) => sketch.theta_sketch_state(), + } + } } impl KeySketch for TupleSketchView<'_, S> { - fn scalars(self) -> SketchScalars { - SketchScalars { - seed_hash: self.seed_hash(), - theta: self.theta64(), - empty: self.is_empty(), - ordered: self.is_ordered(), - num_retained: self.num_retained(), - } + fn metadata(self) -> ThetaSketchMetadata { + ThetaSketchMetadata::from_theta_sketch_state( + self.seed_hash(), + self.theta_sketch_state(), + self.is_ordered(), + self.num_retained(), + ) } fn hashes(self) -> impl Iterator { @@ -217,7 +220,10 @@ where P: SummaryPolicy, { fn from(sketch: &'a TupleSketch

) -> Self { - Self(TupleSketchViewState::Mutable(&sketch.table)) + Self(TupleSketchViewState::Mutable { + table: &sketch.table, + update_state: sketch.update_state, + }) } } @@ -252,6 +258,7 @@ where P: SummaryPolicy, { table: TupleHashTable, + update_state: UpdateSketchState, policy: P, } @@ -284,6 +291,7 @@ where where P: SummaryUpdatePolicy, { + self.update_state = UpdateSketchState::Updated; let policy = &self.policy; self.table.try_insert(key, |existing| match existing { Some(summary) => { @@ -304,7 +312,7 @@ where return 0.0; } let num_retained = self.table.num_retained() as f64; - let theta = self.table.theta() as f64 / MAX_THETA as f64; + let theta = self.theta_sketch_state().theta().get() as f64 / MAX_THETA as f64; num_retained / theta } @@ -318,11 +326,7 @@ where /// An empty sketch reports `MAX_THETA` even when it was built with a sampling probability /// below `1.0`, matching the other DataSketches implementations. pub fn theta64(&self) -> u64 { - if self.is_empty() { - MAX_THETA - } else { - self.table.theta() - } + self.theta_sketch_state().theta().get() } /// Returns the 16-bit seed hash. @@ -332,12 +336,12 @@ where /// Returns `true` if the sketch is empty. pub fn is_empty(&self) -> bool { - self.table.is_empty() + self.theta_sketch_state().is_empty() } /// Returns `true` if the sketch is in estimation mode. pub fn is_estimation_mode(&self) -> bool { - self.theta64() < MAX_THETA + self.theta_sketch_state().is_estimation_mode() } /// Returns the number of retained entries. @@ -358,6 +362,7 @@ where /// Resets the sketch to the empty state. pub fn reset(&mut self) { self.table.reset(); + self.update_state = UpdateSketchState::NeverUpdated; } /// Returns an iterator over retained entries. @@ -392,6 +397,11 @@ where pub fn estimated_size(&self) -> usize { size_of::() + self.table.estimated_size() } + + fn theta_sketch_state(&self) -> ThetaSketchState { + self.update_state + .theta_sketch_state(self.table.retention_theta()) + } } impl

TupleSketch

@@ -416,13 +426,9 @@ where /// assert_eq!(compact.num_retained(), 1); /// ``` pub fn compact(&self, ordered: bool) -> CompactTupleSketch { - let parts = self.table.to_compact_parts(ordered); - CompactTupleSketch::from_parts( - parts.entries, - parts.theta, - parts.seed_hash, - parts.ordered, - parts.empty, + CompactTupleSketch::from_compact_state( + self.table + .to_compact_sketch_state(self.theta_sketch_state(), ordered), ) } } @@ -433,28 +439,12 @@ where /// theta and a 16-bit seed hash. It can be ordered (sorted ascending by hash) or unordered. #[derive(Clone, Debug)] pub struct CompactTupleSketch { - entries: Vec>, - theta: u64, - seed_hash: u16, - ordered: bool, - empty: bool, + compact_state: CompactSketchState>, } impl CompactTupleSketch { - pub(super) fn from_parts( - entries: Vec>, - theta: u64, - seed_hash: u16, - ordered: bool, - empty: bool, - ) -> Self { - Self { - entries, - theta, - seed_hash, - ordered, - empty, - } + pub(super) fn from_compact_state(compact_state: CompactSketchState>) -> Self { + Self { compact_state } } /// Returns a read-only view accepted by Tuple set operations. @@ -468,51 +458,59 @@ impl CompactTupleSketch { return 0.0; } let num_retained = self.num_retained() as f64; - if self.theta == MAX_THETA { + if self.theta64() == MAX_THETA { return num_retained; } - let theta = self.theta as f64 / MAX_THETA as f64; + let theta = self.theta(); num_retained / theta } /// Returns theta as a fraction (0.0 to 1.0). pub fn theta(&self) -> f64 { - self.theta as f64 / MAX_THETA as f64 + self.theta64() as f64 / MAX_THETA as f64 } /// Returns theta as `u64`. pub fn theta64(&self) -> u64 { - self.theta + self.theta_sketch_state().theta().get() } /// Returns `true` if the sketch is empty. pub fn is_empty(&self) -> bool { - self.empty + self.theta_sketch_state().is_empty() } /// Returns `true` if the sketch is in estimation mode. pub fn is_estimation_mode(&self) -> bool { - self.theta < MAX_THETA + self.theta_sketch_state().is_estimation_mode() } /// Returns the number of retained entries. pub fn num_retained(&self) -> usize { - self.entries.len() + self.retained_entries().len() } /// Returns `true` if retained entries are ordered (sorted ascending by hash). pub fn is_ordered(&self) -> bool { - self.ordered + self.compact_state.is_ordered() } /// Returns the 16-bit seed hash. pub fn seed_hash(&self) -> u16 { - self.seed_hash + self.compact_state.seed_hash() } /// Returns an iterator over retained entries. pub fn iter(&self) -> impl Iterator> + '_ { - self.entries.iter() + self.retained_entries().iter() + } + + fn retained_entries(&self) -> &[TupleEntry] { + self.compact_state.retained_entries() + } + + fn theta_sketch_state(&self) -> ThetaSketchState { + self.compact_state.theta_sketch_state() } /// Returns the approximate lower error bound given the number of standard deviations. @@ -540,13 +538,14 @@ impl CompactTupleSketch { /// Returns the estimated size of the sketch in bytes. pub fn estimated_size(&self) -> usize { - size_of::() + self.entries.capacity() * size_of::>() + size_of::() + + self.compact_state.retained_entries_capacity() * size_of::>() } fn preamble_longs(&self) -> u8 { if self.is_estimation_mode() { 3 - } else if self.is_empty() || self.entries.len() == 1 { + } else if self.is_empty() || self.num_retained() == 1 { 1 } else { 2 @@ -575,9 +574,9 @@ impl CompactTupleSketch { where S: TupleSummaryValue, { + let retained_entries = self.retained_entries(); let pre_longs = self.preamble_longs(); - let entries_size: usize = self - .entries + let entries_size: usize = retained_entries .iter() .map(|entry| 8 + entry.summary().serialize_size()) .sum(); @@ -597,17 +596,17 @@ impl CompactTupleSketch { flags |= FLAGS_IS_ORDERED; } bytes.write_u8(flags); - bytes.write_u16_le(self.seed_hash); + bytes.write_u16_le(self.seed_hash()); if pre_longs > 1 { - bytes.write_u32_le(self.entries.len() as u32); + bytes.write_u32_le(retained_entries.len() as u32); bytes.write_u32_le(0); // unused } if self.is_estimation_mode() { - bytes.write_u64_le(self.theta); + bytes.write_u64_le(self.theta64()); } - for entry in &self.entries { + for entry in retained_entries { bytes.write_u64_le(entry.hash()); entry.summary().serialize_value(&mut bytes); } @@ -671,13 +670,9 @@ impl CompactTupleSketch { let ordered = (flags & FLAGS_IS_ORDERED) != 0; if empty { - return Ok(Self::from_parts( - vec![], - MAX_THETA, + return Ok(Self::from_compact_state(CompactSketchState::empty( seed_hash, - ordered, - true, - )); + ))); } check_seed_hash( @@ -687,7 +682,7 @@ impl CompactTupleSketch { ErrorKind::InvalidData, )?; - let mut theta = MAX_THETA; + let mut theta = ThetaThreshold::MAX; let num_entries = if pre_longs == 1 { 1 } else { @@ -698,7 +693,12 @@ impl CompactTupleSketch { .read_u32_le() .map_err(insufficient_data(""))?; if pre_longs > 2 { - theta = cursor.read_u64_le().map_err(insufficient_data("theta"))?; + let value = cursor.read_u64_le().map_err(insufficient_data("theta"))?; + theta = ThetaThreshold::try_new(value).ok_or_else(|| { + Error::deserial(format!( + "corrupted: theta must be in [1, {MAX_THETA}], got {value}" + )) + })?; } n }; @@ -712,19 +712,24 @@ impl CompactTupleSketch { cursor.remaining().len() ))); } - let mut entries = Vec::with_capacity(num_entries); + let mut retained_entries = Vec::with_capacity(num_entries); for _ in 0..num_entries { let hash = cursor .read_u64_le() .map_err(insufficient_data("entry_hash"))?; - if hash == 0 || hash >= theta { + if hash == 0 || hash >= theta.get() { return Err(Error::deserial("corrupted: invalid retained hash value")); } let summary = S::deserialize_value(&mut cursor)?; - entries.push(TupleEntry::new(hash, summary)); + retained_entries.push(TupleEntry::new(hash, summary)); } - Ok(Self::from_parts(entries, theta, seed_hash, ordered, false)) + Ok(Self::from_compact_state(CompactSketchState::non_empty( + retained_entries, + theta, + seed_hash, + ordered, + ))) } } @@ -828,6 +833,7 @@ where self.sampling_probability, self.seed, )?, + update_state: UpdateSketchState::NeverUpdated, policy: self.policy, }) } diff --git a/datasketches/src/thetafamily/tuple/union.rs b/datasketches/src/thetafamily/tuple/union.rs index dd44405c..df7786b7 100644 --- a/datasketches/src/thetafamily/tuple/union.rs +++ b/datasketches/src/thetafamily/tuple/union.rs @@ -104,14 +104,7 @@ where where P::Summary: Clone, { - let result = self.state.to_compact_parts(ordered); - CompactTupleSketch::from_parts( - result.entries, - result.theta, - result.seed_hash, - result.ordered, - result.empty, - ) + CompactTupleSketch::from_compact_state(self.state.to_compact_sketch_state(ordered)) } /// Resets the union to its initial empty state. diff --git a/tests-integration/tests/theta_test/intersection.rs b/tests-integration/tests/theta_test/intersection.rs index dfd8e2c8..87f545fe 100644 --- a/tests-integration/tests/theta_test/intersection.rs +++ b/tests-integration/tests/theta_test/intersection.rs @@ -79,7 +79,8 @@ fn test_update_accepts_compact_sketch() { let r = i.to_sketch(false).unwrap(); assert_eq!(r.estimate(), 0.0); - assert!(!r.is_ordered()); + assert!(r.is_empty()); + assert!(r.is_ordered()); } #[test] diff --git a/tests-integration/tests/theta_test/union.rs b/tests-integration/tests/theta_test/union.rs index 2edac4fe..5a282326 100644 --- a/tests-integration/tests/theta_test/union.rs +++ b/tests-integration/tests/theta_test/union.rs @@ -25,6 +25,7 @@ use googletest::prelude::anything; use googletest::prelude::err; use googletest::prelude::le; use googletest::prelude::near; +use tests_integration::MAX_THETA; use tests_integration::ZERO_HASH_SEED; #[test] @@ -79,6 +80,40 @@ fn test_empty_union() { assert!(!result.is_estimation_mode()); } +#[test] +fn sampled_union_uses_canonical_empty_state_before_updates_and_after_reset() { + for probability in [0.5, 0.1, 0.001] { + let mut union = ThetaUnionBuilder::default() + .sampling_probability(probability) + .build() + .unwrap(); + + let result = union.to_sketch(false); + assert!(result.is_empty()); + assert!(result.is_ordered()); + assert_eq!(result.theta64(), MAX_THETA); + assert!(!result.is_estimation_mode()); + let bytes = result.serialize(); + let restored = CompactThetaSketch::deserialize(&bytes).unwrap(); + assert_eq!(restored.serialize(), bytes); + + let mut input = ThetaSketchBuilder::default().build().unwrap(); + input.update(1u64); + union.update(&input).unwrap(); + assert!(!union.to_sketch(false).is_empty()); + + union.reset(); + let result = union.to_sketch(false); + assert!(result.is_empty()); + assert!(result.is_ordered()); + assert_eq!(result.theta64(), MAX_THETA); + assert!(!result.is_estimation_mode()); + let bytes = result.serialize(); + let restored = CompactThetaSketch::deserialize(&bytes).unwrap(); + assert_eq!(restored.serialize(), bytes); + } +} + #[test] fn test_non_empty_no_retained_keys() { let mut sketch = ThetaSketchBuilder::default() diff --git a/tests-integration/tests/tuple_test/union.rs b/tests-integration/tests/tuple_test/union.rs index c29b5426..e1891551 100644 --- a/tests-integration/tests/tuple_test/union.rs +++ b/tests-integration/tests/tuple_test/union.rs @@ -17,6 +17,7 @@ use datasketches::common::NumStdDev; use datasketches::error::ErrorKind; +use datasketches::tuple::CompactTupleSketch; use datasketches::tuple::DefaultUnionPolicy; use datasketches::tuple::SummaryCombinePolicy; use datasketches::tuple::SummaryPolicy; @@ -25,6 +26,7 @@ use googletest::assert_that; use googletest::prelude::all; use googletest::prelude::ge; use googletest::prelude::le; +use tests_integration::MAX_THETA; use tests_integration::ZERO_HASH_SEED; use crate::default_tuple_sketch_builder; @@ -99,6 +101,40 @@ fn reset_restores_the_initial_empty_state() { assert_eq!(result.estimate(), 0.0); } +#[test] +fn sampled_union_uses_canonical_empty_state_before_updates_and_after_reset() { + for probability in [0.5, 0.1, 0.001] { + let mut union = default_union_builder() + .sampling_probability(probability) + .build() + .unwrap(); + + let result = union.to_sketch(false); + assert!(result.is_empty()); + assert!(result.is_ordered()); + assert_eq!(result.theta64(), MAX_THETA); + assert!(!result.is_estimation_mode()); + let bytes = result.serialize(); + let restored = CompactTupleSketch::::deserialize(&bytes).unwrap(); + assert_eq!(restored.serialize(), bytes); + + let mut input = default_tuple_sketch_builder().build().unwrap(); + input.update(1u64, 1u64); + union.update(&input).unwrap(); + assert!(!union.to_sketch(false).is_empty()); + + union.reset(); + let result = union.to_sketch(false); + assert!(result.is_empty()); + assert!(result.is_ordered()); + assert_eq!(result.theta64(), MAX_THETA); + assert!(!result.is_estimation_mode()); + let bytes = result.serialize(); + let restored = CompactTupleSketch::::deserialize(&bytes).unwrap(); + assert_eq!(restored.serialize(), bytes); + } +} + #[test] fn non_empty_input_requires_the_union_seed() { let mut input = default_tuple_sketch_builder().seed(1).build().unwrap(); From 258a28f7262f294ad9b34babc15463932f884572 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 1 Sep 2026 00:14:40 +0800 Subject: [PATCH 3/4] refactor(theta,tuple): simplify sketch state representation --- .../src/thetafamily/common/a_not_b.rs | 103 ++++++----- .../src/thetafamily/common/hash_table.rs | 75 +++----- .../src/thetafamily/common/intersection.rs | 54 +++--- .../thetafamily/common/jaccard_similarity.rs | 111 +++++------ datasketches/src/thetafamily/common/mod.rs | 4 +- .../src/thetafamily/common/sketch_state.rs | 173 ++---------------- datasketches/src/thetafamily/common/union.rs | 62 +++---- datasketches/src/thetafamily/theta/sketch.rs | 121 ++++++------ datasketches/src/thetafamily/tuple/sketch.rs | 115 ++++++------ tests-integration/tests/theta_test/a_not_b.rs | 34 +--- tests-integration/tests/theta_test/sketch.rs | 57 +++--- tests-integration/tests/theta_test/union.rs | 60 +++--- tests-integration/tests/tuple_test/a_not_b.rs | 45 +---- tests-integration/tests/tuple_test/sketch.rs | 62 +++---- tests-integration/tests/tuple_test/union.rs | 55 ++---- 15 files changed, 435 insertions(+), 696 deletions(-) diff --git a/datasketches/src/thetafamily/common/a_not_b.rs b/datasketches/src/thetafamily/common/a_not_b.rs index 2f3bd0bd..fc0de445 100644 --- a/datasketches/src/thetafamily/common/a_not_b.rs +++ b/datasketches/src/thetafamily/common/a_not_b.rs @@ -23,8 +23,9 @@ use crate::hash::check_seed_hash; use crate::thetacommon::EntrySketch; use crate::thetacommon::KeySketch; use crate::thetacommon::SketchEntry; +use crate::thetacommon::ThetaFamilySketchMetadata; +use crate::thetacommon::constants::MAX_THETA; use crate::thetacommon::sketch_state::CompactSketchState; -use crate::thetacommon::sketch_state::ThetaThreshold; /// Computes `a and not b` for Theta-family sketch views. /// @@ -42,46 +43,46 @@ where A: EntrySketch, B: KeySketch, { - let a_metadata = a.metadata(); - // If A is empty the result is an (empty) copy of A. As with the union and intersection, an // empty input carries no keys, so its seed is not validated. - if a_metadata.is_empty() { - return Ok(compact_state_from_sketch(a, ordered)); - } + let (a_seed_hash, a_theta, a_ordered) = match a.metadata() { + ThetaFamilySketchMetadata::Empty { .. } => { + return Ok(copy_to_compact_state(a, ordered)); + } + ThetaFamilySketchMetadata::NonEmpty { + seed_hash, + theta, + ordered, + .. + } => (seed_hash, theta, ordered), + }; // A is non-empty, so its seed must be compatible. - check_seed_hash( - seed_hash, - a_metadata.seed_hash(), - "A", - ErrorKind::InvalidArgument, - )?; - - let b_metadata = b.metadata(); + check_seed_hash(seed_hash, a_seed_hash, "A", ErrorKind::InvalidArgument)?; // An empty B subtracts nothing, so the result is simply a copy of A. This also covers the // "A is non-empty but has no retained keys" state: B's seed and theta must not influence // the result. - if b_metadata.is_empty() { - return Ok(compact_state_from_sketch(a, ordered)); - } + let (b_seed_hash, b_theta, b_ordered, b_num_retained) = match b.metadata() { + ThetaFamilySketchMetadata::Empty { .. } => { + return Ok(copy_to_compact_state(a, ordered)); + } + ThetaFamilySketchMetadata::NonEmpty { + seed_hash, + theta, + ordered, + num_retained, + } => (seed_hash, theta, ordered, num_retained), + }; // B is non-empty, so its seed must be compatible. - check_seed_hash( - seed_hash, - b_metadata.seed_hash(), - "B", - ErrorKind::InvalidArgument, - )?; + check_seed_hash(seed_hash, b_seed_hash, "B", ErrorKind::InvalidArgument)?; - let theta = a_metadata.theta().min(b_metadata.theta()); + let theta = a_theta.min(b_theta); - let entries: Vec = if b_metadata.num_retained() == 0 { - a.entries() - .filter(|entry| entry.hash() < theta.get()) - .collect() - } else if a_metadata.is_ordered() && b_metadata.is_ordered() { + let entries: Vec = if b_num_retained == 0 { + a.entries().filter(|entry| entry.hash() < theta).collect() + } else if a_ordered && b_ordered { // Both inputs are sorted ascending by hash: merge-scan without a hash set. Only // B hashes below theta can exclude an A entry (A entries are all < theta), so // unexamined B entries at or above theta are harmless. @@ -89,7 +90,7 @@ where let mut entries = vec![]; for entry in a.entries() { let hash = entry.hash(); - if hash >= theta.get() { + if hash >= theta { break; } while let Some(&b_hash) = b_hashes.peek() { @@ -105,11 +106,11 @@ where } entries } else { - let mut b_keys: HashSet = HashSet::with_capacity(b_metadata.num_retained()); + let mut b_keys: HashSet = HashSet::with_capacity(b_num_retained); for hash in b.hashes() { - if hash < theta.get() { + if hash < theta { b_keys.insert(hash); - } else if b_metadata.is_ordered() { + } else if b_ordered { break; } } @@ -117,27 +118,26 @@ where let mut entries = vec![]; for entry in a.entries() { let hash = entry.hash(); - if hash < theta.get() { + if hash < theta { if !b_keys.contains(&hash) { entries.push(entry); } - } else if a_metadata.is_ordered() { + } else if a_ordered { break; } } entries }; - if entries.is_empty() && theta == ThetaThreshold::MAX { + if entries.is_empty() && theta == MAX_THETA { return Ok(CompactSketchState::empty(seed_hash)); } let mut entries = entries; - if ordered && !a_metadata.is_ordered() && entries.len() > 1 { + if ordered && !a_ordered && entries.len() > 1 { entries.sort_unstable_by_key(SketchEntry::hash); } - let out_ordered = - ordered || a_metadata.is_ordered() || (entries.len() == 1 && theta == ThetaThreshold::MAX); + let out_ordered = ordered || a_ordered || (entries.len() == 1 && theta == MAX_THETA); Ok(CompactSketchState::non_empty( entries, @@ -147,21 +147,26 @@ where )) } -fn compact_state_from_sketch(sketch: S, ordered: bool) -> CompactSketchState +fn copy_to_compact_state(sketch: S, ordered: bool) -> CompactSketchState where S: EntrySketch, { - let metadata = sketch.metadata(); - if metadata.is_empty() { - return CompactSketchState::empty(metadata.seed_hash()); - } + let (seed_hash, theta, input_ordered) = match sketch.metadata() { + ThetaFamilySketchMetadata::Empty { seed_hash } => { + return CompactSketchState::empty(seed_hash); + } + ThetaFamilySketchMetadata::NonEmpty { + seed_hash, + theta, + ordered, + .. + } => (seed_hash, theta, ordered), + }; let mut entries: Vec = sketch.entries().collect(); - if ordered && !metadata.is_ordered() && entries.len() > 1 { + if ordered && !input_ordered && entries.len() > 1 { entries.sort_unstable_by_key(SketchEntry::hash); } - let theta = metadata.theta(); - let out_ordered = - ordered || metadata.is_ordered() || (entries.len() == 1 && theta == ThetaThreshold::MAX); - CompactSketchState::non_empty(entries, theta, metadata.seed_hash(), out_ordered) + let out_ordered = ordered || input_ordered || (entries.len() == 1 && theta == MAX_THETA); + CompactSketchState::non_empty(entries, theta, seed_hash, out_ordered) } diff --git a/datasketches/src/thetafamily/common/hash_table.rs b/datasketches/src/thetafamily/common/hash_table.rs index bc71acb1..a115a2e0 100644 --- a/datasketches/src/thetafamily/common/hash_table.rs +++ b/datasketches/src/thetafamily/common/hash_table.rs @@ -27,11 +27,10 @@ use crate::thetacommon::SketchEntry; use crate::thetacommon::constants::HASH_TABLE_REBUILD_THRESHOLD; use crate::thetacommon::constants::HASH_TABLE_RESIZE_THRESHOLD; use crate::thetacommon::constants::MAX_LG_K; +use crate::thetacommon::constants::MAX_THETA; use crate::thetacommon::constants::MIN_LG_K; use crate::thetacommon::constants::STRIDE_MASK; use crate::thetacommon::sketch_state::CompactSketchState; -use crate::thetacommon::sketch_state::ThetaSketchState; -use crate::thetacommon::sketch_state::ThetaThreshold; pub struct SketchHashTableIter<'a, E>(slice::Iter<'a, Option>); @@ -70,7 +69,7 @@ pub struct SketchHashTable { // The operational threshold used to screen future updates. This is intentionally independent // of the sketch's externally visible empty state: a never-updated sketch built with p < 1.0 // must retain this sampling threshold even though its public theta is MAX_THETA. - retention_theta: ThetaThreshold, + retention_theta: u64, entries: Vec>, @@ -126,7 +125,7 @@ where pub fn for_set_operation( lg_cur_size: u8, lg_nom_size: u8, - retention_theta: ThetaThreshold, + retention_theta: u64, seed: u64, seed_hash: u16, ) -> Self { @@ -146,7 +145,7 @@ where lg_nom_size: u8, resize_factor: ResizeFactor, sampling_probability: f32, - retention_theta: ThetaThreshold, + retention_theta: u64, seed: u64, seed_hash: u16, ) -> Self { @@ -197,7 +196,7 @@ where where F: FnOnce(Option<&mut E>) -> Option, { - if hash == 0 || hash >= self.retention_theta.get() { + if hash == 0 || hash >= self.retention_theta { return false; } @@ -283,7 +282,7 @@ where } /// Returns the operational theta used to screen retained entries. - pub fn retention_theta(&self) -> ThetaThreshold { + pub fn retention_theta(&self) -> u64 { self.retention_theta } @@ -292,31 +291,22 @@ where SketchHashTableIter(self.entries.iter()) } - /// Creates canonical compact-sketch state from this table and its owning sketch state. - pub fn to_compact_sketch_state( - &self, - theta_sketch_state: ThetaSketchState, - ordered: bool, - ) -> CompactSketchState + /// Creates compact state for a sketch known by its owner to be non-empty. + pub fn to_non_empty_compact_state(&self, ordered: bool) -> CompactSketchState where E: Clone, { - match theta_sketch_state { - ThetaSketchState::Empty => { - debug_assert_eq!(self.num_retained, 0); - CompactSketchState::empty(self.seed_hash) - } - ThetaSketchState::NonEmpty { theta } => { - debug_assert_eq!(theta, self.retention_theta); - let mut retained_entries: Vec = self.iter_entries().cloned().collect(); - let ordered = - ordered || (retained_entries.len() == 1 && theta == ThetaThreshold::MAX); - if ordered && retained_entries.len() > 1 { - retained_entries.sort_unstable_by_key(SketchEntry::hash); - } - CompactSketchState::non_empty(retained_entries, theta, self.seed_hash, ordered) - } + let mut retained_entries: Vec = self.iter_entries().cloned().collect(); + let ordered = ordered || (retained_entries.len() == 1 && self.retention_theta == MAX_THETA); + if ordered && retained_entries.len() > 1 { + retained_entries.sort_unstable_by_key(SketchEntry::hash); } + CompactSketchState::non_empty( + retained_entries, + self.retention_theta, + self.seed_hash, + ordered, + ) } /// Get log2 of nominal size. @@ -335,7 +325,11 @@ where } /// Sets the operational theta used to screen retained entries. - pub fn set_retention_theta(&mut self, retention_theta: ThetaThreshold) { + pub fn set_retention_theta(&mut self, retention_theta: u64) { + assert!( + (1..=MAX_THETA).contains(&retention_theta), + "theta must be in [1, {MAX_THETA}], got {retention_theta}" + ); self.retention_theta = retention_theta; } @@ -420,7 +414,7 @@ where let (_lesser, kth, _greater) = retained.select_nth_unstable_by_key(k, |e| e.hash()); kth.hash() }; - self.retention_theta = ThetaThreshold::new(kth_hash); + self.retention_theta = kth_hash; retained.truncate(k); let size = 1 << self.lg_cur_size; @@ -464,25 +458,10 @@ pub fn starting_sub_multiple(lg_target: u8, lg_min: u8, lg_resize_factor: u8) -> } /// Computes the initial operational theta from a sampling probability. -pub fn starting_retention_theta(sampling_probability: f32) -> ThetaThreshold { +pub fn starting_retention_theta(sampling_probability: f32) -> u64 { if sampling_probability < 1.0 { - let scaled_theta = (ThetaThreshold::MAX.get() as f64 * sampling_probability as f64) as u64; - // Threshold one and zero screen the same set of usable hashes because hash zero is - // reserved. Keep the state valid when a positive f32 probability rounds below one. - ThetaThreshold::new(scaled_theta.max(1)) + (MAX_THETA as f64 * sampling_probability as f64) as u64 } else { - ThetaThreshold::MAX - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn smallest_positive_probability_has_a_valid_retention_theta() { - let theta = starting_retention_theta(f32::from_bits(1)); - - assert_eq!(theta.get(), 1); + MAX_THETA } } diff --git a/datasketches/src/thetafamily/common/intersection.rs b/datasketches/src/thetafamily/common/intersection.rs index a49ed92f..6634d209 100644 --- a/datasketches/src/thetafamily/common/intersection.rs +++ b/datasketches/src/thetafamily/common/intersection.rs @@ -21,11 +21,11 @@ use crate::hash::check_seed_hash; use crate::hash::compute_seed_hash; use crate::thetacommon::EntrySketch; use crate::thetacommon::SketchEntry; +use crate::thetacommon::ThetaFamilySketchMetadata; use crate::thetacommon::constants::HASH_TABLE_REBUILD_THRESHOLD; +use crate::thetacommon::constants::MAX_THETA; use crate::thetacommon::hash_table::SketchHashTable; use crate::thetacommon::sketch_state::CompactSketchState; -use crate::thetacommon::sketch_state::ThetaSketchState; -use crate::thetacommon::sketch_state::ThetaThreshold; /// Merges an incoming entry into an existing entry with the same hash. /// @@ -63,7 +63,7 @@ where let seed_hash = compute_seed_hash(seed, ErrorKind::InvalidArgument)?; Ok(Self { result_state: IntersectionResultState::Uninitialized, - table: SketchHashTable::for_set_operation(0, 0, ThetaThreshold::MAX, seed, seed_hash), + table: SketchHashTable::for_set_operation(0, 0, MAX_THETA, seed, seed_hash), policy, }) } @@ -78,7 +78,6 @@ where E: Clone, P: IntersectionMergePolicy, { - let metadata = sketch.metadata(); let table_without_entries = |table: &SketchHashTable, retention_theta| { SketchHashTable::for_set_operation( 0, @@ -93,20 +92,28 @@ where return Ok(()); } - if metadata.is_empty() { - self.result_state = IntersectionResultState::Empty; - self.table = table_without_entries(&self.table, ThetaThreshold::MAX); - return Ok(()); - } + let (seed_hash, input_theta, input_ordered, input_num_retained) = match sketch.metadata() { + ThetaFamilySketchMetadata::Empty { .. } => { + self.result_state = IntersectionResultState::Empty; + self.table = table_without_entries(&self.table, MAX_THETA); + return Ok(()); + } + ThetaFamilySketchMetadata::NonEmpty { + seed_hash, + theta, + ordered, + num_retained, + } => (seed_hash, theta, ordered, num_retained), + }; check_seed_hash( self.table.seed_hash(), - metadata.seed_hash(), + seed_hash, "intersection update", ErrorKind::InvalidArgument, )?; - let result_theta = self.table.retention_theta().min(metadata.theta()); + let result_theta = self.table.retention_theta().min(input_theta); self.table.set_retention_theta(result_theta); if self.result_state == IntersectionResultState::NonEmpty && self.table.num_retained() == 0 @@ -114,7 +121,7 @@ where return Ok(()); } - if metadata.num_retained() == 0 { + if input_num_retained == 0 { self.result_state = IntersectionResultState::NonEmpty; self.table = table_without_entries(&self.table, result_theta); return Ok(()); @@ -124,7 +131,7 @@ where if self.result_state == IntersectionResultState::Uninitialized { self.result_state = IntersectionResultState::NonEmpty; let lg_size = SketchHashTable::::lg_size_from_count_for_rebuild( - metadata.num_retained(), + input_num_retained, HASH_TABLE_REBUILD_THRESHOLD, ); // The retained count is at least one here, so lg_size >= 1 and lg_size - 1 below @@ -149,18 +156,18 @@ where } } // Safety check. - if self.table.num_retained() != metadata.num_retained() { + if self.table.num_retained() != input_num_retained { return Err(Error::invalid_argument( "num entries mismatch, possibly corrupted input sketch", )); } } else { - let max_matches = self.table.num_retained().min(metadata.num_retained()); + let max_matches = self.table.num_retained().min(input_num_retained); let mut matched_entries = Vec::with_capacity(max_matches); let mut count = 0; for entry in sketch.entries() { let hash = entry.hash(); - if hash < self.table.retention_theta().get() { + if hash < self.table.retention_theta() { if let Some(existing) = self.table.entry(hash) { if matched_entries.len() == max_matches { return Err(Error::invalid_argument( @@ -171,24 +178,24 @@ where self.policy.merge(&mut merged, entry); matched_entries.push(merged); } - } else if metadata.is_ordered() { + } else if input_ordered { break; // early stop for ordered sketches } count += 1; } // Safety check. - if count > metadata.num_retained() { + if count > input_num_retained { return Err(Error::invalid_argument( "more keys than expected, possibly corrupted input sketch", )); - } else if !metadata.is_ordered() && count < metadata.num_retained() { + } else if !input_ordered && count < input_num_retained { return Err(Error::invalid_argument( "fewer keys than expected, possibly corrupted input sketch", )); } if matched_entries.is_empty() { self.table = table_without_entries(&self.table, result_theta); - if result_theta == ThetaThreshold::MAX { + if result_theta == MAX_THETA { self.result_state = IntersectionResultState::Empty; } } else { @@ -242,10 +249,9 @@ where IntersectionResultState::Empty => { Some(CompactSketchState::empty(self.table.seed_hash())) } - IntersectionResultState::NonEmpty => Some(self.table.to_compact_sketch_state( - ThetaSketchState::non_empty(self.table.retention_theta()), - ordered, - )), + IntersectionResultState::NonEmpty => { + Some(self.table.to_non_empty_compact_state(ordered)) + } } } } diff --git a/datasketches/src/thetafamily/common/jaccard_similarity.rs b/datasketches/src/thetafamily/common/jaccard_similarity.rs index 80c7ed24..263cee4a 100644 --- a/datasketches/src/thetafamily/common/jaccard_similarity.rs +++ b/datasketches/src/thetafamily/common/jaccard_similarity.rs @@ -23,7 +23,7 @@ use crate::hash::compute_seed_hash; use crate::thetacommon::EntrySketch; use crate::thetacommon::KeySketch; use crate::thetacommon::SketchEntry; -use crate::thetacommon::ThetaSketchMetadata; +use crate::thetacommon::ThetaFamilySketchMetadata; use crate::thetacommon::binomial_bounds; use crate::thetacommon::constants::MAX_LG_K; use crate::thetacommon::constants::MAX_THETA; @@ -31,7 +31,6 @@ use crate::thetacommon::constants::MIN_LG_K; use crate::thetacommon::intersection::IntersectionMergePolicy; use crate::thetacommon::intersection::IntersectionState; use crate::thetacommon::sketch_state::CompactSketchState; -use crate::thetacommon::sketch_state::ThetaThreshold; use crate::thetacommon::union::UnionMergePolicy; use crate::thetacommon::union::UnionState; @@ -71,11 +70,7 @@ impl JaccardSimilarity { } } - fn ratio_bounds( - union_count: u64, - intersection_count: u64, - theta: ThetaThreshold, - ) -> Result { + fn ratio_bounds(union_count: u64, intersection_count: u64, theta: u64) -> Result { if intersection_count > union_count { return Err(Error::invalid_argument(format!( "intersection count cannot exceed union count: {intersection_count} > {union_count}" @@ -89,8 +84,8 @@ impl JaccardSimilarity { }); } - let sampling_probability = theta.get() as f64 / MAX_THETA as f64; - if theta == ThetaThreshold::MAX { + let sampling_probability = theta as f64 / MAX_THETA as f64; + if theta == MAX_THETA { return Ok(Self::exact(intersection_count as f64 / union_count as f64)); } @@ -140,7 +135,7 @@ impl KeySketch for KeyEntries where S: KeySketch, { - fn metadata(self) -> ThetaSketchMetadata { + fn metadata(self) -> ThetaFamilySketchMetadata { self.0.metadata() } @@ -165,17 +160,14 @@ where A: KeySketch, B: KeySketch, { - let a_metadata = sketch_a.metadata(); - let b_metadata = sketch_b.metadata(); - if a_metadata.is_empty() && b_metadata.is_empty() { - return Ok(JaccardSimilarity::exact(1.0)); - } - if a_metadata.is_empty() || b_metadata.is_empty() { - return Ok(JaccardSimilarity::exact(0.0)); - } - - let sketch_a_state = (a_metadata.num_retained(), a_metadata.theta()); - let sketch_b_state = (b_metadata.num_retained(), b_metadata.theta()); + let (sketch_a_state, sketch_b_state) = match ( + non_empty_count_and_theta(sketch_a.metadata()), + non_empty_count_and_theta(sketch_b.metadata()), + ) { + (None, None) => return Ok(JaccardSimilarity::exact(1.0)), + (None, _) | (_, None) => return Ok(JaccardSimilarity::exact(0.0)), + (Some(a), Some(b)) => (a, b), + }; let union = compute_union(seed, sketch_a, sketch_b)?; if identical_sets(sketch_a_state, sketch_b_state, &union) { return Ok(JaccardSimilarity::exact(1.0)); @@ -187,11 +179,11 @@ where let intersection = intersection .to_compact_sketch_state(false) .expect("two intersection updates must produce a result"); - let union_theta = union.theta_sketch_state().theta(); + let union_theta = union.theta(); let intersection_count = intersection .retained_entries() .iter() - .filter(|entry| entry.hash < union_theta.get()) + .filter(|entry| entry.hash < union_theta) .count(); JaccardSimilarity::ratio_bounds( @@ -206,17 +198,14 @@ where A: KeySketch, B: KeySketch, { - let a_metadata = sketch_a.metadata(); - let b_metadata = sketch_b.metadata(); - if a_metadata.is_empty() && b_metadata.is_empty() { - return Ok(true); - } - if a_metadata.is_empty() || b_metadata.is_empty() { - return Ok(false); - } - - let sketch_a_state = (a_metadata.num_retained(), a_metadata.theta()); - let sketch_b_state = (b_metadata.num_retained(), b_metadata.theta()); + let (sketch_a_state, sketch_b_state) = match ( + non_empty_count_and_theta(sketch_a.metadata()), + non_empty_count_and_theta(sketch_b.metadata()), + ) { + (None, None) => return Ok(true), + (None, _) | (_, None) => return Ok(false), + (Some(a), Some(b)) => (a, b), + }; let union = compute_union(seed, sketch_a, sketch_b)?; Ok(identical_sets(sketch_a_state, sketch_b_state, &union)) } @@ -230,24 +219,28 @@ where A: KeySketch, B: KeySketch, { - let a_metadata = sketch_a.metadata(); - let b_metadata = sketch_b.metadata(); + let ThetaFamilySketchMetadata::NonEmpty { + seed_hash: a_seed_hash, + num_retained: a_num_retained, + .. + } = sketch_a.metadata() + else { + unreachable!("Jaccard union inputs are known to be non-empty") + }; + let ThetaFamilySketchMetadata::NonEmpty { + seed_hash: b_seed_hash, + num_retained: b_num_retained, + .. + } = sketch_b.metadata() + else { + unreachable!("Jaccard union inputs are known to be non-empty") + }; let seed_hash = compute_seed_hash(seed, ErrorKind::InvalidArgument)?; - check_seed_hash( - seed_hash, - a_metadata.seed_hash(), - "A", - ErrorKind::InvalidData, - )?; - check_seed_hash( - seed_hash, - b_metadata.seed_hash(), - "B", - ErrorKind::InvalidData, - )?; + check_seed_hash(seed_hash, a_seed_hash, "A", ErrorKind::InvalidData)?; + check_seed_hash(seed_hash, b_seed_hash, "B", ErrorKind::InvalidData)?; let mut union = UnionState::new( - union_lg_k(a_metadata.num_retained(), b_metadata.num_retained()), + union_lg_k(a_num_retained, b_num_retained), ResizeFactor::X8, 1.0, seed, @@ -263,15 +256,25 @@ where /// When the union retains no additional keys and preserves both input theta values, each input /// contains exactly the same retained key set represented by the union. fn identical_sets( - sketch_a: (usize, ThetaThreshold), - sketch_b: (usize, ThetaThreshold), + sketch_a: (usize, u64), + sketch_b: (usize, u64), union: &CompactSketchState, ) -> bool { - let union_state = union.theta_sketch_state(); union.retained_entries().len() == sketch_a.0 && union.retained_entries().len() == sketch_b.0 - && union_state.theta() == sketch_a.1 - && union_state.theta() == sketch_b.1 + && union.theta() == sketch_a.1 + && union.theta() == sketch_b.1 +} + +fn non_empty_count_and_theta(metadata: ThetaFamilySketchMetadata) -> Option<(usize, u64)> { + match metadata { + ThetaFamilySketchMetadata::Empty { .. } => None, + ThetaFamilySketchMetadata::NonEmpty { + theta, + num_retained, + .. + } => Some((num_retained, theta)), + } } fn sampling_adjuster(sampling_probability: f64) -> f64 { diff --git a/datasketches/src/thetafamily/common/mod.rs b/datasketches/src/thetafamily/common/mod.rs index 7827ae6b..6482a548 100644 --- a/datasketches/src/thetafamily/common/mod.rs +++ b/datasketches/src/thetafamily/common/mod.rs @@ -34,7 +34,7 @@ pub(super) trait SketchEntry { } pub(super) trait KeySketch: Copy { - fn metadata(self) -> ThetaSketchMetadata; + fn metadata(self) -> ThetaFamilySketchMetadata; fn hashes(self) -> impl Iterator; } @@ -45,4 +45,4 @@ pub(super) trait EntrySketch: KeySketch { fn entries(self) -> impl Iterator; } -pub(super) use self::sketch_state::ThetaSketchMetadata; +pub(super) use self::sketch_state::ThetaFamilySketchMetadata; diff --git a/datasketches/src/thetafamily/common/sketch_state.rs b/datasketches/src/thetafamily/common/sketch_state.rs index bdc93d71..d59f74f6 100644 --- a/datasketches/src/thetafamily/common/sketch_state.rs +++ b/datasketches/src/thetafamily/common/sketch_state.rs @@ -15,171 +15,31 @@ // specific language governing permissions and limitations // under the License. -use std::num::NonZeroU64; - use crate::thetacommon::constants::MAX_THETA; -/// A validated Theta-family retention threshold. -/// -/// Hash zero is reserved and hashes are made non-negative by dropping their high bit, so every -/// usable threshold is in `1..=MAX_THETA`. -#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)] -#[repr(transparent)] -pub struct ThetaThreshold(NonZeroU64); - -impl ThetaThreshold { - pub const MAX: Self = Self(NonZeroU64::new(MAX_THETA).unwrap()); - - /// Creates a threshold if `value` is in the valid Theta-family range. - pub fn try_new(value: u64) -> Option { - let value = NonZeroU64::new(value)?; - (value.get() <= MAX_THETA).then_some(Self(value)) - } - - /// Creates a threshold known by the caller to be valid. - /// - /// # Panics - /// - /// Panics if `value` is outside `1..=MAX_THETA`. - pub fn new(value: u64) -> Self { - Self::try_new(value) - .unwrap_or_else(|| panic!("theta must be in [1, {MAX_THETA}], got {value}")) - } - - pub fn get(self) -> u64 { - self.0.get() - } - - pub fn is_estimation_mode(self) -> bool { - self < Self::MAX - } -} - -/// Canonical state exposed by a Theta-family sketch. -/// -/// `NonEmpty` means that the sketch has observed or represents non-empty input. It may still -/// retain zero entries when every observed hash was screened out by theta. -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub enum ThetaSketchState { - Empty, - NonEmpty { theta: ThetaThreshold }, -} - -impl ThetaSketchState { - pub fn non_empty(theta: ThetaThreshold) -> Self { - Self::NonEmpty { theta } - } - - pub fn theta(self) -> ThetaThreshold { - match self { - Self::Empty => ThetaThreshold::MAX, - Self::NonEmpty { theta } => theta, - } - } - - pub fn is_empty(self) -> bool { - matches!(self, Self::Empty) - } - - pub fn is_estimation_mode(self) -> bool { - matches!(self, Self::NonEmpty { theta } if theta.is_estimation_mode()) - } -} - -/// Whether an update sketch has observed an update call since construction or reset. -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub enum UpdateSketchState { - NeverUpdated, - Updated, -} - -impl UpdateSketchState { - pub fn theta_sketch_state(self, retention_theta: ThetaThreshold) -> ThetaSketchState { - match self { - Self::NeverUpdated => ThetaSketchState::Empty, - Self::Updated => ThetaSketchState::non_empty(retention_theta), - } - } -} - /// Observable metadata consumed by shared Theta-family set operations. /// /// Empty sketches cannot carry a retained count, theta, or ordering claim in this representation. +/// A non-empty sketch may retain zero entries when theta screened every input. #[derive(Clone, Copy, Debug)] -pub enum ThetaSketchMetadata { +pub enum ThetaFamilySketchMetadata { Empty { seed_hash: u16, }, NonEmpty { seed_hash: u16, - theta: ThetaThreshold, + theta: u64, ordered: bool, num_retained: usize, }, } -impl ThetaSketchMetadata { - pub fn from_theta_sketch_state( - seed_hash: u16, - theta_sketch_state: ThetaSketchState, - ordered: bool, - num_retained: usize, - ) -> Self { - match theta_sketch_state { - ThetaSketchState::Empty => { - debug_assert_eq!(num_retained, 0); - Self::Empty { seed_hash } - } - ThetaSketchState::NonEmpty { theta } => Self::NonEmpty { - seed_hash, - theta, - ordered, - num_retained, - }, - } - } - - pub fn seed_hash(self) -> u16 { - match self { - Self::Empty { seed_hash } | Self::NonEmpty { seed_hash, .. } => seed_hash, - } - } - - pub fn theta_sketch_state(self) -> ThetaSketchState { - match self { - Self::Empty { .. } => ThetaSketchState::Empty, - Self::NonEmpty { theta, .. } => ThetaSketchState::non_empty(theta), - } - } - - pub fn is_empty(self) -> bool { - matches!(self, Self::Empty { .. }) - } - - pub fn theta(self) -> ThetaThreshold { - self.theta_sketch_state().theta() - } - - pub fn is_ordered(self) -> bool { - match self { - Self::Empty { .. } => true, - Self::NonEmpty { ordered, .. } => ordered, - } - } - - pub fn num_retained(self) -> usize { - match self { - Self::Empty { .. } => 0, - Self::NonEmpty { num_retained, .. } => num_retained, - } - } -} - /// Canonical in-memory state for a compact Theta-family sketch. /// /// The empty variant deliberately has neither retained entries nor theta. This makes the only /// representable empty state use `MAX_THETA`, report exact mode, and serialize through the -/// canonical empty-image path. +/// canonical empty-image path. The non-empty variant may contain no retained entries after theta +/// screening. #[derive(Clone, Debug)] pub enum CompactSketchState { Empty { @@ -187,7 +47,7 @@ pub enum CompactSketchState { }, NonEmpty { retained_entries: Vec, - theta: ThetaThreshold, + theta: u64, seed_hash: u16, ordered: bool, }, @@ -198,12 +58,7 @@ impl CompactSketchState { Self::Empty { seed_hash } } - pub fn non_empty( - retained_entries: Vec, - theta: ThetaThreshold, - seed_hash: u16, - ordered: bool, - ) -> Self { + pub fn non_empty(retained_entries: Vec, theta: u64, seed_hash: u16, ordered: bool) -> Self { Self::NonEmpty { retained_entries, theta, @@ -212,13 +67,21 @@ impl CompactSketchState { } } - pub fn theta_sketch_state(&self) -> ThetaSketchState { + pub fn theta(&self) -> u64 { match self { - Self::Empty { .. } => ThetaSketchState::Empty, - Self::NonEmpty { theta, .. } => ThetaSketchState::non_empty(*theta), + Self::Empty { .. } => MAX_THETA, + Self::NonEmpty { theta, .. } => *theta, } } + pub fn is_empty(&self) -> bool { + matches!(self, Self::Empty { .. }) + } + + pub fn is_estimation_mode(&self) -> bool { + matches!(self, Self::NonEmpty { theta, .. } if *theta < MAX_THETA) + } + pub fn seed_hash(&self) -> u16 { match self { Self::Empty { seed_hash } | Self::NonEmpty { seed_hash, .. } => *seed_hash, diff --git a/datasketches/src/thetafamily/common/union.rs b/datasketches/src/thetafamily/common/union.rs index 43201a04..58b3b3aa 100644 --- a/datasketches/src/thetafamily/common/union.rs +++ b/datasketches/src/thetafamily/common/union.rs @@ -21,9 +21,10 @@ use crate::error::ErrorKind; use crate::hash::check_seed_hash; use crate::thetacommon::EntrySketch; use crate::thetacommon::SketchEntry; +use crate::thetacommon::ThetaFamilySketchMetadata; +use crate::thetacommon::constants::MAX_THETA; use crate::thetacommon::hash_table::SketchHashTable; use crate::thetacommon::sketch_state::CompactSketchState; -use crate::thetacommon::sketch_state::ThetaThreshold; /// Merges an incoming entry into an existing entry with the same hash. pub trait UnionMergePolicy { @@ -38,14 +39,8 @@ pub trait UnionMergePolicy { pub struct UnionState { table: SketchHashTable, policy: P, - result_state: UnionResultState, -} - -/// State of the value currently represented by a union operator. -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -enum UnionResultState { - Empty, - NonEmpty { theta: ThetaThreshold }, + // None until the union receives a non-empty input sketch. + result_theta: Option, } impl UnionState @@ -61,7 +56,7 @@ where ) -> Result { let table = SketchHashTable::new(lg_k, resize_factor, sampling_probability, seed)?; Ok(Self { - result_state: UnionResultState::Empty, + result_theta: None, table, policy, }) @@ -73,30 +68,30 @@ where S: EntrySketch, P: UnionMergePolicy, { - let metadata = sketch.metadata(); - if metadata.is_empty() { + let ThetaFamilySketchMetadata::NonEmpty { + seed_hash, + theta, + ordered, + .. + } = sketch.metadata() + else { return Ok(()); - } + }; check_seed_hash( self.table.seed_hash(), - metadata.seed_hash(), + seed_hash, "union update", ErrorKind::InvalidArgument, )?; - let current_theta = match self.result_state { - UnionResultState::Empty => self.table.retention_theta(), - UnionResultState::NonEmpty { theta } => theta, - }; - let result_theta = current_theta.min(metadata.theta()); - self.result_state = UnionResultState::NonEmpty { - theta: result_theta, - }; + let current_theta = self.result_theta.unwrap_or(self.table.retention_theta()); + let result_theta = current_theta.min(theta); + self.result_theta = Some(result_theta); for entry in sketch.entries() { let hash = entry.hash(); - if hash < result_theta.get() && hash < self.table.retention_theta().get() { + if hash < result_theta && hash < self.table.retention_theta() { self.table.upsert_entry(hash, |existing| match existing { Some(existing) => { self.policy.merge(existing, entry); @@ -104,13 +99,11 @@ where } None => Some(entry), }); - } else if metadata.is_ordered() { + } else if ordered { break; } } - self.result_state = UnionResultState::NonEmpty { - theta: result_theta.min(self.table.retention_theta()), - }; + self.result_theta = Some(result_theta.min(self.table.retention_theta())); Ok(()) } @@ -120,11 +113,8 @@ where where E: Clone, { - let result_theta = match self.result_state { - UnionResultState::Empty => { - return CompactSketchState::empty(self.table.seed_hash()); - } - UnionResultState::NonEmpty { theta } => theta, + let Some(result_theta) = self.result_theta else { + return CompactSketchState::empty(self.table.seed_hash()); }; let mut theta = result_theta.min(self.table.retention_theta()); @@ -133,7 +123,7 @@ where } else { self.table .iter_entries() - .filter(|entry| entry.hash() < theta.get()) + .filter(|entry| entry.hash() < theta) .cloned() .collect::>() }; @@ -142,11 +132,11 @@ where if retained_entries.len() > nominal_num { let (_, kth, _) = retained_entries.select_nth_unstable_by_key(nominal_num, |entry| entry.hash()); - theta = ThetaThreshold::new(kth.hash()); + theta = kth.hash(); retained_entries.truncate(nominal_num); } - let ordered = ordered || (retained_entries.len() == 1 && theta == ThetaThreshold::MAX); + let ordered = ordered || (retained_entries.len() == 1 && theta == MAX_THETA); if ordered { retained_entries.sort_unstable_by_key(SketchEntry::hash); } @@ -157,7 +147,7 @@ where /// Reset the union to its initial state. pub fn reset(&mut self) { self.table.reset(); - self.result_state = UnionResultState::Empty; + self.result_theta = None; } /// Returns the estimated size of the heap allocations in bytes. diff --git a/datasketches/src/thetafamily/theta/sketch.rs b/datasketches/src/thetafamily/theta/sketch.rs index 192ba8ee..cbd3b5c4 100644 --- a/datasketches/src/thetafamily/theta/sketch.rs +++ b/datasketches/src/thetafamily/theta/sketch.rs @@ -48,7 +48,7 @@ use crate::theta::serialization::V2_PREAMBLE_ESTIMATE; use crate::theta::serialization::V2_PREAMBLE_PRECISE; use crate::thetacommon::EntrySketch; use crate::thetacommon::KeySketch; -use crate::thetacommon::ThetaSketchMetadata; +use crate::thetacommon::ThetaFamilySketchMetadata; use crate::thetacommon::binomial_bounds; use crate::thetacommon::constants::DEFAULT_LG_K; use crate::thetacommon::constants::FLAGS_IS_COMPACT; @@ -58,9 +58,6 @@ use crate::thetacommon::constants::FLAGS_IS_READ_ONLY; use crate::thetacommon::constants::MAX_THETA; use crate::thetacommon::hash_table::SketchHashTableIter; use crate::thetacommon::sketch_state::CompactSketchState; -use crate::thetacommon::sketch_state::ThetaSketchState; -use crate::thetacommon::sketch_state::ThetaThreshold; -use crate::thetacommon::sketch_state::UpdateSketchState; /// Read-only view for Theta sketches. /// @@ -121,12 +118,18 @@ impl<'a> ThetaSketchView<'a> { /// Returns theta as a `u64` threshold. pub fn theta64(&self) -> u64 { - self.theta_sketch_state().theta().get() + match self.0 { + ThetaSketchViewState::Mutable(sketch) => sketch.theta64(), + ThetaSketchViewState::Compact(sketch) => sketch.theta64(), + } } /// Returns `true` if the viewed sketch is empty. pub fn is_empty(&self) -> bool { - self.theta_sketch_state().is_empty() + match self.0 { + ThetaSketchViewState::Mutable(sketch) => sketch.is_empty(), + ThetaSketchViewState::Compact(sketch) => sketch.is_empty(), + } } /// Returns whether retained entries are ordered by ascending hash. @@ -156,23 +159,22 @@ impl<'a> ThetaSketchView<'a> { ThetaSketchViewState::Compact(sketch) => sketch.num_retained(), } } - - fn theta_sketch_state(&self) -> ThetaSketchState { - match self.0 { - ThetaSketchViewState::Mutable(sketch) => sketch.theta_sketch_state(), - ThetaSketchViewState::Compact(sketch) => sketch.theta_sketch_state(), - } - } } impl KeySketch for ThetaSketchView<'_> { - fn metadata(self) -> ThetaSketchMetadata { - ThetaSketchMetadata::from_theta_sketch_state( - self.seed_hash(), - self.theta_sketch_state(), - self.is_ordered(), - self.num_retained(), - ) + fn metadata(self) -> ThetaFamilySketchMetadata { + if self.is_empty() { + ThetaFamilySketchMetadata::Empty { + seed_hash: self.seed_hash(), + } + } else { + ThetaFamilySketchMetadata::NonEmpty { + seed_hash: self.seed_hash(), + theta: self.theta64(), + ordered: self.is_ordered(), + num_retained: self.num_retained(), + } + } } fn hashes(self) -> impl Iterator { @@ -204,7 +206,8 @@ impl<'a> From<&'a CompactThetaSketch> for ThetaSketchView<'a> { #[derive(Debug)] pub struct ThetaSketch { table: ThetaHashTable, - update_state: UpdateSketchState, + // Public emptiness tracks update calls, not retained entries: theta may screen every update. + is_empty: bool, } impl ThetaSketch { @@ -233,7 +236,7 @@ impl ThetaSketch { /// assert!(sketch.estimate() >= 1.0); /// ``` pub fn update(&mut self, value: T) { - self.update_state = UpdateSketchState::Updated; + self.is_empty = false; self.table.try_insert(value); } @@ -253,7 +256,7 @@ impl ThetaSketch { return 0.0; } let num_retained = self.table.num_retained() as f64; - let theta = self.theta_sketch_state().theta().get() as f64 / MAX_THETA as f64; + let theta = self.theta64() as f64 / MAX_THETA as f64; num_retained / theta } @@ -267,7 +270,11 @@ impl ThetaSketch { /// An empty sketch reports `MAX_THETA` even when it was built with a sampling probability /// below `1.0`, matching the other DataSketches implementations. pub fn theta64(&self) -> u64 { - self.theta_sketch_state().theta().get() + if self.is_empty { + MAX_THETA + } else { + self.table.retention_theta() + } } /// Returns the 16-bit seed hash. @@ -277,12 +284,12 @@ impl ThetaSketch { /// Returns `true` if the sketch is empty. pub fn is_empty(&self) -> bool { - self.theta_sketch_state().is_empty() + self.is_empty } /// Returns `true` if the sketch is in estimation mode. pub fn is_estimation_mode(&self) -> bool { - self.theta_sketch_state().is_estimation_mode() + !self.is_empty && self.table.retention_theta() < MAX_THETA } /// Returns the number of retained entries. @@ -303,7 +310,7 @@ impl ThetaSketch { /// Resets the sketch to its empty state. pub fn reset(&mut self) { self.table.reset(); - self.update_state = UpdateSketchState::NeverUpdated; + self.is_empty = true; } /// Returns an iterator over retained entries. @@ -337,10 +344,13 @@ impl ThetaSketch { /// assert_eq!(compact.num_retained(), 1); /// ``` pub fn compact(&self, ordered: bool) -> CompactThetaSketch { - let compact_state = self - .table - .to_compact_sketch_state(self.theta_sketch_state(), ordered) - .map_retained_entries(|entry| entry.hash()); + let compact_state = if self.is_empty() { + debug_assert_eq!(self.num_retained(), 0); + CompactSketchState::empty(self.seed_hash()) + } else { + self.table.to_non_empty_compact_state(ordered) + } + .map_retained_entries(|entry| entry.hash()); CompactThetaSketch::from_compact_state(compact_state) } @@ -421,11 +431,6 @@ impl ThetaSketch { pub fn estimated_size(&self) -> usize { size_of::() + self.table.estimated_size() } - - fn theta_sketch_state(&self) -> ThetaSketchState { - self.update_state - .theta_sketch_state(self.table.retention_theta()) - } } /// Compact (immutable) theta sketch. @@ -467,17 +472,17 @@ impl CompactThetaSketch { /// Returns theta as a `u64`. pub fn theta64(&self) -> u64 { - self.theta_sketch_state().theta().get() + self.compact_state.theta() } /// Returns `true` if this sketch is empty. pub fn is_empty(&self) -> bool { - self.theta_sketch_state().is_empty() + self.compact_state.is_empty() } /// Returns `true` if this sketch is in estimation mode. pub fn is_estimation_mode(&self) -> bool { - self.theta_sketch_state().is_estimation_mode() + self.compact_state.is_estimation_mode() } /// Returns the number of retained entries. @@ -504,10 +509,6 @@ impl CompactThetaSketch { self.compact_state.retained_entries() } - fn theta_sketch_state(&self) -> ThetaSketchState { - self.compact_state.theta_sketch_state() - } - /// Returns the approximate lower error bound for the specified number of standard deviations. pub fn lower_bound(&self, num_std_dev: NumStdDev) -> f64 { if !self.is_estimation_mode() { @@ -728,7 +729,7 @@ impl CompactThetaSketch { fn read_entries( cursor: &mut SketchSlice<'_>, num_entries: usize, - theta: ThetaThreshold, + theta: u64, ) -> Result, Error> { let required_bytes = num_entries .checked_mul(size_of::()) @@ -742,7 +743,7 @@ impl CompactThetaSketch { let mut entries = Vec::with_capacity(num_entries); for _ in 0..num_entries { let hash = cursor.read_u64_le().map_err(insufficient_data("entries"))?; - if hash == 0 || hash >= theta.get() { + if hash == 0 || hash >= theta { return Err(Error::deserial("corrupted: invalid retained hash value")); } entries.push(hash); @@ -750,12 +751,13 @@ impl CompactThetaSketch { Ok(entries) } - fn deserialize_theta(value: u64) -> Result { - ThetaThreshold::try_new(value).ok_or_else(|| { - Error::deserial(format!( + fn deserialize_theta(value: u64) -> Result { + if !(1..=MAX_THETA).contains(&value) { + return Err(Error::deserial(format!( "corrupted: theta must be in [1, {MAX_THETA}], got {value}" - )) - }) + ))); + } + Ok(value) } fn deserialize_v1(mut cursor: SketchSlice<'_>, expected_seed_hash: u16) -> Result { @@ -776,7 +778,7 @@ impl CompactThetaSketch { .map_err(insufficient_data("theta_long"))?, )?; - if num_entries == 0 && theta == ThetaThreshold::MAX { + if num_entries == 0 && theta == MAX_THETA { return Ok(Self::from_compact_state(CompactSketchState::empty( seed_hash, ))); @@ -820,17 +822,14 @@ impl CompactThetaSketch { cursor .read_u32_le() .map_err(insufficient_data(""))?; - let entries = Self::read_entries(&mut cursor, num_entries, ThetaThreshold::MAX)?; + let entries = Self::read_entries(&mut cursor, num_entries, MAX_THETA)?; if num_entries == 0 { return Ok(Self::from_compact_state(CompactSketchState::empty( seed_hash, ))); } Ok(Self::from_compact_state(CompactSketchState::non_empty( - entries, - ThetaThreshold::MAX, - seed_hash, - true, + entries, MAX_THETA, seed_hash, true, ))) } V2_PREAMBLE_ESTIMATE => { @@ -847,7 +846,7 @@ impl CompactThetaSketch { .map_err(insufficient_data("theta_long"))?, )?; let entries = Self::read_entries(&mut cursor, num_entries, theta)?; - if num_entries == 0 && theta == ThetaThreshold::MAX { + if num_entries == 0 && theta == MAX_THETA { return Ok(Self::from_compact_state(CompactSketchState::empty( seed_hash, ))); @@ -886,7 +885,7 @@ impl CompactThetaSketch { "deserialized CompactThetaSketch v3", ErrorKind::InvalidData, )?; - let mut theta = ThetaThreshold::MAX; + let mut theta = MAX_THETA; let num_entries = if pre_longs == 1 { 1 } else { @@ -944,7 +943,7 @@ impl CompactThetaSketch { .map_err(insufficient_data("theta_long"))?, )? } else { - ThetaThreshold::MAX + MAX_THETA }; // unpack num_entries @@ -1007,7 +1006,7 @@ impl CompactThetaSketch { .checked_add(previous) .ok_or_else(|| Error::deserial("Theta entry delta overflows"))?; previous = *e; - if *e == 0 || *e >= theta.get() { + if *e == 0 || *e >= theta { return Err(Error::deserial("corrupted: invalid retained hash value")); } } @@ -1130,7 +1129,7 @@ impl ThetaSketchBuilder { Ok(ThetaSketch { table, - update_state: UpdateSketchState::NeverUpdated, + is_empty: true, }) } } diff --git a/datasketches/src/thetafamily/tuple/sketch.rs b/datasketches/src/thetafamily/tuple/sketch.rs index 8e7f20e4..f4d24097 100644 --- a/datasketches/src/thetafamily/tuple/sketch.rs +++ b/datasketches/src/thetafamily/tuple/sketch.rs @@ -39,7 +39,7 @@ use crate::hash::check_seed_hash; use crate::hash::compute_seed_hash; use crate::thetacommon::EntrySketch; use crate::thetacommon::KeySketch; -use crate::thetacommon::ThetaSketchMetadata; +use crate::thetacommon::ThetaFamilySketchMetadata; use crate::thetacommon::binomial_bounds; use crate::thetacommon::constants::DEFAULT_LG_K; use crate::thetacommon::constants::FLAGS_IS_COMPACT; @@ -49,9 +49,6 @@ use crate::thetacommon::constants::FLAGS_IS_READ_ONLY; use crate::thetacommon::constants::MAX_THETA; use crate::thetacommon::hash_table::SketchHashTableIter; use crate::thetacommon::sketch_state::CompactSketchState; -use crate::thetacommon::sketch_state::ThetaSketchState; -use crate::thetacommon::sketch_state::ThetaThreshold; -use crate::thetacommon::sketch_state::UpdateSketchState; use crate::tuple::hash_table::TupleEntry; use crate::tuple::hash_table::TupleHashTable; use crate::tuple::policy::SummaryPolicy; @@ -87,7 +84,7 @@ pub struct TupleSketchView<'a, S>(TupleSketchViewState<'a, S>); enum TupleSketchViewState<'a, S> { Mutable { table: &'a TupleHashTable, - update_state: UpdateSketchState, + is_empty: bool, }, Compact(&'a CompactTupleSketch), } @@ -142,12 +139,24 @@ impl<'a, S> TupleSketchView<'a, S> { /// Returns theta as a `u64` threshold. pub fn theta64(&self) -> u64 { - self.theta_sketch_state().theta().get() + match self.0 { + TupleSketchViewState::Mutable { table, is_empty } => { + if is_empty { + MAX_THETA + } else { + table.retention_theta() + } + } + TupleSketchViewState::Compact(sketch) => sketch.theta64(), + } } /// Returns `true` if the viewed sketch is empty. pub fn is_empty(&self) -> bool { - self.theta_sketch_state().is_empty() + match self.0 { + TupleSketchViewState::Mutable { is_empty, .. } => is_empty, + TupleSketchViewState::Compact(sketch) => sketch.is_empty(), + } } /// Returns whether retained entries are ordered by ascending hash. @@ -177,26 +186,22 @@ impl<'a, S> TupleSketchView<'a, S> { TupleSketchViewState::Compact(sketch) => sketch.num_retained(), } } - - fn theta_sketch_state(&self) -> ThetaSketchState { - match self.0 { - TupleSketchViewState::Mutable { - table, - update_state, - } => update_state.theta_sketch_state(table.retention_theta()), - TupleSketchViewState::Compact(sketch) => sketch.theta_sketch_state(), - } - } } impl KeySketch for TupleSketchView<'_, S> { - fn metadata(self) -> ThetaSketchMetadata { - ThetaSketchMetadata::from_theta_sketch_state( - self.seed_hash(), - self.theta_sketch_state(), - self.is_ordered(), - self.num_retained(), - ) + fn metadata(self) -> ThetaFamilySketchMetadata { + if self.is_empty() { + ThetaFamilySketchMetadata::Empty { + seed_hash: self.seed_hash(), + } + } else { + ThetaFamilySketchMetadata::NonEmpty { + seed_hash: self.seed_hash(), + theta: self.theta64(), + ordered: self.is_ordered(), + num_retained: self.num_retained(), + } + } } fn hashes(self) -> impl Iterator { @@ -222,7 +227,7 @@ where fn from(sketch: &'a TupleSketch

) -> Self { Self(TupleSketchViewState::Mutable { table: &sketch.table, - update_state: sketch.update_state, + is_empty: sketch.is_empty, }) } } @@ -258,7 +263,8 @@ where P: SummaryPolicy, { table: TupleHashTable, - update_state: UpdateSketchState, + // Public emptiness tracks update calls, not retained entries: theta may screen every update. + is_empty: bool, policy: P, } @@ -291,7 +297,7 @@ where where P: SummaryUpdatePolicy, { - self.update_state = UpdateSketchState::Updated; + self.is_empty = false; let policy = &self.policy; self.table.try_insert(key, |existing| match existing { Some(summary) => { @@ -312,7 +318,7 @@ where return 0.0; } let num_retained = self.table.num_retained() as f64; - let theta = self.theta_sketch_state().theta().get() as f64 / MAX_THETA as f64; + let theta = self.theta64() as f64 / MAX_THETA as f64; num_retained / theta } @@ -326,7 +332,11 @@ where /// An empty sketch reports `MAX_THETA` even when it was built with a sampling probability /// below `1.0`, matching the other DataSketches implementations. pub fn theta64(&self) -> u64 { - self.theta_sketch_state().theta().get() + if self.is_empty { + MAX_THETA + } else { + self.table.retention_theta() + } } /// Returns the 16-bit seed hash. @@ -336,12 +346,12 @@ where /// Returns `true` if the sketch is empty. pub fn is_empty(&self) -> bool { - self.theta_sketch_state().is_empty() + self.is_empty } /// Returns `true` if the sketch is in estimation mode. pub fn is_estimation_mode(&self) -> bool { - self.theta_sketch_state().is_estimation_mode() + !self.is_empty && self.table.retention_theta() < MAX_THETA } /// Returns the number of retained entries. @@ -362,7 +372,7 @@ where /// Resets the sketch to the empty state. pub fn reset(&mut self) { self.table.reset(); - self.update_state = UpdateSketchState::NeverUpdated; + self.is_empty = true; } /// Returns an iterator over retained entries. @@ -397,11 +407,6 @@ where pub fn estimated_size(&self) -> usize { size_of::() + self.table.estimated_size() } - - fn theta_sketch_state(&self) -> ThetaSketchState { - self.update_state - .theta_sketch_state(self.table.retention_theta()) - } } impl

TupleSketch

@@ -426,10 +431,13 @@ where /// assert_eq!(compact.num_retained(), 1); /// ``` pub fn compact(&self, ordered: bool) -> CompactTupleSketch { - CompactTupleSketch::from_compact_state( - self.table - .to_compact_sketch_state(self.theta_sketch_state(), ordered), - ) + let compact_state = if self.is_empty() { + debug_assert_eq!(self.num_retained(), 0); + CompactSketchState::empty(self.seed_hash()) + } else { + self.table.to_non_empty_compact_state(ordered) + }; + CompactTupleSketch::from_compact_state(compact_state) } } @@ -472,17 +480,17 @@ impl CompactTupleSketch { /// Returns theta as `u64`. pub fn theta64(&self) -> u64 { - self.theta_sketch_state().theta().get() + self.compact_state.theta() } /// Returns `true` if the sketch is empty. pub fn is_empty(&self) -> bool { - self.theta_sketch_state().is_empty() + self.compact_state.is_empty() } /// Returns `true` if the sketch is in estimation mode. pub fn is_estimation_mode(&self) -> bool { - self.theta_sketch_state().is_estimation_mode() + self.compact_state.is_estimation_mode() } /// Returns the number of retained entries. @@ -509,10 +517,6 @@ impl CompactTupleSketch { self.compact_state.retained_entries() } - fn theta_sketch_state(&self) -> ThetaSketchState { - self.compact_state.theta_sketch_state() - } - /// Returns the approximate lower error bound given the number of standard deviations. pub fn lower_bound(&self, num_std_dev: NumStdDev) -> f64 { if !self.is_estimation_mode() { @@ -682,7 +686,7 @@ impl CompactTupleSketch { ErrorKind::InvalidData, )?; - let mut theta = ThetaThreshold::MAX; + let mut theta = MAX_THETA; let num_entries = if pre_longs == 1 { 1 } else { @@ -694,11 +698,12 @@ impl CompactTupleSketch { .map_err(insufficient_data(""))?; if pre_longs > 2 { let value = cursor.read_u64_le().map_err(insufficient_data("theta"))?; - theta = ThetaThreshold::try_new(value).ok_or_else(|| { - Error::deserial(format!( + if !(1..=MAX_THETA).contains(&value) { + return Err(Error::deserial(format!( "corrupted: theta must be in [1, {MAX_THETA}], got {value}" - )) - })?; + ))); + } + theta = value; } n }; @@ -717,7 +722,7 @@ impl CompactTupleSketch { let hash = cursor .read_u64_le() .map_err(insufficient_data("entry_hash"))?; - if hash == 0 || hash >= theta.get() { + if hash == 0 || hash >= theta { return Err(Error::deserial("corrupted: invalid retained hash value")); } let summary = S::deserialize_value(&mut cursor)?; @@ -833,7 +838,7 @@ where self.sampling_probability, self.seed, )?, - update_state: UpdateSketchState::NeverUpdated, + is_empty: true, policy: self.policy, }) } diff --git a/tests-integration/tests/theta_test/a_not_b.rs b/tests-integration/tests/theta_test/a_not_b.rs index 5e62c1c9..ff8e64c0 100644 --- a/tests-integration/tests/theta_test/a_not_b.rs +++ b/tests-integration/tests/theta_test/a_not_b.rs @@ -104,7 +104,10 @@ fn test_seed_mismatch_ignored_for_empty_inputs() { #[test] fn test_empty_a_returns_empty() { - let empty = ThetaSketchBuilder::default().build().unwrap(); + let empty = ThetaSketchBuilder::default() + .sampling_probability(0.5) + .build() + .unwrap(); let b = sketch_with_range(0, 1000); let a_not_b = ThetaANotB::default(); @@ -113,6 +116,8 @@ fn test_empty_a_returns_empty() { assert!(r.is_empty()); assert_eq!(r.num_retained(), 0); assert_eq!(r.estimate(), 0.0); + assert_eq!(r.theta64(), MAX_THETA); + assert!(!r.is_estimation_mode()); } #[test] @@ -258,30 +263,3 @@ fn test_estimation_disjoint_returns_a() { assert!(r.is_estimation_mode()); assert_that!(r.estimate(), near(10000.0, 10000.0 * 0.02)); } - -#[test] -fn test_empty_sampled_inputs_produce_a_serializable_empty_result() { - for probability in [1.0, 0.5, 0.1, 0.001] { - let a = ThetaSketchBuilder::default() - .lg_k(12) - .sampling_probability(probability) - .build() - .unwrap(); - let b = ThetaSketchBuilder::default() - .lg_k(12) - .sampling_probability(probability) - .build() - .unwrap(); - - let r = ThetaANotB::default().compute(&a, &b, true).unwrap(); - - assert!(r.is_empty()); - assert_eq!(r.theta64(), MAX_THETA); - assert!(!r.is_estimation_mode()); - - let bytes = r.serialize(); - let restored = CompactThetaSketch::deserialize(&bytes).unwrap(); - assert_eq!(restored.serialize(), bytes); - assert_eq!(restored.theta64(), r.theta64()); - } -} diff --git a/tests-integration/tests/theta_test/sketch.rs b/tests-integration/tests/theta_test/sketch.rs index cf82065b..b900ab1b 100644 --- a/tests-integration/tests/theta_test/sketch.rs +++ b/tests-integration/tests/theta_test/sketch.rs @@ -18,6 +18,7 @@ use datasketches::common::NumStdDev; use datasketches::error::ErrorKind; use datasketches::hash::value::canonical_float; +use datasketches::theta::CompactThetaSketch; use datasketches::theta::ThetaSketchBuilder; use googletest::assert_that; use googletest::prelude::ge; @@ -307,41 +308,7 @@ fn test_bounds_empty_with_sampling() { } #[test] -fn test_empty_sketch_reports_max_theta_for_every_sampling_probability() { - for probability in [1.0, 0.9, 0.5, 0.1, 0.01, 0.001] { - let sketch = ThetaSketchBuilder::default() - .lg_k(12) - .sampling_probability(probability) - .build() - .unwrap(); - - assert!(sketch.is_empty()); - assert_eq!(sketch.theta64(), MAX_THETA); - assert_eq!(sketch.theta(), 1.0); - assert!(!sketch.is_estimation_mode()); - assert_eq!(sketch.theta64(), sketch.compact(true).theta64()); - } -} - -#[test] -fn test_sampling_theta_applies_once_the_sketch_is_non_empty() { - let mut sketch = ThetaSketchBuilder::default() - .lg_k(12) - .sampling_probability(0.5) - .build() - .unwrap(); - - assert_eq!(sketch.theta64(), MAX_THETA); - - sketch.update(1u64); - - assert!(!sketch.is_empty()); - assert!(sketch.is_estimation_mode()); - assert_that!(sketch.theta64(), lt(MAX_THETA)); -} - -#[test] -fn test_compact_preserves_logical_non_empty_after_screened_update() { +fn test_sampling_state_transitions_through_compaction_and_reset() { let screened_value = (0u64..) .find(|candidate| { let mut sketch = ThetaSketchBuilder::default() @@ -359,13 +326,33 @@ fn test_compact_preserves_logical_non_empty_after_screened_update() { .sampling_probability(0.5) .build() .unwrap(); + + assert!(sketch.is_empty()); + assert_eq!(sketch.theta64(), MAX_THETA); + assert!(!sketch.is_estimation_mode()); + let empty_compact = sketch.compact(false); + assert!(empty_compact.is_empty()); + assert!(empty_compact.is_ordered()); + let bytes = empty_compact.serialize(); + assert_eq!( + CompactThetaSketch::deserialize(&bytes).unwrap().serialize(), + bytes + ); + sketch.update(screened_value); assert!(!sketch.is_empty()); assert_eq!(sketch.num_retained(), 0); + assert!(sketch.is_estimation_mode()); + assert_that!(sketch.theta64(), lt(MAX_THETA)); let compact = sketch.compact(false); assert!(!compact.is_empty()); assert_eq!(compact.num_retained(), 0); assert_eq!(compact.theta64(), sketch.theta64()); + + sketch.reset(); + assert!(sketch.is_empty()); + assert_eq!(sketch.theta64(), MAX_THETA); + assert!(!sketch.is_estimation_mode()); } diff --git a/tests-integration/tests/theta_test/union.rs b/tests-integration/tests/theta_test/union.rs index 5a282326..14df5bfe 100644 --- a/tests-integration/tests/theta_test/union.rs +++ b/tests-integration/tests/theta_test/union.rs @@ -66,52 +66,36 @@ fn assert_estimate_close(sketch: &CompactThetaSketch, expected: f64, tolerance: #[test] fn test_empty_union() { - let sketch = ThetaSketchBuilder::default().build().unwrap(); - let mut union = ThetaUnionBuilder::default().build().unwrap(); + let sketch = ThetaSketchBuilder::default() + .sampling_probability(0.5) + .build() + .unwrap(); + let mut union = ThetaUnionBuilder::default() + .sampling_probability(0.5) + .build() + .unwrap(); let result = union.to_sketch(true); assert_eq!(result.num_retained(), 0); assert!(result.is_empty()); + assert!(result.is_ordered()); + assert_eq!(result.theta64(), MAX_THETA); assert!(!result.is_estimation_mode()); union.update(&sketch).unwrap(); let result = union.to_sketch(true); assert_eq!(result.num_retained(), 0); assert!(result.is_empty()); - assert!(!result.is_estimation_mode()); -} + assert_eq!(result.theta64(), MAX_THETA); -#[test] -fn sampled_union_uses_canonical_empty_state_before_updates_and_after_reset() { - for probability in [0.5, 0.1, 0.001] { - let mut union = ThetaUnionBuilder::default() - .sampling_probability(probability) - .build() - .unwrap(); - - let result = union.to_sketch(false); - assert!(result.is_empty()); - assert!(result.is_ordered()); - assert_eq!(result.theta64(), MAX_THETA); - assert!(!result.is_estimation_mode()); - let bytes = result.serialize(); - let restored = CompactThetaSketch::deserialize(&bytes).unwrap(); - assert_eq!(restored.serialize(), bytes); - - let mut input = ThetaSketchBuilder::default().build().unwrap(); - input.update(1u64); - union.update(&input).unwrap(); - assert!(!union.to_sketch(false).is_empty()); - - union.reset(); - let result = union.to_sketch(false); - assert!(result.is_empty()); - assert!(result.is_ordered()); - assert_eq!(result.theta64(), MAX_THETA); - assert!(!result.is_estimation_mode()); - let bytes = result.serialize(); - let restored = CompactThetaSketch::deserialize(&bytes).unwrap(); - assert_eq!(restored.serialize(), bytes); - } + let mut input = ThetaSketchBuilder::default().build().unwrap(); + input.update(1u64); + union.update(&input).unwrap(); + union.reset(); + let result = union.to_sketch(false); + assert!(result.is_empty()); + assert!(result.is_ordered()); + assert_eq!(result.theta64(), MAX_THETA); + assert!(!result.is_estimation_mode()); } #[test] @@ -750,9 +734,9 @@ fn test_corner_case_union_states() { #[test] fn test_union_estimated_size() { let mut union = ThetaUnionBuilder::default().build().unwrap(); - assert_eq!(union.estimated_size(), 1096); + assert_eq!(union.estimated_size(), 1104); let sketch = sketch_with_range(12, 0, 1000); union.update(&sketch).unwrap(); - assert_eq!(union.estimated_size(), 65608); + assert_eq!(union.estimated_size(), 65616); } diff --git a/tests-integration/tests/tuple_test/a_not_b.rs b/tests-integration/tests/tuple_test/a_not_b.rs index eb9e9840..60bd308b 100644 --- a/tests-integration/tests/tuple_test/a_not_b.rs +++ b/tests-integration/tests/tuple_test/a_not_b.rs @@ -109,15 +109,18 @@ fn input_and_result_ordering_preserve_entries() { #[test] fn empty_inputs_do_not_impose_a_seed() { - let empty_other_seed = default_tuple_sketch_builder().seed(2).build().unwrap(); + let empty_other_seed = default_tuple_sketch_builder() + .sampling_probability(0.5) + .seed(2) + .build() + .unwrap(); let non_empty = tuple_sketch_with_range(0, 10); let op = TupleANotB::default(); - assert!( - op.compute(&empty_other_seed, &non_empty, true) - .unwrap() - .is_empty() - ); + let result = op.compute(&empty_other_seed, &non_empty, true).unwrap(); + assert!(result.is_empty()); + assert_eq!(result.theta64(), MAX_THETA); + assert!(!result.is_estimation_mode()); assert_eq!( op.compute(&non_empty, &empty_other_seed, true) .unwrap() @@ -186,33 +189,3 @@ fn estimation_bounds_cover_the_true_difference() { assert!(result.is_estimation_mode()); assert_that!(25_000.0, all!(ge(lower), le(upper))); } - -#[test] -fn empty_sampled_inputs_produce_a_serializable_empty_result() { - for probability in [1.0, 0.5, 0.1, 0.001] { - let a = default_tuple_sketch_builder() - .lg_k(12) - .sampling_probability(probability) - .build() - .unwrap(); - let b = default_tuple_sketch_builder() - .lg_k(12) - .sampling_probability(probability) - .build() - .unwrap(); - - assert_eq!(a.theta64(), MAX_THETA); - assert!(!a.is_estimation_mode()); - - let result = TupleANotB::default().compute(&a, &b, true).unwrap(); - - assert!(result.is_empty()); - assert_eq!(result.theta64(), MAX_THETA); - assert!(!result.is_estimation_mode()); - - let bytes = result.serialize(); - let restored = CompactTupleSketch::::deserialize(&bytes).unwrap(); - assert_eq!(restored.serialize(), bytes); - assert_eq!(restored.theta64(), result.theta64()); - } -} diff --git a/tests-integration/tests/tuple_test/sketch.rs b/tests-integration/tests/tuple_test/sketch.rs index 30bd74cc..194060d2 100644 --- a/tests-integration/tests/tuple_test/sketch.rs +++ b/tests-integration/tests/tuple_test/sketch.rs @@ -190,40 +190,6 @@ fn empty_sampled_sketch_has_zero_bounds() { assert_eq!(sketch.upper_bound(NumStdDev::Three), 0.0); } -#[test] -fn empty_sketch_reports_max_theta_for_every_sampling_probability() { - for probability in [1.0, 0.9, 0.5, 0.1, 0.01, 0.001] { - let sketch = default_tuple_sketch_builder() - .lg_k(12) - .sampling_probability(probability) - .build() - .unwrap(); - - assert!(sketch.is_empty()); - assert_eq!(sketch.theta64(), MAX_THETA); - assert_eq!(sketch.theta(), 1.0); - assert!(!sketch.is_estimation_mode()); - assert_eq!(sketch.theta64(), sketch.compact(true).theta64()); - } -} - -#[test] -fn sampling_theta_applies_once_the_sketch_is_non_empty() { - let mut sketch = default_tuple_sketch_builder() - .lg_k(12) - .sampling_probability(0.5) - .build() - .unwrap(); - - assert_eq!(sketch.theta64(), MAX_THETA); - - sketch.update(1u64, 1u64); - - assert!(!sketch.is_empty()); - assert!(sketch.is_estimation_mode()); - assert_that!(sketch.theta64(), lt(MAX_THETA)); -} - fn sorted_entries<'a>(entries: impl Iterator>) -> Vec<(u64, u64)> { let mut entries: Vec<_> = entries .map(|entry| (entry.hash(), *entry.summary())) @@ -266,7 +232,7 @@ fn compact_preserves_state_in_exact_and_estimation_modes() { } #[test] -fn compact_preserves_logical_non_empty_after_screened_update() { +fn sampling_state_transitions_through_compaction_and_reset() { let screened_value = (0u64..) .find(|candidate| { let mut sketch = default_tuple_sketch_builder() @@ -282,10 +248,36 @@ fn compact_preserves_logical_non_empty_after_screened_update() { .sampling_probability(0.5) .build() .unwrap(); + + assert!(sketch.is_empty()); + assert_eq!(sketch.theta64(), MAX_THETA); + assert!(!sketch.is_estimation_mode()); + let empty_compact = sketch.compact(false); + assert!(empty_compact.is_empty()); + assert!(empty_compact.is_ordered()); + let bytes = empty_compact.serialize(); + assert_eq!( + CompactTupleSketch::::deserialize(&bytes) + .unwrap() + .serialize(), + bytes + ); + sketch.update(screened_value, 1u64); + + assert!(!sketch.is_empty()); + assert_eq!(sketch.num_retained(), 0); + assert!(sketch.is_estimation_mode()); + assert_that!(sketch.theta64(), lt(MAX_THETA)); + let compact = sketch.compact(false); assert!(!compact.is_empty()); assert_eq!(compact.num_retained(), 0); assert_eq!(compact.theta64(), sketch.theta64()); + + sketch.reset(); + assert!(sketch.is_empty()); + assert_eq!(sketch.theta64(), MAX_THETA); + assert!(!sketch.is_estimation_mode()); } diff --git a/tests-integration/tests/tuple_test/union.rs b/tests-integration/tests/tuple_test/union.rs index e1891551..90b68b01 100644 --- a/tests-integration/tests/tuple_test/union.rs +++ b/tests-integration/tests/tuple_test/union.rs @@ -17,7 +17,6 @@ use datasketches::common::NumStdDev; use datasketches::error::ErrorKind; -use datasketches::tuple::CompactTupleSketch; use datasketches::tuple::DefaultUnionPolicy; use datasketches::tuple::SummaryCombinePolicy; use datasketches::tuple::SummaryPolicy; @@ -89,52 +88,28 @@ fn accepts_mutable_and_compact_inputs() { #[test] fn reset_restores_the_initial_empty_state() { let input = tuple_sketch_with_range(0, 100); - let mut union = default_union_builder().build().unwrap(); + let mut union = default_union_builder() + .sampling_probability(0.5) + .build() + .unwrap(); - assert!(union.to_sketch(true).is_empty()); + let initial = union.to_sketch(false); + assert!(initial.is_empty()); + assert!(initial.is_ordered()); + assert_eq!(initial.theta64(), MAX_THETA); + assert!(!initial.is_estimation_mode()); union.update(&input).unwrap(); assert!(!union.to_sketch(true).is_empty()); union.reset(); - let result = union.to_sketch(true); + let result = union.to_sketch(false); assert!(result.is_empty()); + assert!(result.is_ordered()); + assert_eq!(result.theta64(), MAX_THETA); + assert!(!result.is_estimation_mode()); assert_eq!(result.estimate(), 0.0); } -#[test] -fn sampled_union_uses_canonical_empty_state_before_updates_and_after_reset() { - for probability in [0.5, 0.1, 0.001] { - let mut union = default_union_builder() - .sampling_probability(probability) - .build() - .unwrap(); - - let result = union.to_sketch(false); - assert!(result.is_empty()); - assert!(result.is_ordered()); - assert_eq!(result.theta64(), MAX_THETA); - assert!(!result.is_estimation_mode()); - let bytes = result.serialize(); - let restored = CompactTupleSketch::::deserialize(&bytes).unwrap(); - assert_eq!(restored.serialize(), bytes); - - let mut input = default_tuple_sketch_builder().build().unwrap(); - input.update(1u64, 1u64); - union.update(&input).unwrap(); - assert!(!union.to_sketch(false).is_empty()); - - union.reset(); - let result = union.to_sketch(false); - assert!(result.is_empty()); - assert!(result.is_ordered()); - assert_eq!(result.theta64(), MAX_THETA); - assert!(!result.is_estimation_mode()); - let bytes = result.serialize(); - let restored = CompactTupleSketch::::deserialize(&bytes).unwrap(); - assert_eq!(restored.serialize(), bytes); - } -} - #[test] fn non_empty_input_requires_the_union_seed() { let mut input = default_tuple_sketch_builder().seed(1).build().unwrap(); @@ -211,9 +186,9 @@ fn estimation_bounds_cover_the_true_union() { #[test] fn union_estimated_size_grows_with_updates() { let mut union = default_union_builder().build().unwrap(); - assert_eq!(union.estimated_size(), 2120); + assert_eq!(union.estimated_size(), 2128); let sketch = tuple_sketch_with_range(0, 1000); union.update(&sketch).unwrap(); - assert_eq!(union.estimated_size(), 131144); + assert_eq!(union.estimated_size(), 131152); } From fda159c9d19368de4f987c4eeb78878f452a34cd Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 1 Sep 2026 01:24:57 +0800 Subject: [PATCH 4/4] refactor(theta,tuple): remove internal metadata re-export --- datasketches/src/thetafamily/common/a_not_b.rs | 2 +- datasketches/src/thetafamily/common/intersection.rs | 2 +- datasketches/src/thetafamily/common/jaccard_similarity.rs | 2 +- datasketches/src/thetafamily/common/mod.rs | 4 +--- datasketches/src/thetafamily/common/union.rs | 2 +- datasketches/src/thetafamily/theta/sketch.rs | 2 +- datasketches/src/thetafamily/tuple/sketch.rs | 2 +- 7 files changed, 7 insertions(+), 9 deletions(-) diff --git a/datasketches/src/thetafamily/common/a_not_b.rs b/datasketches/src/thetafamily/common/a_not_b.rs index fc0de445..d1caeea5 100644 --- a/datasketches/src/thetafamily/common/a_not_b.rs +++ b/datasketches/src/thetafamily/common/a_not_b.rs @@ -23,9 +23,9 @@ use crate::hash::check_seed_hash; use crate::thetacommon::EntrySketch; use crate::thetacommon::KeySketch; use crate::thetacommon::SketchEntry; -use crate::thetacommon::ThetaFamilySketchMetadata; use crate::thetacommon::constants::MAX_THETA; use crate::thetacommon::sketch_state::CompactSketchState; +use crate::thetacommon::sketch_state::ThetaFamilySketchMetadata; /// Computes `a and not b` for Theta-family sketch views. /// diff --git a/datasketches/src/thetafamily/common/intersection.rs b/datasketches/src/thetafamily/common/intersection.rs index 6634d209..486b2cb7 100644 --- a/datasketches/src/thetafamily/common/intersection.rs +++ b/datasketches/src/thetafamily/common/intersection.rs @@ -21,11 +21,11 @@ use crate::hash::check_seed_hash; use crate::hash::compute_seed_hash; use crate::thetacommon::EntrySketch; use crate::thetacommon::SketchEntry; -use crate::thetacommon::ThetaFamilySketchMetadata; use crate::thetacommon::constants::HASH_TABLE_REBUILD_THRESHOLD; use crate::thetacommon::constants::MAX_THETA; use crate::thetacommon::hash_table::SketchHashTable; use crate::thetacommon::sketch_state::CompactSketchState; +use crate::thetacommon::sketch_state::ThetaFamilySketchMetadata; /// Merges an incoming entry into an existing entry with the same hash. /// diff --git a/datasketches/src/thetafamily/common/jaccard_similarity.rs b/datasketches/src/thetafamily/common/jaccard_similarity.rs index 263cee4a..ec67ba02 100644 --- a/datasketches/src/thetafamily/common/jaccard_similarity.rs +++ b/datasketches/src/thetafamily/common/jaccard_similarity.rs @@ -23,7 +23,6 @@ use crate::hash::compute_seed_hash; use crate::thetacommon::EntrySketch; use crate::thetacommon::KeySketch; use crate::thetacommon::SketchEntry; -use crate::thetacommon::ThetaFamilySketchMetadata; use crate::thetacommon::binomial_bounds; use crate::thetacommon::constants::MAX_LG_K; use crate::thetacommon::constants::MAX_THETA; @@ -31,6 +30,7 @@ use crate::thetacommon::constants::MIN_LG_K; use crate::thetacommon::intersection::IntersectionMergePolicy; use crate::thetacommon::intersection::IntersectionState; use crate::thetacommon::sketch_state::CompactSketchState; +use crate::thetacommon::sketch_state::ThetaFamilySketchMetadata; use crate::thetacommon::union::UnionMergePolicy; use crate::thetacommon::union::UnionState; diff --git a/datasketches/src/thetafamily/common/mod.rs b/datasketches/src/thetafamily/common/mod.rs index 6482a548..dbc273ee 100644 --- a/datasketches/src/thetafamily/common/mod.rs +++ b/datasketches/src/thetafamily/common/mod.rs @@ -34,7 +34,7 @@ pub(super) trait SketchEntry { } pub(super) trait KeySketch: Copy { - fn metadata(self) -> ThetaFamilySketchMetadata; + fn metadata(self) -> sketch_state::ThetaFamilySketchMetadata; fn hashes(self) -> impl Iterator; } @@ -44,5 +44,3 @@ pub(super) trait EntrySketch: KeySketch { fn entries(self) -> impl Iterator; } - -pub(super) use self::sketch_state::ThetaFamilySketchMetadata; diff --git a/datasketches/src/thetafamily/common/union.rs b/datasketches/src/thetafamily/common/union.rs index 58b3b3aa..0131b55a 100644 --- a/datasketches/src/thetafamily/common/union.rs +++ b/datasketches/src/thetafamily/common/union.rs @@ -21,10 +21,10 @@ use crate::error::ErrorKind; use crate::hash::check_seed_hash; use crate::thetacommon::EntrySketch; use crate::thetacommon::SketchEntry; -use crate::thetacommon::ThetaFamilySketchMetadata; use crate::thetacommon::constants::MAX_THETA; use crate::thetacommon::hash_table::SketchHashTable; use crate::thetacommon::sketch_state::CompactSketchState; +use crate::thetacommon::sketch_state::ThetaFamilySketchMetadata; /// Merges an incoming entry into an existing entry with the same hash. pub trait UnionMergePolicy { diff --git a/datasketches/src/thetafamily/theta/sketch.rs b/datasketches/src/thetafamily/theta/sketch.rs index cbd3b5c4..88f002e9 100644 --- a/datasketches/src/thetafamily/theta/sketch.rs +++ b/datasketches/src/thetafamily/theta/sketch.rs @@ -48,7 +48,6 @@ use crate::theta::serialization::V2_PREAMBLE_ESTIMATE; use crate::theta::serialization::V2_PREAMBLE_PRECISE; use crate::thetacommon::EntrySketch; use crate::thetacommon::KeySketch; -use crate::thetacommon::ThetaFamilySketchMetadata; use crate::thetacommon::binomial_bounds; use crate::thetacommon::constants::DEFAULT_LG_K; use crate::thetacommon::constants::FLAGS_IS_COMPACT; @@ -58,6 +57,7 @@ use crate::thetacommon::constants::FLAGS_IS_READ_ONLY; use crate::thetacommon::constants::MAX_THETA; use crate::thetacommon::hash_table::SketchHashTableIter; use crate::thetacommon::sketch_state::CompactSketchState; +use crate::thetacommon::sketch_state::ThetaFamilySketchMetadata; /// Read-only view for Theta sketches. /// diff --git a/datasketches/src/thetafamily/tuple/sketch.rs b/datasketches/src/thetafamily/tuple/sketch.rs index f4d24097..796e17ee 100644 --- a/datasketches/src/thetafamily/tuple/sketch.rs +++ b/datasketches/src/thetafamily/tuple/sketch.rs @@ -39,7 +39,6 @@ use crate::hash::check_seed_hash; use crate::hash::compute_seed_hash; use crate::thetacommon::EntrySketch; use crate::thetacommon::KeySketch; -use crate::thetacommon::ThetaFamilySketchMetadata; use crate::thetacommon::binomial_bounds; use crate::thetacommon::constants::DEFAULT_LG_K; use crate::thetacommon::constants::FLAGS_IS_COMPACT; @@ -49,6 +48,7 @@ use crate::thetacommon::constants::FLAGS_IS_READ_ONLY; use crate::thetacommon::constants::MAX_THETA; use crate::thetacommon::hash_table::SketchHashTableIter; use crate::thetacommon::sketch_state::CompactSketchState; +use crate::thetacommon::sketch_state::ThetaFamilySketchMetadata; use crate::tuple::hash_table::TupleEntry; use crate::tuple::hash_table::TupleHashTable; use crate::tuple::policy::SummaryPolicy;