You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ur...@apache.org on 2026/07/24 08:12:28 UTC
(spark) branch branch-4.x updated: [SPARK-58291][SQL] Guard statistical-aggregate merge against empty-buffer overflow to NaN
This is an automated email from the ASF dual-hosted git repository.
uros-b pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new a202b27770a4 [SPARK-58291][SQL] Guard statistical-aggregate merge against empty-buffer overflow to NaN
a202b27770a4 is described below
commit a202b27770a4e0ed28b989c695037b7a8d52e316
Author: Xiduo You <ul...@gmail.com>
AuthorDate: Fri Jul 24 08:11:58 2026 +0000
[SPARK-58291][SQL] Guard statistical-aggregate merge against empty-buffer overflow to NaN
### What changes were proposed in this pull request?
The `merge` expressions of `PearsonCorrelation` (`Corr`), `Covariance`, and `CentralMomentAgg` overflow to `Infinity` and then to `NaN` via `Infinity * 0` when one side is an empty aggregation buffer whose counterpart holds a large average.
Concretely, the merge computes terms like `delta * deltaN * n1 * n2` (and the `dx * dxN`, `dy * dyN` variants). When one buffer is empty (`n1 == 0` or `n2 == 0`), that term is mathematically zero, but `delta * deltaN` can overflow to `Infinity` *before* being multiplied by the zero count, and `Infinity * 0 = NaN` corrupts the merged moments.
This PR forces the `delta`/`dx`/`dy` terms to `0.0` when either side is empty, and sets the merged average directly to the non-empty side's average, so the empty-side terms vanish cleanly. This mirrors the guard the other aggregates already apply.
This is extracted from #57363 as a standalone correctness fix, independent of the `spark.sql.execution.replaceHashWithSortAgg` default flip, so it can be reviewed, merged, and backported on its own.
### Why are the changes needed?
`var_pop`, `covar_pop`, `regr_sxy`, `corr`, and `regr_r2` can return `NaN` for finite inputs when a partial/final aggregate merge encounters an empty buffer against a large-magnitude average. The correct results are finite (e.g. zero variance for equal values). The bug is only masked when adjacent aggregates are combined into `Complete` mode (which skips the merge), so the same query returns different results depending on whether the merge path runs.
### Does this PR introduce _any_ user-facing change?
Yes. Queries using `corr`, `covar_pop`/`covar_samp`, `var_pop`/`var_samp`, `stddev_*`, `skewness`, `kurtosis`, and the `regr_*` family that previously returned `NaN` for large finite inputs now return the correct finite result.
### How was this patch tested?
Two new regression tests in `DataFrameAggregateSuite`:
- `SPARK-58291`-covered `var_pop`/`covar_pop`/`regr_sxy` on two equal large values (expect `0.0`).
- `corr`/`regr_r2` on two large-magnitude finite points (expect a finite result, not `NaN`).
Both assert results are stable across AQE on/off and `combineAdjacentAggregation` on/off, and both fail without the fix on the merge path. `catalyst/compile` and the new tests pass.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 4.8)
Closes #57460 from ulysses-you/agg-merge-overflow-fix.
Authored-by: Xiduo You <ul...@gmail.com>
Signed-off-by: Uros Bojanic <22...@users.noreply.github.com>
(cherry picked from commit 6158bae0b016b6adc5f88bae70edd57a745072ae)
Signed-off-by: Uros Bojanic <22...@users.noreply.github.com>
---
.../expressions/aggregate/CentralMomentAgg.scala | 11 +++-
.../sql/catalyst/expressions/aggregate/Corr.scala | 13 ++--
.../expressions/aggregate/Covariance.scala | 13 ++--
.../aggregate/MergeEmptyBufferSuite.scala | 71 ++++++++++++++++++++++
.../apache/spark/sql/DataFrameAggregateSuite.scala | 53 ++++++++++++++++
5 files changed, 151 insertions(+), 10 deletions(-)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala
index eb4f2c119152..81da51886e22 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/CentralMomentAgg.scala
@@ -85,9 +85,16 @@ abstract class CentralMomentAgg(child: Expression, nullOnDivideByZero: Boolean)
val n1 = n.left
val n2 = n.right
val newN = n1 + n2
- val delta = avg.right - avg.left
+ // When one side is an empty buffer (n1 == 0 or n2 == 0), the combined buffer is just the other
+ // side. Force `delta` to 0.0 in that case so that the `delta * deltaN * n1 * n2` terms vanish
+ // cleanly; otherwise a large `avg` difference can overflow `delta * deltaN` to Infinity before
+ // it is multiplied by the zero count, producing NaN (e.g. `Infinity * 0 = NaN`) and corrupting
+ // the merged moments. `newAvg` is set directly to the non-empty side's average.
+ val isEmptyLeft = n1 === Literal(0.0)
+ val isEmptyRight = n2 === Literal(0.0)
+ val delta = If(isEmptyLeft || isEmptyRight, Literal(0.0), avg.right - avg.left)
val deltaN = If(newN === 0.0, 0.0, delta / newN)
- val newAvg = avg.left + deltaN * n2
+ val newAvg = If(isEmptyLeft, avg.right, If(isEmptyRight, avg.left, avg.left + deltaN * n2))
// higher order moments computed according to:
// https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Higher-order_statistics
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Corr.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Corr.scala
index 72f682c00b4a..bc78dfdf8cec 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Corr.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Corr.scala
@@ -64,12 +64,17 @@ abstract class PearsonCorrelation(x: Expression, y: Expression, nullOnDivideByZe
val n1 = n.left
val n2 = n.right
val newN = n1 + n2
- val dx = xAvg.right - xAvg.left
+ // See CentralMomentAgg for the rationale: when one side is an empty buffer (n1 == 0 or
+ // n2 == 0), force dx/dy to 0.0 so the `dx * dyN * n1 * n2` (and `dx * dxN`, `dy * dyN`) terms
+ // vanish cleanly instead of overflowing to Infinity and then NaN-ing out via `Infinity * 0`.
+ val isEmptyLeft = n1 === Literal(0.0)
+ val isEmptyRight = n2 === Literal(0.0)
+ val dx = If(isEmptyLeft || isEmptyRight, Literal(0.0), xAvg.right - xAvg.left)
+ val dy = If(isEmptyLeft || isEmptyRight, Literal(0.0), yAvg.right - yAvg.left)
val dxN = If(newN === 0.0, 0.0, dx / newN)
- val dy = yAvg.right - yAvg.left
val dyN = If(newN === 0.0, 0.0, dy / newN)
- val newXAvg = xAvg.left + dxN * n2
- val newYAvg = yAvg.left + dyN * n2
+ val newXAvg = If(isEmptyLeft, xAvg.right, If(isEmptyRight, xAvg.left, xAvg.left + dxN * n2))
+ val newYAvg = If(isEmptyLeft, yAvg.right, If(isEmptyRight, yAvg.left, yAvg.left + dyN * n2))
val newCk = ck.left + ck.right + dx * dyN * n1 * n2
val newXMk = xMk.left + xMk.right + dx * dxN * n1 * n2
val newYMk = yMk.left + yMk.right + dy * dyN * n1 * n2
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Covariance.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Covariance.scala
index 7d047aca0606..9fa2ac3a4d5b 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Covariance.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/aggregate/Covariance.scala
@@ -58,12 +58,17 @@ abstract class Covariance(val left: Expression, val right: Expression, nullOnDiv
val n1 = n.left
val n2 = n.right
val newN = n1 + n2
- val dx = xAvg.right - xAvg.left
+ // See CentralMomentAgg for the rationale: when one side is an empty buffer (n1 == 0 or
+ // n2 == 0), force dx/dy to 0.0 so the `dx * dyN * n1 * n2` term vanishes cleanly instead of
+ // overflowing to Infinity and then NaN-ing out via `Infinity * 0`.
+ val isEmptyLeft = n1 === Literal(0.0)
+ val isEmptyRight = n2 === Literal(0.0)
+ val dx = If(isEmptyLeft || isEmptyRight, Literal(0.0), xAvg.right - xAvg.left)
+ val dy = If(isEmptyLeft || isEmptyRight, Literal(0.0), yAvg.right - yAvg.left)
val dxN = If(newN === 0.0, 0.0, dx / newN)
- val dy = yAvg.right - yAvg.left
val dyN = If(newN === 0.0, 0.0, dy / newN)
- val newXAvg = xAvg.left + dxN * n2
- val newYAvg = yAvg.left + dyN * n2
+ val newXAvg = If(isEmptyLeft, xAvg.right, If(isEmptyRight, xAvg.left, xAvg.left + dxN * n2))
+ val newYAvg = If(isEmptyLeft, yAvg.right, If(isEmptyRight, yAvg.left, yAvg.left + dyN * n2))
val newCk = ck.left + ck.right + dx * dyN * n1 * n2
Seq(newN, newXAvg, newYAvg, newCk)
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/aggregate/MergeEmptyBufferSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/aggregate/MergeEmptyBufferSuite.scala
new file mode 100644
index 000000000000..319503e1d70c
--- /dev/null
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/aggregate/MergeEmptyBufferSuite.scala
@@ -0,0 +1,71 @@
+/*
+ * 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.
+ */
+package org.apache.spark.sql.catalyst.expressions.aggregate
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference, JoinedRow}
+import org.apache.spark.sql.types.DoubleType
+
+/**
+ * Focused regression coverage for the empty-buffer guards in the `merge` expressions of the
+ * statistical aggregates. Merging an empty buffer into a populated one (or vice versa) is an
+ * identity, but without the guards a large average makes the intermediate `delta * deltaN` term
+ * overflow to `Infinity` before it is multiplied by the zero count, and `Infinity * 0 = NaN`
+ * corrupts the moments. The end-to-end `DataFrameAggregateSuite` cases use `repartition(1)`, which
+ * only produces the empty-left merge; here both the `isEmptyLeft` and `isEmptyRight` branches are
+ * exercised directly by driving the merger projection with the empty buffer on each side, under
+ * both interpreted and generated execution.
+ */
+class MergeEmptyBufferSuite extends TestWithAndWithoutCodegen {
+ private val x = AttributeReference("x", DoubleType, nullable = true)()
+ private val y = AttributeReference("y", DoubleType, nullable = true)()
+
+ // Merge the empty buffer on each side and assert the populated buffer is returned unchanged (an
+ // empty-buffer merge is an identity), which also proves no term overflowed to Infinity/NaN.
+ private def checkMergeWithEmptyIsIdentity(
+ evaluator: DeclarativeAggregateEvaluator, populated: InternalRow): Unit = {
+ val empty = evaluator.initialize()
+ val joiner = new JoinedRow
+ // isEmptyLeft: empty buffer as the left (accumulator) operand.
+ assert(evaluator.merger(joiner(empty, populated)).copy() === populated)
+ // isEmptyRight: empty buffer as the right (input) operand.
+ assert(evaluator.merger(joiner(populated, empty)).copy() === populated)
+ }
+
+ testBothCodegenAndInterpreted("SPARK-58291: var_pop empty-buffer merge is overflow-safe") {
+ val evaluator = DeclarativeAggregateEvaluator(VariancePop(x, nullOnDivideByZero = true), Seq(x))
+ // Two equal, very large finite values: a populated buffer with a huge average but zero
+ // variance, which is exactly the shape that overflowed `delta * deltaN`.
+ val populated = evaluator.update(InternalRow(1.0e155), InternalRow(1.0e155))
+ checkMergeWithEmptyIsIdentity(evaluator, populated)
+ }
+
+ testBothCodegenAndInterpreted("SPARK-58291: covar_pop empty-buffer merge is overflow-safe") {
+ val evaluator =
+ DeclarativeAggregateEvaluator(CovPopulation(x, y, nullOnDivideByZero = true), Seq(x, y))
+ val populated = evaluator.update(InternalRow(1.0e155, 1.0e155), InternalRow(1.0e155, 1.0e155))
+ checkMergeWithEmptyIsIdentity(evaluator, populated)
+ }
+
+ testBothCodegenAndInterpreted("SPARK-58291: corr empty-buffer merge is overflow-safe") {
+ val evaluator = DeclarativeAggregateEvaluator(Corr(x, y, nullOnDivideByZero = true), Seq(x, y))
+ val populated = evaluator.update(
+ InternalRow(1.0e155, 1.0e-150),
+ InternalRow(1.000000000000001e155, 2.0e-150))
+ checkMergeWithEmptyIsIdentity(evaluator, populated)
+ }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
index 36c2a9540109..32517af68c91 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala
@@ -589,6 +589,59 @@ class DataFrameAggregateSuite extends SharedSparkSession
Row(null, null, null, null, null))
}
+ test("SPARK-58291: empty-buffer merge must not overflow to NaN for statistical aggregates") {
+ // A single-partition group with two equal, very large finite values has zero variance, so
+ // var_pop / covar_pop / regr_sxy must be 0.0. Previously, when adjacent Partial/Final
+ // aggregates were NOT combined (the old default), the Final merge of the non-empty Partial
+ // buffer into the empty Final buffer computed `delta * deltaN * n1 * n2` where `n1 == 0`;
+ // `delta * deltaN` overflowed to Infinity and `Infinity * 0 = NaN`, corrupting the moments.
+ // CombineAdjacentAggregation (Complete mode) sidesteps the merge and returned 0.0, so the two
+ // configurations disagreed. The merge fix makes both paths return 0.0.
+ // This must hold with and without AQE, and with combining on and off.
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val df = Seq(1e155, 1e155).toDF("a").repartition(1)
+ Seq(true, false).foreach { combine =>
+ withSQLConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> combine.toString) {
+ checkAnswer(
+ df.selectExpr("var_pop(a)", "covar_pop(a, a)", "regr_sxy(a, a)"),
+ Row(0.0, 0.0, 0.0))
+ }
+ }
+ }
+ }
+ }
+
+ test("SPARK-58291: empty-buffer merge must not overflow to NaN for Pearson correlation") {
+ // Two finite points are perfectly linearly correlated, so corr / regr_r2 are finite. The
+ // Pearson merge computes `dx * dxN * n1 * n2` (and the dy variant) for xMk / yMk. When one
+ // side is an empty buffer (n1 == 0 or n2 == 0) and the other has a large average, `dx * dxN`
+ // overflows to Infinity before being multiplied by the zero count, and `Infinity * 0 = NaN`
+ // corrupts the merged moments. The old default (no combining) hit this Final-merge path and
+ // returned NaN, while CombineAdjacentAggregation (Complete mode) sidestepped the merge, so the
+ // two configurations disagreed. The merge fix makes both paths return the same finite result.
+ // This must hold with and without AQE, and with combining on and off.
+ val data = Seq((1.0e155, 1.0e-150), (1.000000000000001e155, 2.0e-150))
+ Seq(true, false).foreach { aqe =>
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString) {
+ val df = data.toDF("x", "y").repartition(1)
+ // Reference result from the combined (Complete-mode) path, which never runs the merge.
+ val expected = withSQLConf(
+ SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> "true") {
+ df.selectExpr("corr(x, y)", "regr_r2(y, x)").collect()
+ }
+ expected.head.toSeq.foreach { v =>
+ assert(!v.asInstanceOf[Double].isNaN, "corr / regr_r2 must be finite, not NaN")
+ }
+ Seq(true, false).foreach { combine =>
+ withSQLConf(SQLConf.COMBINE_ADJACENT_AGGREGATION_ENABLED.key -> combine.toString) {
+ checkAnswer(df.selectExpr("corr(x, y)", "regr_r2(y, x)"), expected.toSeq)
+ }
+ }
+ }
+ }
+ }
+
test("collect functions") {
val df = Seq((1, 2), (2, 2), (3, 4)).toDF("a", "b")
checkAnswer(
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org