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 2d2cfc492c3c [SPARK-58550][ML] Delay GaussianMixture aggregation
allocations
2d2cfc492c3c is described below
commit 2d2cfc492c3c79c1dbdf1863bec63be764d8bdea
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Wed Aug 5 08:48:26 2026 +0000
[SPARK-58550][ML] Delay GaussianMixture aggregation allocations
### What changes were proposed in this pull request?
This PR delays allocation of `ExpectationSum` in `GaussianMixture` until an
aggregation partition receives its first point. Empty partitions and
intermediate tree branches use `null` as the aggregation identity.
### Why are the changes needed?
The previous `treeAggregate` zero value eagerly allocated the mixture
weights, means, and one dense covariance matrix per cluster. That can be
substantial for large models and is also serialized to every partition, even
empty ones.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Not run (not requested). Added coverage using an RDD with an empty
partition. Static checks: `git diff --check`, ASCII scan, and changed-line
length scan.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #57750 from zhengruifeng/gaussian-mixture-lazy-aggregate-dev3.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Uros Bojanic <[email protected]>
(cherry picked from commit 811e666add4063f45e9cf04b55015c2ee9521501)
Signed-off-by: Uros Bojanic <[email protected]>
---
.../spark/mllib/clustering/GaussianMixture.scala | 18 +++++++++++++++---
.../spark/mllib/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 562e5b3995cf..9103d83db41d 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 2ba987b96ef7..10c86472ac0d 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
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
