From 571db962225d23d7757cca6a83be1bc01266e7da Mon Sep 17 00:00:00 2001
From: zz_y <zeyingz@umd.edu>
Date: Thu, 15 Oct 2026 20:42:22 +0000
Subject: [PATCH] perf(weighted_frequency): rescore retained candidates without
 copying them

With a candidate heap, every update cloned each retained candidate
(its identity tuple and hash key) to rebuild the heap, and scoring each
candidate collected its row estimates into a fresh Vec. At capacity 100
with two-label Utf8 identities that was about 66 us (CMS) and 77 us
(CountSketch) per update in a release build.

The update now moves the candidates out of the heap, adds the new key
only when it is not already retained, and scores them with one reused
buffer (Count-Min takes the row minimum without any buffer). The same
benchmark runs at about 33 us and 43 us. Every retained score is still
re-read after each update, so the candidates and their scores are
unchanged; the golden snapshots still match.
---
 CHANGELOG.md                       |  5 ++-
 src/sketches/weighted_frequency.rs | 62 +++++++++++++++++++++++-------
 2 files changed, 52 insertions(+), 15 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a3fef430..a8de475e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -19,7 +19,10 @@ signals a backwards-compatible change.
   capacity 0 builds a frequency sketch that retains no candidates.
   `WeightedFrequency::check_shape` and `check_weight` validate a shape or a
   weight without building or updating a state, and `clear` empties a state
-  in place.
+  in place. An update with a candidate heap moves the retained candidates
+  into the rescoring instead of copying them, and scores them with one
+  buffer, about halving its cost at capacity 100; the retained candidates
+  and their scores are unchanged.
 
 - **`Vector3D<T>`**, filled in from a stub. A `rows x cols` grid whose every
   cell is a contiguous `depth`-length record, addressed by the same
diff --git a/src/sketches/weighted_frequency.rs b/src/sketches/weighted_frequency.rs
index 8e7c6949..51ea13ff 100644
--- a/src/sketches/weighted_frequency.rs
+++ b/src/sketches/weighted_frequency.rs
@@ -202,15 +202,21 @@ impl WeightedFrequency {
         })
     }
     fn estimate_key(&self, key: &[u8]) -> f64 {
-        let mut estimates = self
-            .indexes(key)
-            .map(|(i, sign)| self.cells[i] * sign)
-            .collect::<Vec<_>>();
+        self.estimate_key_with(key, &mut Vec::new())
+    }
+    /// [`Self::estimate_key`] with a caller-owned buffer for the CountSketch
+    /// row estimates, so scoring many candidates allocates once.
+    fn estimate_key_with(&self, key: &[u8], rows: &mut Vec<f64>) -> f64 {
         match self.algorithm {
-            FrequencyAlgorithm::Cms => estimates.into_iter().fold(f64::INFINITY, f64::min),
+            FrequencyAlgorithm::Cms => self
+                .indexes(key)
+                .map(|(i, _)| self.cells[i])
+                .fold(f64::INFINITY, f64::min),
             FrequencyAlgorithm::CountSketch => {
-                let middle = estimates.len() / 2;
-                *estimates.select_nth_unstable_by(middle, f64::total_cmp).1
+                rows.clear();
+                rows.extend(self.indexes(key).map(|(i, sign)| self.cells[i] * sign));
+                let middle = rows.len() / 2;
+                *rows.select_nth_unstable_by(middle, f64::total_cmp).1
             }
         }
     }
@@ -218,8 +224,9 @@ impl WeightedFrequency {
         candidates.sort_by(|a, b| a.key.cmp(&b.key));
         candidates.dedup_by(|a, b| a.key == b.key);
         self.candidates.clear();
+        let mut rows = Vec::with_capacity(self.depth);
         for mut candidate in candidates {
-            candidate.score = self.estimate_key(&candidate.key);
+            candidate.score = self.estimate_key_with(&candidate.key, &mut rows);
             self.candidates.push(candidate);
             if self.candidates.len() > self.capacity {
                 self.candidates.pop();
@@ -258,12 +265,17 @@ impl WeightedFrequency {
         if self.capacity == 0 {
             return Ok(());
         }
-        let mut candidates = self.candidates.iter().cloned().collect::<Vec<_>>();
-        candidates.push(Candidate {
-            identity,
-            key,
-            score: 0.0,
-        });
+        // Every retained score is re-read, since this update may have moved
+        // cells they share. The candidates move into the rescoring instead of
+        // being copied, and a key already retained is not added twice.
+        let mut candidates = std::mem::take(&mut self.candidates).into_vec();
+        if !candidates.iter().any(|candidate| candidate.key == key) {
+            candidates.push(Candidate {
+                identity,
+                key,
+                score: 0.0,
+            });
+        }
         self.retain(candidates);
         Ok(())
     }
@@ -617,6 +629,28 @@ mod tests {
         }
     }
 
+    // Under heavy collisions every update re-reads each retained score, and a
+    // key updated again stays one candidate: the retained scores always equal
+    // the point estimates.
+    #[test]
+    fn retained_scores_track_every_update() {
+        for algorithm in [FrequencyAlgorithm::Cms, FrequencyAlgorithm::CountSketch] {
+            let mut state = WeightedFrequency::new(algorithm, 8, 3, 3).unwrap();
+            for step in 0..60i64 {
+                let key = [FrequencyIdentity::Int64(step % 7)];
+                state.update(&key, 0.5 + (step % 3) as f64).unwrap();
+                let rows = state.topk(8);
+                assert!(rows.len() <= 3);
+                for (identity, score) in &rows {
+                    assert_eq!(*score, state.estimate(identity).unwrap(), "step {step}");
+                }
+                let mut keys = rows.iter().map(|(id, _)| id.clone()).collect::<Vec<_>>();
+                keys.dedup();
+                assert_eq!(keys.len(), rows.len());
+            }
+        }
+    }
+
     // Overflow errors are atomic, and candidate storage never exceeds its configured bound.
     #[test]
     fn overflow_and_candidate_capacity() {