[SPARK-58291][SQL] Guard statistical-aggregate merge against empty-buffer overflow to NaN

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

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

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)

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)
+  }
+}

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(