This is an automated email from the ASF dual-hosted git repository.
uros-b pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 811e666add40 [SPARK-58550][ML] Delay GaussianMixture aggregation
allocations
811e666add40 is described below
commit 811e666add4063f45e9cf04b55015c2ee9521501
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]>
---
.../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]
