From 74a0b68195063d8493dbcb0c5480f14bf0ae80d5 Mon Sep 17 00:00:00 2001
From: Bryan O'Sullivan <bos@voxelai.com>
Date: Wed, 21 Oct 2026 00:11:46 -0700
Subject: [PATCH] threshold_total: treat a zero count as absent

---
 .../src/operators/threshold.rs                |  8 ++-
 differential-dataflow/tests/threshold.rs      | 55 +++++++++++++++++++
 2 files changed, 61 insertions(+), 2 deletions(-)
 create mode 100644 differential-dataflow/tests/threshold.rs

diff --git a/differential-dataflow/src/operators/threshold.rs b/differential-dataflow/src/operators/threshold.rs
index c9502976f..e41457b29 100644
--- a/differential-dataflow/src/operators/threshold.rs
+++ b/differential-dataflow/src/operators/threshold.rs
@@ -26,6 +26,8 @@ pub trait ThresholdTotal<'scope, T: Timestamp + TotalOrder + Lattice, K: Exchang
         ;
     /// Reduces the collection to one occurrence of each distinct element.
     ///
+    /// As with `threshold`, a zero count thresholds to zero, whatever `thresh` returns for it.
+    ///
     /// # Examples
     ///
     /// ```
@@ -41,8 +43,10 @@ pub trait ThresholdTotal<'scope, T: Timestamp + TotalOrder + Lattice, K: Exchang
     /// ```
     fn threshold_total<R2: Abelian+'static, F: FnMut(&K,&R)->R2+'static>(self, mut thresh: F) -> VecCollection<'scope, T, K, R2> {
         self.threshold_semigroup(move |key, new, old| {
-            let mut new = thresh(key, new);
-            if let Some(old) = old {
+            // A key whose count is zero may or may not still have updates in the trace, so
+            // treat a zero count as absent rather than call `thresh` on it.
+            let mut new = if new.is_zero() { R2::zero() } else { thresh(key, new) };
+            if let Some(old) = old.filter(|old| !old.is_zero()) {
                 let mut add = thresh(key, old);
                 add.negate();
                 new.plus_equals(&add);
diff --git a/differential-dataflow/tests/threshold.rs b/differential-dataflow/tests/threshold.rs
new file mode 100644
index 000000000..c7d75943c
--- /dev/null
+++ b/differential-dataflow/tests/threshold.rs
@@ -0,0 +1,55 @@
+use timely::dataflow::operators::capture::{Capture, Extract};
+use differential_dataflow::VecCollection;
+use differential_dataflow::input::InputSession;
+use differential_dataflow::operators::ThresholdTotal;
+
+/// Feeds `updates` to `logic`, one time per batch if `per_time` is set, and returns the
+/// consolidated output.
+fn run<F>(updates: &[(u64, u64, isize)], per_time: bool, logic: F) -> Vec<(u64, u64, isize)>
+where
+    F: for<'scope> FnOnce(VecCollection<'scope, u64, u64, isize>) -> VecCollection<'scope, u64, u64, isize> + Send + Sync + 'static,
+{
+    let updates = updates.to_vec();
+    let (send, recv) = std::sync::mpsc::channel();
+    timely::execute_directly(move |worker| {
+        let mut input = InputSession::<u64, u64, isize>::new();
+        let probe = worker.dataflow::<u64, _, _>(|scope| {
+            let (probe, output) = logic(input.to_collection(scope)).probe();
+            output.inner.capture_into(send);
+            probe
+        });
+        for (key, time, diff) in updates {
+            input.update_at(key, time, diff);
+            if per_time {
+                input.advance_to(time + 1);
+                input.flush();
+                worker.step_while(|| probe.less_than(&(time + 1)));
+            }
+        }
+    });
+    let mut output: Vec<_> = recv.extract().into_iter().flat_map(|(_, data)| data).collect();
+    differential_dataflow::consolidation::consolidate_updates(&mut output);
+    output
+}
+
+/// A key whose count returns to zero leaves the output, and returns with its count, whether or
+/// not the updates that cancel are still separate in the trace.
+#[test]
+fn distinct_total_zero_count() {
+    let updates = [(0, 1, 1), (0, 2, -1), (0, 3, 1), (0, 4, 1), (0, 5, -2)];
+    let expected = vec![(0, 1, 1), (0, 2, -1), (0, 3, 1), (0, 5, -1)];
+    for per_time in [false, true] {
+        assert_eq!(run(&updates, per_time, |c| c.distinct_total()), expected, "per_time: {per_time}");
+    }
+}
+
+/// `threshold_total` treats a zero count as absent, as `threshold` does, even when `thresh`
+/// maps zero to something else.
+#[test]
+fn threshold_total_zero_count() {
+    let updates = [(0, 1, 1), (0, 2, -1), (0, 3, 2)];
+    let expected = vec![(0, 1, 2), (0, 2, -2), (0, 3, 3)];
+    for per_time in [false, true] {
+        assert_eq!(run(&updates, per_time, |c| c.threshold_total(|_, c| c + 1)), expected, "per_time: {per_time}");
+    }
+}