diff --git a/datasketches/src/frequencies/reverse_purge_item_hash_map.rs b/datasketches/src/frequencies/reverse_purge_item_hash_map.rs index 79ed290..0e59745 100644 --- a/datasketches/src/frequencies/reverse_purge_item_hash_map.rs +++ b/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 ReversePurgeItemHashMap { } /// 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(&self, key: &Q) -> u64 + where + T: Borrow, + 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 ReversePurgeItemHashMap { /// 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 ReversePurgeItemHashMap { } } + /// Adds `adjust_amount` to the value for a borrowed `key`, inserting if absent. + pub fn adjust_or_put_value_ref(&mut self, key: &Q, adjust_amount: u64) + where + T: Borrow, + Q: Eq + Hash + ToOwned + ?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 ReversePurgeItemHashMap { self.states[probe] > 0 } - fn hash_probe(&self, key: &T) -> usize { + fn find_probe_or_empty(&self, key: &Q) -> (usize, usize) + where + T: Borrow, + 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(item: &T) -> u64 { +fn hash_item(item: &T) -> u64 { let mut hasher = MurmurHash3X64128::default(); item.hash(&mut hasher); hasher.finish() diff --git a/datasketches/src/frequencies/sketch.rs b/datasketches/src/frequencies/sketch.rs index fc2b297..f46a231 100644 --- a/datasketches/src/frequencies/sketch.rs +++ b/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 FrequentItemsSketch { /// sketch.update_with_count(10, 2); /// assert!(sketch.estimate(&10) >= 2); /// ``` - pub fn estimate(&self, item: &T) -> u64 { + pub fn estimate(&self, item: &Q) -> u64 + where + T: Borrow, + Q: Eq + Hash + ?Sized, + { let value = self.hash_map.get(item); if value > 0 { value + self.offset } else { 0 } } @@ -167,7 +172,11 @@ impl FrequentItemsSketch { /// /// 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(&self, item: &Q) -> u64 + where + T: Borrow, + Q: Eq + Hash + ?Sized, + { self.hash_map.get(item) } @@ -175,7 +184,11 @@ impl FrequentItemsSketch { /// /// 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(&self, item: &Q) -> u64 + where + T: Borrow, + Q: Eq + Hash + ?Sized, + { self.hash_map.get(item) + self.offset } @@ -263,6 +276,57 @@ impl FrequentItemsSketch { 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::::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(&mut self, item: &Q) + where + T: Borrow, + Q: Eq + Hash + ToOwned + ?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::::new(64); + /// sketch.update_with_count_ref("gzip", 3); + /// assert!(sketch.estimate("gzip") >= 3); + /// ``` + pub fn update_with_count_ref(&mut self, item: &Q, count: u64) + where + T: Borrow, + Q: Eq + Hash + ToOwned + ?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 diff --git a/datasketches/tests/frequencies_update_test.rs b/datasketches/tests/frequencies_update_test.rs index f15f219..87c702a 100644 --- a/datasketches/tests/frequencies_update_test.rs +++ b/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::::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 = FrequentItemsSketch::new(8);