From 4df51f5a942c8f479e9dc86b60d041aaea8880e2 Mon Sep 17 00:00:00 2001
From: Ruifeng Zheng <ruifengz@apache.org>
Date: Tue, 4 Aug 2026 08:43:02 +0000
Subject: [PATCH] [SPARK-58550][MLLIB] Delay GaussianMixture aggregation
 allocations

---
 .../mllib/clustering/GaussianMixture.scala     | 18 +++++++++++++++---
 .../clustering/GaussianMixtureSuite.scala      |  4 ++--
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/clustering/GaussianMixture.scala b/mllib/src/main/scala/org/apache/spark/mllib/clustering/GaussianMixture.scala
index 562e5b3995cfb..9103d83db41dc 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/clustering/GaussianMixture.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/clustering/GaussianMixture.scala
@@ -201,10 +201,22 @@ class GaussianMixture private (
       val compute = sc.broadcast(ExpectationSum.add(weights, gaussians)_)
 
       // aggregate the cluster contribution for all sample points
+      // Avoid allocating and serializing a large zero value for empty partitions.
       val sums = breezeData.treeAggregate[ExpectationSum](
-        zeroValue = ExpectationSum.zero(k, d),
-        seqOp = (agg: ExpectationSum, v: BV[Double]) => compute.value(agg, v),
-        combOp = (agg1: ExpectationSum, agg2: ExpectationSum) => agg1 += agg2,
+        zeroValue = null.asInstanceOf[ExpectationSum],
+        seqOp = (maybeAgg: ExpectationSum, v: BV[Double]) => {
+          val agg = if (maybeAgg == null) ExpectationSum.zero(k, d) else maybeAgg
+          compute.value(agg, v)
+        },
+        combOp = (agg1: ExpectationSum, agg2: ExpectationSum) => {
+          if (agg1 == null) {
+            agg2
+          } else if (agg2 == null) {
+            agg1
+          } else {
+            agg1 += agg2
+          }
+        },
         depth = 2,
         finalAggregateOnExecutor = true)
 
diff --git a/mllib/src/test/scala/org/apache/spark/mllib/clustering/GaussianMixtureSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/clustering/GaussianMixtureSuite.scala
index 2ba987b96ef79..10c86472ac0dd 100644
--- a/mllib/src/test/scala/org/apache/spark/mllib/clustering/GaussianMixtureSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/mllib/clustering/GaussianMixtureSuite.scala
@@ -40,12 +40,12 @@ class GaussianMixtureSuite extends SparkFunSuite with MLlibTestSparkContext {
     }
   }
 
-  test("single cluster") {
+  test("single cluster with empty partitions") {
     val data = sc.parallelize(Seq(
       Vectors.dense(6.0, 9.0),
       Vectors.dense(5.0, 10.0),
       Vectors.dense(4.0, 11.0)
-    ))
+    ), 4)
 
     // expectations
     val Ew = 1.0