(spark) branch master updated: [SPARK-58705][ML] Optimize AFTSurvivalRegressionModel transform closure

[ruifengz]

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 cc3d23166b41 [SPARK-58705][ML] Optimize AFTSurvivalRegressionModel 
transform closure
cc3d23166b41 is described below

commit cc3d23166b413e7eaec2851a203025a2e00eeed4
Author: Ruifeng Zheng <ruife...@apache.org>
AuthorDate: Wed Aug 12 10:24:03 2026 +0800

    [SPARK-58705][ML] Optimize AFTSurvivalRegressionModel transform closure
    
    ### What changes were proposed in this pull request?
    
    This PR snapshots the prediction and quantile state in local functions 
while constructing `AFTSurvivalRegressionModel.transform` UDFs. The UDF 
closures therefore retain only the required vectors and scalar rather than the 
complete model.
    
    ### Why are the changes needed?
    
    Avoiding model capture reduces transform-closure retained memory, 
particularly for Spark Connect server workloads.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    `build/sbt -java-home /usr/lib/jvm/java-17-openjdk-amd64 "mllib / Compile / 
compile"`
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: Codex (GPT-5)
    
    Closes #57918 from zhengruifeng/SPARK-58705-aft-transform-closure.
    
    Authored-by: Ruifeng Zheng <ruife...@apache.org>
    Signed-off-by: Ruifeng Zheng <ruife...@foxmail.com>
---
 .../ml/regression/AFTSurvivalRegression.scala      | 51 ++++++++++++++--------
 1 file changed, 34 insertions(+), 17 deletions(-)

diff --git 
a/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala
 
b/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala
index d573d7d524a8..0e6846e6085d 100644
--- 
a/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala
+++ 
b/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala
@@ -408,23 +408,13 @@ class AFTSurvivalRegressionModel private[ml] (
     }
   }
 
-  private def lambda2Quantiles(lambda: Double): Vector = {
-    val quantiles = _quantiles.copy
-    BLAS.scal(lambda, quantiles)
-    quantiles
-  }
-
   @Since("2.0.0")
-  def predictQuantiles(features: Vector): Vector = {
-    // scale parameter for the Weibull distribution of lifetime
-    val lambda = predict(features)
-    lambda2Quantiles(lambda)
-  }
+  def predictQuantiles(features: Vector): Vector =
+    AFTSurvivalRegressionModel.predictQuantiles(features, coefficients, 
intercept, _quantiles)
 
   @Since("2.0.0")
-  def predict(features: Vector): Double = {
-    math.exp(BLAS.dot(coefficients, features) + intercept)
-  }
+  def predict(features: Vector): Double =
+    AFTSurvivalRegressionModel.predict(features, coefficients, intercept)
 
   @Since("2.0.0")
   override def transform(dataset: Dataset[_]): DataFrame = {
@@ -432,19 +422,29 @@ class AFTSurvivalRegressionModel private[ml] (
 
     var predictionColNames = Seq.empty[String]
     var predictionColumns = Seq.empty[Column]
+    val localCoefficients = coefficients
+    val localIntercept = intercept
 
     if ($(predictionCol).nonEmpty) {
-      val predCol = udf(predict _).apply(col($(featuresCol)))
+      val predCol = udf((features: Vector) =>
+        AFTSurvivalRegressionModel.predict(features, localCoefficients, 
localIntercept))
+        .apply(col($(featuresCol)))
       predictionColNames :+= $(predictionCol)
       predictionColumns :+= predCol
         .as($(predictionCol), outputSchema($(predictionCol)).metadata)
     }
 
     if (hasQuantilesCol) {
+      val localQuantiles = _quantiles
       val quanCol = if ($(predictionCol).nonEmpty) {
-        udf(lambda2Quantiles _).apply(predictionColumns.head)
+        udf((lambda: Double) =>
+          AFTSurvivalRegressionModel.lambda2Quantiles(lambda, localQuantiles))
+          .apply(predictionColumns.head)
       } else {
-        udf(predictQuantiles _).apply(col($(featuresCol)))
+        udf((features: Vector) =>
+          AFTSurvivalRegressionModel.predictQuantiles(
+            features, localCoefficients, localIntercept, localQuantiles)
+        ).apply(col($(featuresCol)))
       }
       predictionColNames :+= $(quantilesCol)
       predictionColumns :+= quanCol
@@ -501,6 +501,23 @@ class AFTSurvivalRegressionModel private[ml] (
 object AFTSurvivalRegressionModel extends 
MLReadable[AFTSurvivalRegressionModel] {
   private[ml] case class Data(coefficients: Vector, intercept: Double, scale: 
Double)
 
+  private def lambda2Quantiles(lambda: Double, quantiles: Vector): Vector = {
+    val scaledQuantiles = quantiles.copy
+    BLAS.scal(lambda, scaledQuantiles)
+    scaledQuantiles
+  }
+
+  private def predict(features: Vector, coefficients: Vector, intercept: 
Double): Double =
+    math.exp(BLAS.dot(coefficients, features) + intercept)
+
+  private def predictQuantiles(
+      features: Vector,
+      coefficients: Vector,
+      intercept: Double,
+      quantiles: Vector): Vector = {
+    lambda2Quantiles(predict(features, coefficients, intercept), quantiles)
+  }
+
   private[ml] def serializeData(data: Data, dos: DataOutputStream): Unit = {
     import ReadWriteUtils._
     serializeVector(data.coefficients, dos)


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@spark.apache.org
For additional commands, e-mail: commits-h...@spark.apache.org