From 384220164f60fee7eedd3e2324f276c2ede40de5 Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Thu, 23 Jul 2026 18:30:27 +0800 Subject: [PATCH 1/2] [SPARK-58291][SQL] Guard statistical-aggregate merge against empty-buffer overflow The Pearson correlation, covariance, and central-moment merge expressions overflow to Infinity and then NaN via `Infinity * 0` when one side is an empty buffer with a large average: `delta * deltaN` (and the `dx`/`dy` variants) overflow before being multiplied by the zero count. Force the delta terms to 0.0 and set the merged average directly to the non-empty side when either buffer is empty, mirroring the guard each other aggregate already applies. Extracted from #57363 as a standalone correctness fix, independent of the replaceHashWithSortAgg default flip, so it can be backported on its own. Co-Authored-By: Claude Opus 4.8 --- .../aggregate/CentralMomentAgg.scala | 11 +++- .../catalyst/expressions/aggregate/Corr.scala | 13 +++-- .../expressions/aggregate/Covariance.scala | 13 +++-- .../spark/sql/DataFrameAggregateSuite.scala | 53 +++++++++++++++++++ 4 files changed, 80 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 eb4f2c1191523..81da51886e225 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 72f682c00b4ac..bc78dfdf8cec1 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 7d047aca06068..9fa2ac3a4d5b1 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/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameAggregateSuite.scala index 36c2a95401090..32517af68c916 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( From 199c3e9da0da101293ca2c7553578b153c63bf8d Mon Sep 17 00:00:00 2001 From: Xiduo You Date: Fri, 24 Jul 2026 09:25:26 +0800 Subject: [PATCH 2/2] [SPARK-58291][SQL][TESTS] Add focused merge coverage for both empty-buffer directions Address review (sunchao): the end-to-end DataFrameAggregateSuite cases use repartition(1), which only produces the empty-left merge, so the new isEmptyRight branch of the guard was not independently exercised. Add a Catalyst-level MergeEmptyBufferSuite that drives the merger projection with the empty buffer on each side, for var_pop / covar_pop / corr, under both interpreted and generated execution. Mutation-tested: reverting the guard fails every case; isolating the isEmptyRight assertion alone still fails on the reverted code, confirming the branch is covered. Co-Authored-By: Claude Opus 4.8 --- .../aggregate/MergeEmptyBufferSuite.scala | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) create mode 100644 sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/aggregate/MergeEmptyBufferSuite.scala 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 0000000000000..319503e1d70cd --- /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) + } +}