fix: match Spark variance updates for large nearby values

native/spark-expr/src/agg_funcs/correlation.rs

@@ -144,8 +144,10 @@ impl CorrelationAccumulator {
     pub fn try_new(null_on_divide_by_zero: bool) -> Result<Self> {
         Ok(Self {
             covar: CovarianceAccumulator::try_new(StatsType::Population, null_on_divide_by_zero)?,
-            stddev1: StddevAccumulator::try_new(StatsType::Population, null_on_divide_by_zero)?,
-            stddev2: StddevAccumulator::try_new(StatsType::Population, null_on_divide_by_zero)?,
+            stddev1: StddevAccumulator::try_new(StatsType::Population, null_on_divide_by_zero)?
+                .with_pearson_update(),
+            stddev2: StddevAccumulator::try_new(StatsType::Population, null_on_divide_by_zero)?
+                .with_pearson_update(),
             null_on_divide_by_zero,
         })
     }
@@ -280,8 +282,10 @@ impl CorrelationGroupsAccumulator {
         // that intent explicit.
         Self {
             covar: CovarianceGroupsAccumulator::new(StatsType::Population, false),
-            var1: VarianceGroupsAccumulator::new(StatsType::Population, false),
-            var2: VarianceGroupsAccumulator::new(StatsType::Population, false),
+            var1: VarianceGroupsAccumulator::new(StatsType::Population, false)
+                .with_pearson_update(),
+            var2: VarianceGroupsAccumulator::new(StatsType::Population, false)
+                .with_pearson_update(),
             null_on_divide_by_zero,
         }
     }

native/spark-expr/src/agg_funcs/covariance.rs

@@ -513,6 +513,34 @@ mod groups_tests {
         assert!((evaluate(&mut acc)[0].unwrap() - 4.0).abs() < 1e-12);
     }
 
+    #[test]
+    fn large_offset_covariance_merge() {
+        let mut scalar = CovarianceAccumulator::try_new(StatsType::Population, true).unwrap();
+        let mut grouped = pop();
+        for (value, count) in [(0.0, 0), (1e17 - 32.0, 3), (1e17 - 16.0, 2), (0.0, 0)] {
+            let mut xs = vec![Some(value); count];
+            xs.push(None);
+            let ys = xs.iter().map(|v| v.map(|v| -v)).collect::<Vec<_>>();
+            let values: Vec<ArrayRef> = vec![
+                Arc::new(Float64Array::from(xs)),
+                Arc::new(Float64Array::from(ys)),
+            ];
+            let mut partial = pop();
+            partial
+                .update_batch(&values, &vec![0; count + 1], None, 2)
+                .unwrap();
+            let state = partial.state(EmitTo::All).unwrap();
+            scalar.merge_batch(&state).unwrap();
+            grouped.merge_batch(&state, &[0, 1], 2).unwrap();
+        }
+        let expected = -61.44000000000001;
+        assert_eq!(
+            scalar.evaluate().unwrap(),
+            ScalarValue::Float64(Some(expected))
+        );
+        assert_eq!(evaluate(&mut grouped), vec![Some(expected), None]);
+    }
+
     #[test]
     fn null_in_either_column_skipped() {
         let mut acc = pop();

native/spark-expr/src/agg_funcs/regr.rs

@@ -306,8 +306,10 @@ impl RegrR2Accumulator {
     fn try_new(constant_dependent_is_perfect_fit: bool) -> Result<Self> {
         Ok(Self {
             covar: CovarianceAccumulator::try_new(StatsType::Population, false)?,
-            var_y: VarianceAccumulator::try_new(StatsType::Population, false)?,
-            var_x: VarianceAccumulator::try_new(StatsType::Population, false)?,
+            var_y: VarianceAccumulator::try_new(StatsType::Population, false)?
+                .with_pearson_update(),
+            var_x: VarianceAccumulator::try_new(StatsType::Population, false)?
+                .with_pearson_update(),
             constant_dependent_is_perfect_fit,
         })
     }

native/spark-expr/src/agg_funcs/stddev.rs

@@ -163,6 +163,11 @@ impl StddevAccumulator {
     pub fn get_m2(&self) -> f64 {
         self.variance.get_m2()
     }
+
+    pub(super) fn with_pearson_update(mut self) -> Self {
+        self.variance = self.variance.with_pearson_update();
+        self
+    }
 }
 
 impl Accumulator for StddevAccumulator {

native/spark-expr/src/agg_funcs/variance.rs

@@ -29,6 +29,8 @@ use datafusion::physical_expr::expressions::StatsType;
 use std::mem::size_of;
 use std::sync::Arc;
 
+use super::welford::VarianceUpdate;
+
 /// VAR_SAMP and VAR_POP aggregate expression
 /// The implementation mostly is the same as the DataFusion's implementation. The reason
 /// we have our own implementation is that DataFusion has UInt64 for state_field `count`,
@@ -144,6 +146,7 @@ pub struct VarianceAccumulator {
     count: f64,
     stats_type: StatsType,
     null_on_divide_by_zero: bool,
+    update: VarianceUpdate,
 }
 
 impl VarianceAccumulator {
@@ -155,9 +158,15 @@ impl VarianceAccumulator {
             count: 0_f64,
             stats_type: s_type,
             null_on_divide_by_zero,
+            update: VarianceUpdate::CentralMoment,
         })
     }
 
+    pub(super) fn with_pearson_update(mut self) -> Self {
+        self.update = VarianceUpdate::Pearson;
+        self
+    }
+
     pub fn get_count(&self) -> f64 {
         self.count
     }
@@ -184,7 +193,8 @@ impl Accumulator for VarianceAccumulator {
         let arr = downcast_value!(&values[0], Float64Array).iter().flatten();
 
         for value in arr {
-            let (c, m, m2) = super::welford::variance_update(self.count, self.mean, self.m2, value);
+            let (c, m, m2) =
+                super::welford::variance_update(self.count, self.mean, self.m2, value, self.update);
             self.count = c;
             self.mean = m;
             self.m2 = m2;
@@ -271,6 +281,7 @@ pub(crate) struct VarianceGroupsAccumulator {
     pub(super) m2s: Vec<f64>,
     stats_type: StatsType,
     null_on_divide_by_zero: bool,
+    update: VarianceUpdate,
 }
 
 impl VarianceGroupsAccumulator {
@@ -281,9 +292,15 @@ impl VarianceGroupsAccumulator {
             m2s: Vec::new(),
             stats_type,
             null_on_divide_by_zero,
+            update: VarianceUpdate::CentralMoment,
         }
     }
 
+    pub(super) fn with_pearson_update(mut self) -> Self {
+        self.update = VarianceUpdate::Pearson;
+        self
+    }
+
     fn resize(&mut self, total_num_groups: usize) {
         self.counts.resize(total_num_groups, 0.0);
         self.means.resize(total_num_groups, 0.0);
@@ -327,6 +344,7 @@ impl GroupsAccumulator for VarianceGroupsAccumulator {
                 self.means[group_index],
                 self.m2s[group_index],
                 value,
+                self.update,
             );
             self.counts[group_index] = c;
             self.means[group_index] = m;
@@ -421,6 +439,93 @@ mod groups_tests {
             .collect()
     }
 
+    #[test]
+    fn large_offset_variance() {
+        for pair in [[1e16, 1e16 + 2.0], [1e16 + 2.0, 1e16], [-1e16, -1e16 - 2.0]] {
+            for (stats, expected) in [(StatsType::Population, 1.0), (StatsType::Sample, 2.0)] {
+                let values: ArrayRef =
+                    Arc::new(Float64Array::from(vec![Some(pair[0]), None, Some(pair[1])]));
+                for batch_size in [1, 3] {
+                    let mut scalar = VarianceAccumulator::try_new(stats, true).unwrap();
+                    let mut grouped = VarianceGroupsAccumulator::new(stats, true);
+                    for offset in (0..3).step_by(batch_size) {
+                        let batch = [values.slice(offset, batch_size)];
+                        scalar.update_batch(&batch).unwrap();
+                        grouped
+                            .update_batch(&batch, &vec![0; batch_size], None, 2)
+                            .unwrap();
+                    }
+                    assert_eq!(
+                        scalar.evaluate().unwrap(),
+                        ScalarValue::Float64(Some(expected))
+                    );
+                    let states = grouped.state(EmitTo::All).unwrap();
+                    let mut merged = VarianceGroupsAccumulator::new(stats, true);
+                    merged.merge_batch(&states, &[0, 1], 2).unwrap();
+                    assert_eq!(evaluate(&mut merged), vec![Some(expected), None]);
+                }
+            }
+        }
+    }
+
+    #[test]
+    fn large_offset_variance_merge() {
+        for (partitions, population, sample) in [
+            (
+                [(1e17 - 32.0, 3), (1e17 - 16.0, 2)],
+                61.44000000000001,
+                76.80000000000001,
+            ),
+            ([(1e17 - 96.0, 3), (1e17 - 32.0, 3)], 1024.0, 1228.8),
+        ] {
+            for sign in [1.0, -1.0] {
+                for reverse in [false, true] {
+                    let mut partitions = partitions;
+                    if reverse {
+                        partitions.reverse();
+                    }
+                    for (stats, expected) in [
+                        (StatsType::Population, population),
+                        (StatsType::Sample, sample),
+                    ] {
+                        let mut scalar = VarianceAccumulator::try_new(stats, true).unwrap();
+                        let mut grouped = VarianceGroupsAccumulator::new(stats, true);
+                        // Include empty partials before and after the nonempty states.
+                        for (value, count) in
+                            [(0.0, 0)].into_iter().chain(partitions).chain([(0.0, 0)])
+                        {
+                            let mut values = vec![Some(sign * value); count];
+                            values.push(None);
+                            let values: ArrayRef = Arc::new(Float64Array::from(values));
+                            let mut partial = VarianceAccumulator::try_new(stats, true).unwrap();
+                            partial.update_batch(&[Arc::clone(&values)]).unwrap();
+                            let state = partial
+                                .state()
+                                .unwrap()
+                                .iter()
+                                .map(|v| v.to_array_of_size(1).unwrap())
+                                .collect::<Vec<_>>();
+                            scalar.merge_batch(&state).unwrap();
+
+                            let mut partial = VarianceGroupsAccumulator::new(stats, true);
+                            partial
+                                .update_batch(&[values], &vec![0; count + 1], None, 2)
+                                .unwrap();
+                            grouped
+                                .merge_batch(&partial.state(EmitTo::All).unwrap(), &[0, 1], 2)
+                                .unwrap();
+                        }
+                        assert_eq!(
+                            scalar.evaluate().unwrap(),
+                            ScalarValue::Float64(Some(expected))
+                        );
+                        assert_eq!(evaluate(&mut grouped), vec![Some(expected), None]);
+                    }
+                }
+            }
+        }
+    }
+
     #[test]
     fn pop_variance_single_group() {
         let mut acc = pop_acc();

native/spark-expr/src/agg_funcs/welford.rs

@@ -23,12 +23,30 @@
 use arrow::buffer::NullBuffer;
 use datafusion::physical_expr::expressions::StatsType;
 
+#[derive(Debug, Clone, Copy)]
+pub(super) enum VarianceUpdate {
+    CentralMoment,
+    Pearson,
+}
+
 #[inline]
-pub(crate) fn variance_update(count: f64, mean: f64, m2: f64, value: f64) -> (f64, f64, f64) {
+pub(super) fn variance_update(
+    count: f64,
+    mean: f64,
+    m2: f64,
+    value: f64,
+    update: VarianceUpdate,
+) -> (f64, f64, f64) {
     let new_count = count + 1.0;
     let delta1 = value - mean;
-    let new_mean = delta1 / new_count + mean;
-    let delta2 = value - new_mean;
+    let delta_n = delta1 / new_count;
+    let new_mean = mean + delta_n;
+    // Match Spark's CentralMomentAgg without subtracting the rounded new mean.
+    // PearsonCorrelation (also used by regr_r2) deliberately uses that subtraction.
+    let delta2 = match update {
+        VarianceUpdate::CentralMoment => delta1 - delta_n,
+        VarianceUpdate::Pearson => value - new_mean,
+    };
     let new_m2 = m2 + delta1 * delta2;
     (new_count, new_mean, new_m2)
 }
@@ -53,9 +71,16 @@ pub(crate) fn variance_merge(
     m2_b: f64,
 ) -> (f64, f64, f64) {
     let new_count = count_a + count_b;
-    let new_mean = mean_a * count_a / new_count + mean_b * count_b / new_count;
-    let delta = mean_a - mean_b;
-    let new_m2 = m2_a + m2_b + delta * delta * count_a * count_b / new_count;
+    // CentralMomentAgg and PearsonCorrelation use the same merge expressions.
+    // Preserve Spark's operation order to avoid rounding large means differently.
+    let delta = mean_b - mean_a;
+    let delta_n = if new_count == 0.0 {
+        0.0
+    } else {
+        delta / new_count
+    };
+    let new_mean = mean_a + delta_n * count_b;
+    let new_m2 = m2_a + m2_b + delta * delta_n * count_a * count_b;
     (new_count, new_mean, new_m2)
 }
 
@@ -149,10 +174,21 @@ pub(crate) fn covariance_merge(
     c_b: f64,
 ) -> (f64, f64, f64, f64) {
     let new_count = count_a + count_b;
-    let new_mean1 = mean1_a * count_a / new_count + mean1_b * count_b / new_count;
-    let new_mean2 = mean2_a * count_a / new_count + mean2_b * count_b / new_count;
-    let delta1 = mean1_a - mean1_b;
-    let delta2 = mean2_a - mean2_b;
-    let new_c = c_a + c_b + delta1 * delta2 * count_a * count_b / new_count;
+    // Keep covariance aligned with variance and Spark's Covariance/PearsonCorrelation.
+    let delta1 = mean1_b - mean1_a;
+    let delta2 = mean2_b - mean2_a;
+    let delta1_n = if new_count == 0.0 {
+        0.0
+    } else {
+        delta1 / new_count
+    };
+    let delta2_n = if new_count == 0.0 {
+        0.0
+    } else {
+        delta2 / new_count
+    };
+    let new_mean1 = mean1_a + delta1_n * count_b;
+    let new_mean2 = mean2_a + delta2_n * count_b;
+    let new_c = c_a + c_b + delta1 * delta2_n * count_a * count_b;
     (new_count, new_mean1, new_mean2, new_c)
 }