This is an automated email from the ASF dual-hosted git repository.
zhengruifeng 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 80afd1d95b4e [SPARK-58587][ML] Move final IsotonicRegression PAV pass
to executor
80afd1d95b4e is described below
commit 80afd1d95b4ea1256301a6f6b5e46e74ec7b395a
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Thu Aug 6 08:09:38 2026 +0800
[SPARK-58587][ML] Move final IsotonicRegression PAV pass to executor
### What changes were proposed in this pull request?
This PR moves the final Pool Adjacent Violators (PAV) pass in
`IsotonicRegression` from the
driver to a single executor partition. It replaces the driver-side
in-memory sort of all partial
results with `sortBy(..., numPartitions = 1)`, allowing Spark shuffle sort
to spill while ordering
partial results. The driver collects only the final PAV output.
The PR also adds regression coverage for a monotonicity violation that
spans two input
partitions.
### Why are the changes needed?
The final PAV pass can substantially compress partial results. Running it
before `collect()`
reduces driver memory and data transfer when that compression occurs. The
final PAV pass remains
single-partition, but the preceding sort no longer materializes every
partial result on the driver.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Added a regression test for merging partial PAV results across partitions.
Static checks completed successfully:
- `git diff --check`
- ASCII scan of the changed Scala files
- changed-line length scan
The test suite was not run because it was not requested.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #57788 from zhengruifeng/ml_iso_fit.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../org/apache/spark/mllib/regression/IsotonicRegression.scala | 9 +++++----
.../apache/spark/mllib/regression/IsotonicRegressionSuite.scala | 7 +++++++
2 files changed, 12 insertions(+), 4 deletions(-)
diff --git
a/mllib/src/main/scala/org/apache/spark/mllib/regression/IsotonicRegression.scala
b/mllib/src/main/scala/org/apache/spark/mllib/regression/IsotonicRegression.scala
index 456580ffa531..78442c5c52c0 100644
---
a/mllib/src/main/scala/org/apache/spark/mllib/regression/IsotonicRegression.scala
+++
b/mllib/src/main/scala/org/apache/spark/mllib/regression/IsotonicRegression.scala
@@ -477,7 +477,8 @@ class IsotonicRegression private (private var isotonic:
Boolean) extends Seriali
private def parallelPoolAdjacentViolators(
input: RDD[(Double, Double, Double)]): Array[(Double, Double, Double)] =
{
val keyedInput = input.keyBy(_._2)
- val parallelStepResult = keyedInput
+
+ keyedInput
// Points with same or adjacent features must collocate within the same
partition.
.partitionBy(new RangePartitioner(keyedInput.getNumPartitions,
keyedInput))
.values
@@ -486,10 +487,10 @@ class IsotonicRegression private (private var isotonic:
Boolean) extends Seriali
// Aggregate points with equal features into a single point.
.map(makeUnique)
.flatMap(poolAdjacentViolators)
+ // Sort partial results with a spill-capable shuffle before the final
PAV pass.
+ .sortBy(_._2, ascending = true, numPartitions = 1)
+ .mapPartitions(p => poolAdjacentViolators(p.toArray).iterator)
.collect()
- // Sort again because collect() doesn't promise ordering.
- .sortBy(_._2)
- poolAdjacentViolators(parallelStepResult)
}
/**
diff --git
a/mllib/src/test/scala/org/apache/spark/mllib/regression/IsotonicRegressionSuite.scala
b/mllib/src/test/scala/org/apache/spark/mllib/regression/IsotonicRegressionSuite.scala
index 0a4f81bd7aa9..2500a52ac632 100644
---
a/mllib/src/test/scala/org/apache/spark/mllib/regression/IsotonicRegressionSuite.scala
+++
b/mllib/src/test/scala/org/apache/spark/mllib/regression/IsotonicRegressionSuite.scala
@@ -176,6 +176,13 @@ class IsotonicRegressionSuite extends SparkFunSuite with
MLlibTestSparkContext w
assert(model.predictions === Array(1, 2, 3, 4, 5))
}
+ test("isotonic regression merges partial results across partitions") {
+ val model = runIsotonicRegressionOnInput(generateIsotonicInput(Seq(1, 3,
2, 4)), true, 2)
+
+ assert(model.boundaries === Array(0, 1, 2, 3))
+ assert(model.predictions === Array(1, 2.5, 2.5, 4))
+ }
+
test("weighted isotonic regression") {
val model = runIsotonicRegression(Seq(1, 2, 3, 4, 2), Seq(1, 1, 1, 1, 2),
true)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]