feat(frequencies): add borrowed-key update methods to the frequent items sketch

datasketches/src/frequencies/reverse_purge_item_hash_map.rs

@@ -20,6 +20,7 @@
 //! This linear-probing hash map supports a reverse purge operation that removes
 //! keys with non-positive counts by scanning clusters from the back to the front.
 
+use std::borrow::Borrow;
 use std::hash::Hash;
 use std::hash::Hasher;
 
@@ -59,8 +60,15 @@ impl<T: Eq + Hash> ReversePurgeItemHashMap<T> {
     }
 
     /// Returns the value for `key`, or zero if the key is not present.
-    pub fn get(&self, key: &T) -> u64 {
-        let probe = self.hash_probe(key);
+    ///
+    /// The key may be any borrowed form of `T` (e.g. `&str` for a `String` map),
+    /// matching the `HashMap::get` convention.
+    pub fn get<Q>(&self, key: &Q) -> u64
+    where
+        T: Borrow<Q>,
+        Q: Eq + Hash + ?Sized,
+    {
+        let (probe, _) = self.find_probe_or_empty(key);
         if self.states[probe] > 0 {
             return self.values[probe];
         }
@@ -69,21 +77,7 @@ impl<T: Eq + Hash> ReversePurgeItemHashMap<T> {
 
     /// Adds `adjust_amount` to the value for `key`, inserting if absent.
     pub fn adjust_or_put_value(&mut self, key: T, adjust_amount: 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 {
-            let matches = self.keys[probe]
-                .as_ref()
-                .map(|existing| existing == &key)
-                .unwrap_or(false);
-            if matches {
-                break;
-            }
-            probe = (probe + 1) & mask;
-            drift += 1;
-            debug_assert!(drift < DRIFT_LIMIT, "drift limit exceeded");
-        }
+        let (probe, drift) = self.find_probe_or_empty(&key);
         if self.states[probe] == 0 {
             self.keys[probe] = Some(key);
             self.values[probe] = adjust_amount;
@@ -94,6 +88,23 @@ impl<T: Eq + Hash> ReversePurgeItemHashMap<T> {
         }
     }
 
+    /// Adds `adjust_amount` to the value for a borrowed `key`, inserting if absent.
+    pub fn adjust_or_put_value_ref<Q>(&mut self, key: &Q, adjust_amount: u64)
+    where
+        T: Borrow<Q>,
+        Q: Eq + Hash + ToOwned<Owned = T> + ?Sized,
+    {
+        let (probe, drift) = self.find_probe_or_empty(key);
+        if self.states[probe] == 0 {
+            self.keys[probe] = Some(key.to_owned());
+            self.values[probe] = adjust_amount;
+            self.states[probe] = drift as u16;
+            self.num_active += 1;
+        } else {
+            self.values[probe] += adjust_amount;
+        }
+    }
+
     /// Removes all keys with non-positive counts.
     fn keep_only_positive_counts(&mut self) {
         let len = self.keys.len();
@@ -228,20 +239,27 @@ impl<T: Eq + Hash> ReversePurgeItemHashMap<T> {
         self.states[probe] > 0
     }
 
-    fn hash_probe(&self, key: &T) -> usize {
+    fn find_probe_or_empty<Q>(&self, key: &Q) -> (usize, usize)
+    where
+        T: Borrow<Q>,
+        Q: Eq + Hash + ?Sized,
+    {
         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 {
             let matches = self.keys[probe]
                 .as_ref()
-                .map(|existing| existing == key)
+                .map(|existing| existing.borrow() == key)
                 .unwrap_or(false);
             if matches {
                 break;
             }
             probe = (probe + 1) & mask;
+            drift += 1;
+            debug_assert!(drift < DRIFT_LIMIT, "drift limit exceeded");
         }
-        probe
+        (probe, drift)
     }
 
     fn hash_delete(&mut self, mut delete_probe: usize) {
@@ -312,7 +330,7 @@ impl<'a, T> Iterator for ReversePurgeItemIter<'a, T> {
 }
 
 #[inline]
-fn hash_item<T: Hash>(item: &T) -> u64 {
+fn hash_item<T: Hash + ?Sized>(item: &T) -> u64 {
     let mut hasher = MurmurHash3X64128::default();
     item.hash(&mut hasher);
     hasher.finish()

datasketches/src/frequencies/sketch.rs

@@ -17,6 +17,7 @@
 
 //! Frequent items sketch implementations.
 
+use std::borrow::Borrow;
 use std::hash::Hash;
 
 use crate::codec::SketchBytes;
@@ -158,7 +159,11 @@ impl<T: Eq + Hash> FrequentItemsSketch<T> {
     /// sketch.update_with_count(10, 2);
     /// assert!(sketch.estimate(&10) >= 2);
     /// ```
-    pub fn estimate(&self, item: &T) -> u64 {
+    pub fn estimate<Q>(&self, item: &Q) -> u64
+    where
+        T: Borrow<Q>,
+        Q: Eq + Hash + ?Sized,
+    {
         let value = self.hash_map.get(item);
         if value > 0 { value + self.offset } else { 0 }
     }
@@ -167,15 +172,23 @@ impl<T: Eq + Hash> FrequentItemsSketch<T> {
     ///
     /// This value is guaranteed to be no larger than the true frequency. If the item is not
     /// tracked, the lower bound is zero.
-    pub fn lower_bound(&self, item: &T) -> u64 {
+    pub fn lower_bound<Q>(&self, item: &Q) -> u64
+    where
+        T: Borrow<Q>,
+        Q: Eq + Hash + ?Sized,
+    {
         self.hash_map.get(item)
     }
 
     /// Returns the guaranteed upper bound frequency for an item.
     ///
     /// This value is guaranteed to be no smaller than the true frequency. If the item is tracked,
     /// this is `item_count + offset`.
-    pub fn upper_bound(&self, item: &T) -> u64 {
+    pub fn upper_bound<Q>(&self, item: &Q) -> u64
+    where
+        T: Borrow<Q>,
+        Q: Eq + Hash + ?Sized,
+    {
         self.hash_map.get(item) + self.offset
     }
 
@@ -263,6 +276,57 @@ impl<T: Eq + Hash> FrequentItemsSketch<T> {
         self.maybe_resize_or_purge();
     }
 
+    /// Updates the sketch with a borrowed item and a count of one.
+    ///
+    /// Equivalent to [`update`](Self::update) but takes the item by reference and
+    /// only allocates an owned item when it is newly inserted, so updating an
+    /// already-tracked item is allocation free.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use datasketches::frequencies::FrequentItemsSketch;
+    /// let mut sketch = FrequentItemsSketch::<String>::new(64);
+    /// sketch.update_ref("nginx");
+    /// sketch.update_ref("nginx"); // no allocation on the second hit
+    /// assert!(sketch.estimate("nginx") >= 2);
+    /// ```
+    pub fn update_ref<Q>(&mut self, item: &Q)
+    where
+        T: Borrow<Q>,
+        Q: Eq + Hash + ToOwned<Owned = T> + ?Sized,
+    {
+        self.update_with_count_ref(item, 1);
+    }
+
+    /// Updates the sketch with a borrowed item and count.
+    ///
+    /// Equivalent to [`update_with_count`](Self::update_with_count) but takes the
+    /// item by reference and only allocates an owned item when it is newly
+    /// inserted. A count of zero is a no-op.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use datasketches::frequencies::FrequentItemsSketch;
+    /// let mut sketch = FrequentItemsSketch::<String>::new(64);
+    /// sketch.update_with_count_ref("gzip", 3);
+    /// assert!(sketch.estimate("gzip") >= 3);
+    /// ```
+    pub fn update_with_count_ref<Q>(&mut self, item: &Q, count: u64)
+    where
+        T: Borrow<Q>,
+        Q: Eq + Hash + ToOwned<Owned = T> + ?Sized,
+    {
+        if count == 0 {
+            return;
+        }
+        assert!(count > 0, "count may not be negative");
+        self.stream_weight += count;
+        self.hash_map.adjust_or_put_value_ref(item, count);
+        self.maybe_resize_or_purge();
+    }
+
     /// Merges another sketch into this one.
     ///
     /// The other sketch may have a different map size. The merged sketch respects the

datasketches/tests/frequencies_update_test.rs

@@ -119,6 +119,34 @@ fn test_items_one_item() {
     assert_eq!(sketch.upper_bound(&item), 1);
 }
 
+#[test]
+fn test_items_borrowed_key_updates_and_queries() {
+    let mut sketch = FrequentItemsSketch::<String>::new(16);
+
+    sketch.update_ref("alpha");
+    sketch.update_ref("alpha");
+    sketch.update_with_count_ref("beta", 3);
+    sketch.update_with_count_ref("ignored", 0);
+    for item in [
+        "gamma", "delta", "epsilon", "zeta", "eta", "theta", "iota", "kappa", "lambda", "mu",
+    ] {
+        sketch.update_ref(item);
+    }
+
+    assert!(!sketch.is_empty());
+    assert_eq!(sketch.total_weight(), 15);
+    assert_eq!(sketch.num_active_items(), 12);
+    assert_eq!(sketch.estimate("alpha"), 2);
+    assert_eq!(sketch.lower_bound("alpha"), 2);
+    assert_eq!(sketch.upper_bound("alpha"), 2);
+    assert_eq!(sketch.estimate("beta"), 3);
+    assert_eq!(sketch.estimate("ignored"), 0);
+    assert_eq!(sketch.estimate("missing"), 0);
+
+    let owned = "alpha".to_string();
+    assert_eq!(sketch.estimate(&owned), 2);
+}
+
 #[test]
 fn test_longs_several_items_no_resize_no_purge() {
     let mut sketch: FrequentItemsSketch<i64> = FrequentItemsSketch::new(8);