From ac3aa3d3a496d6c6802e0563ae40b8344f6a0d3a Mon Sep 17 00:00:00 2001
From: edubraqd <eduardoalcantara.sp@gmail.com>
Date: Wed, 2 Sep 2026 23:32:58 -0300
Subject: [PATCH] fix: wrap instead of panicking on overflow in the Duration
 avg accumulator

`DurationAvgAccumulator` accumulated its running sum with plain `+=` / `-`,
which panics with "attempt to add with overflow" in debug builds (and wraps
silently in release builds) once the sum of the input durations no longer fits
in `i64`, for example `avg(x)` over a window that contains two `i64::MAX`
durations.

The `sum` aggregate and every other `avg` accumulator deliberately wrap on
overflow (see the docs on `add_avg_sum`). Use `wrapping_add` / `wrapping_sub`
here too so the behaviour is consistent across build profiles.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 datafusion/functions-aggregate/src/average.rs | 44 +++++++++++++++++--
 1 file changed, 41 insertions(+), 3 deletions(-)

diff --git a/datafusion/functions-aggregate/src/average.rs b/datafusion/functions-aggregate/src/average.rs
index 8ce893c17dfd4..5a9a75b30668d 100644
--- a/datafusion/functions-aggregate/src/average.rs
+++ b/datafusion/functions-aggregate/src/average.rs
@@ -884,7 +884,10 @@ impl Accumulator for DurationAvgAccumulator {
 
         if let Some(x) = sum_value {
             let v = self.sum.get_or_insert(0);
-            *v += x;
+            // Wraps on overflow, matching the `sum` aggregate and the other
+            // `avg` accumulators; the checked arithmetic would panic in debug
+            // builds and silently wrap in release builds otherwise.
+            *v = v.wrapping_add(x);
         }
         Ok(())
     }
@@ -940,7 +943,10 @@ impl Accumulator for DurationAvgAccumulator {
 
         if let Some(x) = sum_value {
             let v = self.sum.get_or_insert(0);
-            *v += x;
+            // Wraps on overflow, matching the `sum` aggregate and the other
+            // `avg` accumulators; the checked arithmetic would panic in debug
+            // builds and silently wrap in release builds otherwise.
+            *v = v.wrapping_add(x);
         }
         Ok(())
     }
@@ -957,7 +963,7 @@ impl Accumulator for DurationAvgAccumulator {
         };
 
         if let Some(x) = sum_value {
-            self.sum = Some(self.sum.unwrap() - x);
+            self.sum = Some(self.sum.unwrap_or(0).wrapping_sub(x));
         }
         Ok(())
     }
@@ -1552,6 +1558,38 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn avg_duration_wraps_on_overflow() -> Result<()> {
+        // The sum of these values does not fit in `i64`. Like `sum` and the
+        // other `avg` accumulators the running total wraps; this used to
+        // panic with "attempt to add with overflow" in debug builds.
+        let duration = DataType::Duration(TimeUnit::Second);
+        let values: ArrayRef =
+            Arc::new(DurationSecondArray::from(vec![i64::MAX, i64::MAX]));
+
+        let mut acc = avg_accumulator(&duration, &duration)?;
+        acc.update_batch(std::slice::from_ref(&values))?;
+        assert_eq!(acc.evaluate()?, ScalarValue::DurationSecond(Some(-1)));
+
+        // Merging two such states wraps as well.
+        let mut merged = avg_accumulator(&duration, &duration)?;
+        let state = acc.state()?;
+        let state_arrays: Vec<ArrayRef> = state
+            .iter()
+            .map(|v| v.to_array_of_size(1))
+            .collect::<Result<_>>()?;
+        merged.merge_batch(&state_arrays)?;
+        merged.merge_batch(&state_arrays)?;
+        assert_eq!(merged.evaluate()?, ScalarValue::DurationSecond(Some(-1)));
+
+        // Retracting the batch in sliding-window mode brings the sum back to
+        // where it started instead of underflowing.
+        acc.retract_batch(&[values])?;
+        assert_eq!(acc.evaluate()?, ScalarValue::DurationSecond(None));
+
+        Ok(())
+    }
+
     #[test]
     fn avg_accumulator_evaluate_and_state_types() -> Result<()> {
         for case in avg_cases()? {