chore: refactor VarianceAccumulator, add tests and benchmark (#23977)

## Which issue does this PR close?

- N/A

## Rationale for this change

This PR adds a benchmark for `VarianceAccumulator::update_batch`, and
`retract_batch`, adds some unit tests, and refactors the `retract_batch`
code to add a helper.

## What changes are included in this PR?

See above.

## Are these changes tested?

Yes, new tests added.

## Are there any user-facing changes?

No.

Author: neilconway

datafusion/functions-aggregate/Cargo.toml

@@ -100,5 +100,9 @@ harness = false
 name = "sliding_max"
 harness = false
 
+[[bench]]
+name = "variance"
+harness = false
+
 [features]
 force_hash_collisions = ["datafusion-common/force_hash_collisions"]

datafusion/functions-aggregate/benches/variance.rs

@@ -0,0 +1,83 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::hint::black_box;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, Float64Array};
+use criterion::{BatchSize, Criterion, criterion_group, criterion_main};
+use datafusion_expr::Accumulator;
+use datafusion_functions_aggregate::variance::VarianceAccumulator;
+use datafusion_functions_aggregate_common::stats::StatsType;
+
+const BATCH_SIZE: usize = 8192;
+
+fn batch_array(null_stride: Option<usize>) -> ArrayRef {
+    let values = (0..BATCH_SIZE)
+        .map(|idx| {
+            if null_stride.is_some_and(|stride| idx % stride == 0) {
+                None
+            } else {
+                Some(idx as f64)
+            }
+        })
+        .collect::<Vec<_>>();
+    Arc::new(Float64Array::from(values)) as ArrayRef
+}
+
+fn update_bench(c: &mut Criterion, name: &str, batch: &ArrayRef) {
+    c.bench_function(name, |b| {
+        b.iter(|| {
+            let mut acc = VarianceAccumulator::try_new(StatsType::Sample).unwrap();
+            acc.update_batch(std::slice::from_ref(batch)).unwrap();
+            black_box(acc.evaluate().unwrap())
+        })
+    });
+}
+
+fn retract_bench(c: &mut Criterion, name: &str, batch: &ArrayRef) {
+    c.bench_function(name, |b| {
+        b.iter_batched(
+            || {
+                let mut acc = VarianceAccumulator::try_new(StatsType::Sample).unwrap();
+                // Accumulate two batches so that retracting one leaves the
+                // accumulator with rows remaining, as in a sliding window.
+                acc.update_batch(std::slice::from_ref(batch)).unwrap();
+                acc.update_batch(std::slice::from_ref(batch)).unwrap();
+                acc
+            },
+            |mut acc| {
+                acc.retract_batch(std::slice::from_ref(batch)).unwrap();
+                black_box(acc.evaluate().unwrap())
+            },
+            BatchSize::SmallInput,
+        )
+    });
+}
+
+fn variance_benchmark(c: &mut Criterion) {
+    let no_nulls = batch_array(None);
+    let with_nulls = batch_array(Some(10));
+
+    update_bench(c, "variance update_batch f64 no_nulls", &no_nulls);
+    update_bench(c, "variance update_batch f64 with_nulls", &with_nulls);
+    retract_bench(c, "variance retract_batch f64 no_nulls", &no_nulls);
+    retract_bench(c, "variance retract_batch f64 with_nulls", &with_nulls);
+}
+
+criterion_group!(benches, variance_benchmark);
+criterion_main!(benches);

datafusion/functions-aggregate/src/variance.rs

@@ -326,6 +326,23 @@ fn update(count: u64, mean: f64, m2: f64, value: f64) -> (u64, f64, f64) {
     (new_count, new_mean, new_m2)
 }
 
+/// Inverse of [`update`]: removes a previously accumulated value. Retracting
+/// from a state with one or zero values resets the state to empty.
+#[inline]
+fn retract(count: u64, mean: f64, m2: f64, value: f64) -> (u64, f64, f64) {
+    if count <= 1 {
+        return (0, 0.0, 0.0);
+    }
+
+    let new_count = count - 1;
+    let delta1 = mean - value;
+    let new_mean = delta1 / new_count as f64 + mean;
+    let delta2 = new_mean - value;
+    let new_m2 = m2 - delta1 * delta2;
+
+    (new_count, new_mean, new_m2)
+}
+
 impl Accumulator for VarianceAccumulator {
     fn state(&mut self) -> Result<Vec<ScalarValue>> {
         Ok(vec![
@@ -348,22 +365,8 @@ impl Accumulator for VarianceAccumulator {
     fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
         let arr = as_float64_array(&values[0])?;
         for value in arr.iter().flatten() {
-            if self.count <= 1 {
-                self.count = 0;
-                self.mean = 0.0;
-                self.m2 = 0.0;
-                continue;
-            }
-
-            let new_count = self.count - 1;
-            let delta1 = self.mean - value;
-            let new_mean = delta1 / new_count as f64 + self.mean;
-            let delta2 = new_mean - value;
-            let new_m2 = self.m2 - delta1 * delta2;
-
-            self.count -= 1;
-            self.mean = new_mean;
-            self.m2 = new_m2;
+            (self.count, self.mean, self.m2) =
+                retract(self.count, self.mean, self.m2, value)
         }
 
         Ok(())
@@ -691,6 +694,75 @@ mod tests {
 
     use super::*;
 
+    #[test]
+    fn update_batch_ignores_nulls() -> Result<()> {
+        // An array with nulls must accumulate the same values as a dense
+        // array of its non-null values.
+        let dense: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0]));
+        let sparse: ArrayRef = Arc::new(Float64Array::from(vec![
+            Some(1.0),
+            None,
+            Some(2.0),
+            Some(3.0),
+            None,
+            Some(4.0),
+        ]));
+
+        let mut dense_acc = VarianceAccumulator::try_new(StatsType::Sample)?;
+        dense_acc.update_batch(std::slice::from_ref(&dense))?;
+        let mut sparse_acc = VarianceAccumulator::try_new(StatsType::Sample)?;
+        sparse_acc.update_batch(std::slice::from_ref(&sparse))?;
+
+        // Sample variance of {1, 2, 3, 4} is 5/3 (all steps are exact in f64).
+        assert_eq!(dense_acc.evaluate()?, ScalarValue::Float64(Some(5.0 / 3.0)));
+        assert_eq!(dense_acc.evaluate()?, sparse_acc.evaluate()?);
+        Ok(())
+    }
+
+    #[test]
+    fn retract_batch_ignores_nulls() -> Result<()> {
+        let values: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0]));
+        let dense_retract: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0]));
+        let sparse_retract: ArrayRef =
+            Arc::new(Float64Array::from(vec![Some(1.0), None, Some(2.0)]));
+
+        let mut dense_acc = VarianceAccumulator::try_new(StatsType::Sample)?;
+        dense_acc.update_batch(std::slice::from_ref(&values))?;
+        dense_acc.retract_batch(std::slice::from_ref(&dense_retract))?;
+        let mut sparse_acc = VarianceAccumulator::try_new(StatsType::Sample)?;
+        sparse_acc.update_batch(std::slice::from_ref(&values))?;
+        sparse_acc.retract_batch(std::slice::from_ref(&sparse_retract))?;
+
+        // Sample variance of the remaining {3, 4} is 0.5 (all steps are exact
+        // in f64).
+        assert_eq!(dense_acc.evaluate()?, ScalarValue::Float64(Some(0.5)));
+        assert_eq!(dense_acc.evaluate()?, sparse_acc.evaluate()?);
+        Ok(())
+    }
+
+    #[test]
+    fn retract_batch_resets_when_underflowing() -> Result<()> {
+        // Retracting more values than were accumulated resets to the empty
+        // state, with or without nulls in the retracted batch.
+        let values: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0]));
+        let dense_retract: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0]));
+        let sparse_retract: ArrayRef = Arc::new(Float64Array::from(vec![
+            Some(1.0),
+            None,
+            Some(2.0),
+            Some(3.0),
+        ]));
+
+        for retract in [&dense_retract, &sparse_retract] {
+            let mut acc = VarianceAccumulator::try_new(StatsType::Sample)?;
+            acc.update_batch(std::slice::from_ref(&values))?;
+            acc.retract_batch(std::slice::from_ref(retract))?;
+            assert_eq!(acc.get_count(), 0);
+            assert_eq!(acc.evaluate()?, ScalarValue::Float64(None));
+        }
+        Ok(())
+    }
+
     #[test]
     fn test_groups_accumulator_merge_empty_states() -> Result<()> {
         let state_1 = vec![