perf(theta): gather and bucket-sort hashes when compacting update sketches

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)
 

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<u64> {
+        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

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<u64>) {
+    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;

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<u64> = sketch.iter().map(|entry| entry.hash()).collect();
+        expected.sort_unstable();
+
+        let ordered = sketch.compact(true);
+        assert!(ordered.is_ordered());
+        let hashes: Vec<u64> = ordered.iter().map(|entry| entry.hash()).collect();
+        assert_eq!(hashes, expected);
+
+        let unordered = sketch.compact(false);
+        let mut hashes: Vec<u64> = 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();