This is an automated email from the ASF dual-hosted git repository.
zhengruifeng 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 825fa296ffc4 [SPARK-58855][ML] Optimize IsotonicRegressionModel
transform closure
825fa296ffc4 is described below
commit 825fa296ffc47a6f673a5e06ac5ab3dbcc814af8
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Wed Aug 19 18:42:15 2026 +0800
[SPARK-58855][ML] Optimize IsotonicRegressionModel transform closure
### What changes were proposed in this pull request?
This PR optimizes `IsotonicRegressionModel.transform` by snapshotting the
underlying
`boundaries` and `predictions` arrays before creating the transform UDF.
The double and vector
feature transform paths now call an array-backed companion helper instead
of closing over the
legacy `mllib.regression.IsotonicRegressionModel` wrapper.
The legacy `mllib.regression.IsotonicRegressionModel` prediction path also
reuses the helper, and
its RDD prediction path snapshots the arrays before creating the map
closure.
### Why are the changes needed?
Avoiding the model capture reduces transform-closure retained memory, which
is useful for Spark
Connect server workloads and follows the closure-size optimizations tracked
under SPARK-58584.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Ran local static checks:
```
git diff --check HEAD~1 HEAD
grep -rn -P "[^\x00-\x7F]"
mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala
mllib/src/main/scala/org/apache/spark/mllib/regression/IsotonicRegression.scala
awk 'length>100 && $0 !~ /^[[:space:]]*(import|package) / && $0 !~
/https?:\/\// {print FILENAME":"FNR": "length" chars"}'
mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala
mllib/src/main/scala/org/apache/spark/mllib/regression/IsotonicRegression.scala
```
Compiled MLlib:
```
build/sbt -java-home /usr/lib/jvm/java-17-openjdk-amd64 mllib/compile
```
Measured serialized `ScalaUDF.function` closure size with a temporary local
probe:
| Case | Before | After | Delta |
|---|---:|---:|---:|
| `isotonic.double` | 3862 B | 1354 B | -2508 B, -64.9% |
| `isotonic.vector` | 3931 B | 1547 B | -2384 B, -60.6% |
No unit test was added because this is an internal refactor of existing
prediction logic.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #58101 from zhengruifeng/spark-58855-isotonic-transform-closure.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit 2b030fa7f7ecbafc69a2aa3e3480d1a30009fc92)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../spark/ml/regression/IsotonicRegression.scala | 37 ++++++++++++++++++++--
.../mllib/regression/IsotonicRegression.scala | 33 ++++---------------
2 files changed, 41 insertions(+), 29 deletions(-)
diff --git
a/mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala
b/mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala
index d6c7af65d841..72fe7706d610 100644
---
a/mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala
+++
b/mllib/src/main/scala/org/apache/spark/ml/regression/IsotonicRegression.scala
@@ -18,6 +18,7 @@
package org.apache.spark.ml.regression
import java.io.{DataInputStream, DataOutputStream}
+import java.util.Arrays.binarySearch
import org.apache.hadoop.fs.Path
@@ -261,19 +262,27 @@ class IsotonicRegressionModel private[ml] (
@Since("2.0.0")
override def transform(dataset: Dataset[_]): DataFrame = {
val outputSchema = transformSchema(dataset.schema, logging = true)
+ val localBoundaries = oldModel.boundaries
+ val localPredictions = oldModel.predictions
val predict = dataset.schema($(featuresCol)).dataType match {
case DoubleType =>
- udf { feature: Double => oldModel.predict(feature) }
+ udf { feature: Double =>
+ IsotonicRegressionModel.predict(localBoundaries, localPredictions,
feature)
+ }
case _: VectorUDT =>
val idx = $(featureIndex)
- udf { features: Vector => oldModel.predict(features(idx)) }
+ udf { features: Vector =>
+ IsotonicRegressionModel.predict(localBoundaries, localPredictions,
features(idx))
+ }
}
dataset.withColumn($(predictionCol), predict(col($(featuresCol))),
outputSchema($(predictionCol)).metadata)
}
@Since("3.0.0")
- def predict(value: Double): Double = oldModel.predict(value)
+ def predict(value: Double): Double = {
+ IsotonicRegressionModel.predict(oldModel.boundaries, oldModel.predictions,
value)
+ }
@Since("1.5.0")
override def transformSchema(schema: StructType): StructType = {
@@ -304,6 +313,28 @@ object IsotonicRegressionModel extends
MLReadable[IsotonicRegressionModel] {
predictions: Array[Double],
isotonic: Boolean)
+ private[spark] def predict(
+ boundaries: Array[Double],
+ predictions: Array[Double],
+ testData: Double): Double = {
+ val foundIndex = binarySearch(boundaries, testData)
+ val insertIndex = -foundIndex - 1
+
+ if (insertIndex == 0) {
+ predictions.head
+ } else if (insertIndex == boundaries.length) {
+ predictions.last
+ } else if (foundIndex < 0) {
+ val x1 = boundaries(insertIndex - 1)
+ val y1 = predictions(insertIndex - 1)
+ val x2 = boundaries(insertIndex)
+ val y2 = predictions(insertIndex)
+ y1 + (y2 - y1) * (testData - x1) / (x2 - x1)
+ } else {
+ predictions(foundIndex)
+ }
+ }
+
private[ml] def serializeData(data: Data, dos: DataOutputStream): Unit = {
import ReadWriteUtils._
serializeDoubleArray(data.boundaries, dos)
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 78442c5c52c0..e4aec6d57710 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
@@ -18,7 +18,6 @@ package org.apache.spark.mllib.regression
import java.io.Serializable
import java.lang.{Double => JDouble}
-import java.util.Arrays.binarySearch
import scala.collection.mutable.ArrayBuffer
import scala.jdk.CollectionConverters._
@@ -30,6 +29,7 @@ import org.json4s.jackson.JsonMethods._
import org.apache.spark.{RangePartitioner, SparkContext}
import org.apache.spark.annotation.Since
import org.apache.spark.api.java.{JavaDoubleRDD, JavaRDD}
+import org.apache.spark.ml.regression.{IsotonicRegressionModel =>
NewIsotonicRegressionModel}
import org.apache.spark.mllib.linalg.{Vector, Vectors}
import org.apache.spark.mllib.util.{Loader, Saveable}
import org.apache.spark.rdd.RDD
@@ -89,7 +89,11 @@ class IsotonicRegressionModel @Since("1.3.0") (
*/
@Since("1.3.0")
def predict(testData: RDD[Double]): RDD[Double] = {
- testData.map(predict)
+ val localBoundaries = boundaries
+ val localPredictions = predictions
+ testData.map { value =>
+ NewIsotonicRegressionModel.predict(localBoundaries, localPredictions,
value)
+ }
}
/**
@@ -124,30 +128,7 @@ class IsotonicRegressionModel @Since("1.3.0") (
*/
@Since("1.3.0")
def predict(testData: Double): Double = {
-
- def linearInterpolation(x1: Double, y1: Double, x2: Double, y2: Double, x:
Double): Double = {
- y1 + (y2 - y1) * (x - x1) / (x2 - x1)
- }
-
- val foundIndex = binarySearch(boundaries, testData)
- val insertIndex = -foundIndex - 1
-
- // Find if the index was lower than all values,
- // higher than all values, in between two values or exact match.
- if (insertIndex == 0) {
- predictions.head
- } else if (insertIndex == boundaries.length) {
- predictions.last
- } else if (foundIndex < 0) {
- linearInterpolation(
- boundaries(insertIndex - 1),
- predictions(insertIndex - 1),
- boundaries(insertIndex),
- predictions(insertIndex),
- testData)
- } else {
- predictions(foundIndex)
- }
+ NewIsotonicRegressionModel.predict(boundaries, predictions, testData)
}
/** A convenient method for boundaries called by the Python API. */
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
