diff --git a/native/spark-expr/src/agg_funcs/correlation.rs b/native/spark-expr/src/agg_funcs/correlation.rs index d69ac61def..4dbdef6b51 100644 --- a/native/spark-expr/src/agg_funcs/correlation.rs +++ b/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 { 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, } } diff --git a/native/spark-expr/src/agg_funcs/covariance.rs b/native/spark-expr/src/agg_funcs/covariance.rs index 548c118720..058c2576f5 100644 --- a/native/spark-expr/src/agg_funcs/covariance.rs +++ b/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::>(); + let values: Vec = 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(); diff --git a/native/spark-expr/src/agg_funcs/regr.rs b/native/spark-expr/src/agg_funcs/regr.rs index 72511fc987..7c73181708 100644 --- a/native/spark-expr/src/agg_funcs/regr.rs +++ b/native/spark-expr/src/agg_funcs/regr.rs @@ -306,8 +306,10 @@ impl RegrR2Accumulator { fn try_new(constant_dependent_is_perfect_fit: bool) -> Result { 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, }) } diff --git a/native/spark-expr/src/agg_funcs/stddev.rs b/native/spark-expr/src/agg_funcs/stddev.rs index bbceaa72dc..1ef31f7a10 100644 --- a/native/spark-expr/src/agg_funcs/stddev.rs +++ b/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 { diff --git a/native/spark-expr/src/agg_funcs/variance.rs b/native/spark-expr/src/agg_funcs/variance.rs index 57a8f6da50..0c722d71de 100644 --- a/native/spark-expr/src/agg_funcs/variance.rs +++ b/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, 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::>(); + 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(); diff --git a/native/spark-expr/src/agg_funcs/welford.rs b/native/spark-expr/src/agg_funcs/welford.rs index bcc44b2988..dc3f8551e3 100644 --- a/native/spark-expr/src/agg_funcs/welford.rs +++ b/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) } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala index 55c2ebdb57..5526acfe18 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometAggregateSuite.scala @@ -2343,6 +2343,88 @@ class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper { } } + test("statistical aggregates with large nearby values") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + "spark.sql.files.minPartitionNum" -> "1", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native") { + for (values <- Seq(Seq(1e16, 1e16 + 2), Seq(1e16 + 2, 1e16), Seq(-1e16, -1e16 - 2))) { + // One ordered file keeps both values in the same partial accumulator. Splitting + // them across files would only exercise merging two single-row states. + withTempPath { path => + (Seq(Some(values.head), None, Some(values.last))) + .map(v => (0, v)) + .toDF("g", "v") + .coalesce(1) + .write + .parquet(path.getCanonicalPath) + withParquetTable(path.getCanonicalPath, "large_moments") { + for (groupBy <- Seq("", " GROUP BY g")) { + val query = "SELECT var_pop(v), var_samp(v), stddev_pop(v), stddev_samp(v) " + + "FROM large_moments" + groupBy + val (_, cometPlan) = checkSparkAnswerAndOperator(query) + val aggregates = cometPlan.collect { case a: CometHashAggregateExec => a } + assert(aggregates.exists(_.modes.contains(Partial))) + assert(aggregates.exists(_.modes.contains(Final))) + checkAnswer(sql(query), Seq(Row(1.0, 2.0, 1.0, math.sqrt(2.0)))) + + // CORR and REGR_R2 use PearsonCorrelation's update, while REGR_SXX/SYY + // and the variance used by slope/intercept follow CentralMomentAgg. + checkSparkAnswerWithTolAndNumOfAggregates( + "SELECT corr(v, v), regr_r2(v, v), regr_sxx(v, v), regr_syy(v, v), " + + "regr_slope(v, v), regr_intercept(v, v) FROM large_moments" + groupBy, + 2) + } + } + } + } + } + } + + test("statistical aggregates merge large nearby values across partitions") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + SQLConf.FILES_MAX_PARTITION_BYTES.key -> "1048576", + SQLConf.FILES_OPEN_COST_IN_BYTES.key -> "1048576", + CometConf.COMET_SHUFFLE_ENABLED.key -> "true", + CometConf.COMET_SHUFFLE_MODE.key -> "native") { + withTempPath { path => + // Two constant-valued files produce separate partials with zero M2. The + // old merge returns 576 instead of 1024, regardless of which partial arrives first. + for (value <- Seq(1e17 - 96, 1e17 - 32)) { + (Seq.fill(3)((0, Option(value))) ++ Seq((0, None), (1, None))) + .toDF("g", "v") + .coalesce(1) + .write + .mode("append") + .parquet(path.getCanonicalPath) + } + withParquetTable(path.getCanonicalPath, "merged_moments") { + assert(spark.table("merged_moments").rdd.getNumPartitions == 2) + for (groupBy <- Seq("", " GROUP BY g")) { + val query = "SELECT var_pop(v), var_samp(v), stddev_pop(v), stddev_samp(v) " + + "FROM merged_moments" + groupBy + val (_, cometPlan) = checkSparkAnswerAndOperator(query) + val aggregates = cometPlan.collect { case a: CometHashAggregateExec => a } + assert(aggregates.exists(_.modes.contains(Partial))) + assert(aggregates.exists(_.modes.contains(Final))) + val expected = Seq(Row(1024.0, 1228.8, 32.0, math.sqrt(1228.8))) ++ + (if (groupBy.isEmpty) Seq.empty else Seq(Row(null, null, null, null))) + checkAnswer(sql(query), expected) + checkSparkAnswerWithTolAndNumOfAggregates( + "SELECT covar_pop(v, -v), covar_samp(v, -v), corr(v, -v), regr_r2(v, -v), " + + "regr_sxx(v, -v), regr_syy(v, -v), regr_sxy(v, -v), " + + "regr_slope(v, -v), regr_intercept(v, -v) FROM merged_moments" + groupBy, + 2) + } + } + } + } + } + test("var_pop and var_samp") { withSQLConf(CometConf.COMET_SHUFFLE_ENABLED.key -> "true") { Seq("native", "jvm").foreach { cometShuffleMode =>