perf(frequencies): skip key comparisons when rehashing

CHANGELOG.md

@@ -12,6 +12,7 @@ All significant changes to this project will be documented in this file.
 
 ### Improvements
 
+* `FrequentItemsSketch` updates are faster for streams of distinct items while the map is growing.
 * `HllUnion` is faster when adopting HLL8 inputs without downsampling or merging multiple dense inputs before querying or serializing the result.
 * `KllSketch::rank` is faster for repeated queries on compacted sketches.
 * Repeated scalar `TDigest::rank` queries are faster, including small digests and data with repeated values.

datasketches/src/frequencies/reverse_purge_item_hash_map.rs

@@ -171,12 +171,30 @@ impl<T: Eq + Hash> ReversePurgeItemHashMap<T> {
         for i in 0..old_keys.len() {
             if old_states[i] > 0 {
                 if let Some(key) = old_keys[i].take() {
-                    self.adjust_or_put_value(key, old_values[i]);
+                    self.insert_absent(key, old_values[i]);
                 }
             }
         }
     }
 
+    /// Inserts a key that is known to be absent, such as when rehashing distinct keys.
+    ///
+    /// This probes the same slots as [`Self::adjust_or_put_value`] but skips key comparisons.
+    fn insert_absent(&mut self, key: T, value: u64) {
+        let mask = self.keys.len() - 1;
+        let mut probe = (hash_item(&key) as usize) & mask;
+        let mut drift: usize = 1;
+        while self.states[probe] > 0 {
+            probe = (probe + 1) & mask;
+            drift += 1;
+            debug_assert!(drift < DRIFT_LIMIT, "drift limit exceeded");
+        }
+        self.keys[probe] = Some(key);
+        self.values[probe] = value;
+        self.states[probe] = drift as u16;
+        self.num_active += 1;
+    }
+
     /// Returns the length of the underlying arrays.
     pub fn len(&self) -> usize {
         self.keys.len()

tests-integration/tests/frequencies_test/main.rs

@@ -15,4 +15,5 @@
 // specific language governing permissions and limitations
 // under the License.
 
+mod rehash;
 mod update;

tests-integration/tests/frequencies_test/rehash.rs

@@ -0,0 +1,111 @@
+// 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.
+
+//! Growing the item map must keep every item with its exact weight.
+
+use std::hash::Hash;
+use std::hash::Hasher;
+
+use datasketches::frequencies::FrequentItemsSketch;
+
+fn weight(i: u64) -> u64 {
+    i % 5 + 1
+}
+
+// Stay below the purge threshold and check all counts at each growth boundary.
+fn check_growth<T: Eq + Hash>(max_map_size: usize, distinct: u64, item: impl Fn(u64) -> T) {
+    let mut sketch = FrequentItemsSketch::new(max_map_size).unwrap();
+    let initial_capacity = sketch.current_map_capacity();
+    let mut total_weight = 0;
+    for i in 0..distinct {
+        let capacity = sketch.current_map_capacity();
+        sketch.update_with_count(item(i), weight(i));
+        total_weight += weight(i);
+        if sketch.current_map_capacity() != capacity || i + 1 == distinct {
+            assert_eq!(sketch.total_weight(), total_weight);
+            assert_eq!(sketch.num_active_items() as u64, i + 1);
+            assert_eq!(sketch.maximum_error(), 0);
+            for j in 0..=i {
+                assert_eq!(sketch.estimate(&item(j)), weight(j), "item {j} after {i}");
+            }
+        }
+    }
+    assert!(sketch.current_map_capacity() > initial_capacity);
+}
+
+#[test]
+fn test_longs_growth_keeps_every_item() {
+    check_growth(1 << 12, 2_000, |i| (i * 7_919 + 13) as i64);
+}
+
+#[test]
+fn test_items_growth_keeps_every_item() {
+    check_growth(1 << 12, 2_000, |i| format!("item-{i}"));
+}
+
+#[derive(Debug, PartialEq, Eq)]
+struct CollidingItem(u64);
+
+impl Hash for CollidingItem {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        0u64.hash(state);
+    }
+}
+
+#[test]
+fn test_colliding_items_survive_growth_updates_and_purge() {
+    let mut sketch = FrequentItemsSketch::new(64).unwrap();
+    for i in 0..48 {
+        sketch.update_with_count(CollidingItem(i), if i % 4 == 0 { 100 } else { 1 });
+    }
+    assert_eq!(sketch.current_map_capacity(), sketch.maximum_map_capacity());
+    assert_eq!(sketch.maximum_error(), 0);
+    for i in 0..48 {
+        assert_eq!(
+            sketch.estimate(&CollidingItem(i)),
+            if i % 4 == 0 { 100 } else { 1 }
+        );
+    }
+
+    for i in (0..48).step_by(4) {
+        sketch.update_with_count(CollidingItem(i), 7);
+    }
+    assert_eq!(sketch.num_active_items(), 48);
+    assert_eq!(sketch.maximum_error(), 0);
+
+    // The new key triggers a purge with median 1. Deleting light items from
+    // the collision chain must preserve access to every displaced heavy item.
+    sketch.update(CollidingItem(48));
+    assert_eq!(sketch.maximum_error(), 1);
+    assert_eq!(sketch.num_active_items(), 12);
+    for i in 0..49 {
+        if i < 48 && i % 4 == 0 {
+            assert_eq!(sketch.estimate(&CollidingItem(i)), 107);
+            assert_eq!(sketch.lower_bound(&CollidingItem(i)), 106);
+        } else {
+            assert_eq!(sketch.estimate(&CollidingItem(i)), 0);
+            assert_eq!(sketch.upper_bound(&CollidingItem(i)), 1);
+        }
+    }
+
+    sketch.update_with_count(CollidingItem(44), 3);
+    sketch.update_with_count(CollidingItem(1), 5);
+    assert_eq!(sketch.estimate(&CollidingItem(44)), 110);
+    assert_eq!(sketch.estimate(&CollidingItem(1)), 6);
+    assert_eq!(sketch.num_active_items(), 13);
+    assert_eq!(sketch.total_weight(), 12 * 107 + 37 + 3 + 5);
+}

tests-integration/tests/serde_tests/frequencies.rs

@@ -64,6 +64,23 @@ fn test_longs_round_trip() {
     assert_eq!(restored.maximum_error(), sketch.maximum_error());
 }
 
+#[test]
+fn test_growth_survives_serialization() {
+    let mut sketch = FrequentItemsSketch::<i64>::new(1 << 12).unwrap();
+    for i in 0..2_000 {
+        sketch.update_with_count(i * 7_919 + 13, (i % 5 + 1) as u64);
+    }
+    let bytes = sketch.serialize();
+    let restored = FrequentItemsSketch::<i64>::deserialize(&bytes).unwrap();
+    assert_eq!(restored.serialize(), bytes);
+    assert_eq!(restored.total_weight(), sketch.total_weight());
+    assert_eq!(restored.num_active_items(), 2_000);
+    assert_eq!(restored.maximum_error(), 0);
+    for i in 0..2_000 {
+        assert_eq!(restored.estimate(&(i * 7_919 + 13)), (i % 5 + 1) as u64);
+    }
+}
+
 #[test]
 fn test_items_round_trip() {
     let mut sketch = FrequentItemsSketch::new(32).unwrap();