From 55ea6fdf30a47363fcfd7c634ec86a104805cc00 Mon Sep 17 00:00:00 2001 From: David Cromberge Date: Tue, 10 Nov 2026 13:24:53 +0000 Subject: [PATCH] perf(theta): gather and bucket-sort hashes when compacting update sketches Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 1 + .../src/thetafamily/common/hash_table.rs | 18 +++++++ datasketches/src/thetafamily/theta/sketch.rs | 54 +++++++++++++++++-- tests-integration/tests/theta_test/sketch.rs | 39 ++++++++++++++ 4 files changed, 109 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 36602987..6934b09b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,7 @@ All significant changes to this project will be documented in this file. * `BloomFilter::insert` is faster for streams of new items when multiple hashes are configured. * Compressed `CompactThetaSketch` serialization and deserialization are faster. * Tuple sketch serialization and deserialization are faster. +* `ThetaSketch::compact` is faster, especially for ordered results. ## v0.6.0 (2026-10-29) diff --git a/datasketches/src/thetafamily/common/hash_table.rs b/datasketches/src/thetafamily/common/hash_table.rs index f57d786f..18967dd2 100644 --- a/datasketches/src/thetafamily/common/hash_table.rs +++ b/datasketches/src/thetafamily/common/hash_table.rs @@ -307,6 +307,24 @@ where ) } + /// Get the retained hashes in table order. + #[cfg(feature = "theta")] + pub fn retained_hashes(&self) -> Vec { + if self.num_retained * 16 <= self.entries.len() { + let mut hashes = Vec::with_capacity(self.num_retained); + hashes.extend(self.entries.iter().flatten().map(SketchEntry::hash)); + return hashes; + } + let mut hashes = vec![0u64; self.num_retained + 1]; + let mut len = 0; + for slot in &self.entries { + hashes[len] = slot.as_ref().map_or(0, SketchEntry::hash); + len += usize::from(slot.is_some()); + } + hashes.truncate(len); + hashes + } + /// Get log2 of nominal size. pub fn lg_nom_size(&self) -> u8 { self.lg_nom_size diff --git a/datasketches/src/thetafamily/theta/sketch.rs b/datasketches/src/thetafamily/theta/sketch.rs index 82197dff..ae26c33b 100644 --- a/datasketches/src/thetafamily/theta/sketch.rs +++ b/datasketches/src/thetafamily/theta/sketch.rs @@ -342,9 +342,14 @@ impl ThetaSketch { 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()); + let mut hashes = self.table.retained_hashes(); + let theta = self.table.retention_theta(); + let ordered = ordered || (hashes.len() == 1 && theta == MAX_THETA); + if ordered { + sort_uniform_hashes(&mut hashes); + } + CompactSketchState::non_empty(hashes, theta, self.seed_hash(), ordered) + }; CompactThetaSketch::from_compact_state(compact_state) } @@ -1094,6 +1099,49 @@ fn unpack_deltas( Ok(entries) } +fn sort_uniform_hashes(hashes: &mut Vec) { + const MAX_BUCKET: usize = 32; + let len = hashes.len(); + let Some(&max) = hashes.iter().max() else { + return; + }; + if len < 256 { + hashes.sort_unstable(); + return; + } + let lg_buckets = len.ilog2(); + let shift = (u64::BITS - max.leading_zeros()).saturating_sub(lg_buckets); + let mut starts = vec![0u32; (1 << lg_buckets) + 1]; + for &hash in hashes.iter() { + starts[(hash >> shift) as usize + 1] += 1; + } + let mut largest = 0; + for i in 1..starts.len() { + largest = largest.max(starts[i]); + starts[i] += starts[i - 1]; + } + if largest as usize > MAX_BUCKET { + hashes.sort_unstable(); + return; + } + let mut sorted = vec![0u64; len]; + for &hash in hashes.iter() { + let next = &mut starts[(hash >> shift) as usize]; + sorted[*next as usize] = hash; + *next += 1; + } + for i in 1..len { + let hash = sorted[i]; + let mut j = i; + while j > 0 && sorted[j - 1] > hash { + sorted[j] = sorted[j - 1]; + j -= 1; + } + sorted[j] = hash; + } + *hashes = sorted; +} + #[cfg(test)] mod tests { use googletest::assert_that; diff --git a/tests-integration/tests/theta_test/sketch.rs b/tests-integration/tests/theta_test/sketch.rs index b900ab1b..3545a6cb 100644 --- a/tests-integration/tests/theta_test/sketch.rs +++ b/tests-integration/tests/theta_test/sketch.rs @@ -16,6 +16,7 @@ // under the License. use datasketches::common::NumStdDev; +use datasketches::common::ResizeFactor; use datasketches::error::ErrorKind; use datasketches::hash::value::canonical_float; use datasketches::theta::CompactThetaSketch; @@ -167,6 +168,44 @@ fn test_reset() { assert_eq!(sketch.upper_bound(NumStdDev::One), 0.0); } +#[test] +fn test_compact_ordered_sorts_retained_hashes() { + for (lg_k, n, p, rf) in [ + (12, 1, 1.0, ResizeFactor::X8), + (12, 2, 1.0, ResizeFactor::X8), + (12, 300, 1.0, ResizeFactor::X1), + (5, 300, 1.0, ResizeFactor::X8), + (8, 255, 1.0, ResizeFactor::X8), + (12, 3000, 1.0, ResizeFactor::X8), + (12, 100_000, 1.0, ResizeFactor::X8), + (12, 100_000, 0.01, ResizeFactor::X8), + (16, 1_000_000, 1.0, ResizeFactor::X8), + ] { + let mut sketch = ThetaSketchBuilder::default() + .lg_k(lg_k) + .sampling_probability(p) + .resize_factor(rf) + .build() + .unwrap(); + for i in 0..n { + sketch.update(i as u64); + } + let mut expected: Vec = sketch.iter().map(|entry| entry.hash()).collect(); + expected.sort_unstable(); + + let ordered = sketch.compact(true); + assert!(ordered.is_ordered()); + let hashes: Vec = ordered.iter().map(|entry| entry.hash()).collect(); + assert_eq!(hashes, expected); + + let unordered = sketch.compact(false); + let mut hashes: Vec = unordered.iter().map(|entry| entry.hash()).collect(); + hashes.sort_unstable(); + assert_eq!(hashes, expected); + assert_eq!(unordered.theta64(), sketch.theta64()); + } +} + #[test] fn test_iterator() { let mut sketch = ThetaSketchBuilder::default().lg_k(12).build().unwrap();