(spark) branch master updated: [SPARK-58854][ML] Optimize GeneralizedLinearRegressionModel transform closure

[ru...@apache.org]

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 1f0641f4bc92 [SPARK-58854][ML] Optimize GeneralizedLinearRegressionModel transform closure
1f0641f4bc92 is described below

commit 1f0641f4bc92a12ac983d6282a82d67ea36d59a9
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Wed Aug 19 18:40:58 2026 +0800

    [SPARK-58854][ML] Optimize GeneralizedLinearRegressionModel transform closure
    
    ### What changes were proposed in this pull request?
    
    This PR optimizes `GeneralizedLinearRegressionModel.transform` by snapshotting the prediction
    state before creating transform UDFs. The UDFs now capture only the required coefficients,
    intercept, and `FamilyAndLink` instead of the full model.
    
    The public prediction methods continue to use the same prediction logic through companion helper
    methods.
    
    ### 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
    grep -rn -P "[^\x00-\x7F]" mllib/src/main/scala/org/apache/spark/ml/regression/GeneralizedLinearRegression.scala
    awk 'length>100 && $0 !~ /^[[:space:]]*(import|package) / && $0 !~ /https?:\/\// {print FILENAME":"FNR": "length" chars"}' mllib/src/main/scala/org/apache/spark/ml/regression/GeneralizedLinearRegression.scala
    ```
    
    Measured serialized `ScalaUDF.function` closure size with a temporary local probe:
    
    | Case | Before | After | Delta |
    |---|---:|---:|---:|
    | `directPrediction` | 6753 B | 3033 B | -3720 B, -55.1% |
    | `linkAndPrediction` total | 13450 B | 9559 B | -3891 B, -28.9% |
    
    For `linkAndPrediction`, the two UDF closures measured individually as:
    
    | UDF | Before | After | Delta |
    |---:|---:|---:|---:|
    | `udf#0` | 6669 B | 7482 B | +813 B, +12.2% |
    | `udf#1` | 6781 B | 2077 B | -4704 B, -69.4% |
    
    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 #58100 from zhengruifeng/spark-58854-glr-transform-closure.
    
    Authored-by: Ruifeng Zheng <ru...@apache.org>
    Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
---
 .../regression/GeneralizedLinearRegression.scala   | 37 ++++++++++++++++++----
 1 file changed, 31 insertions(+), 6 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/GeneralizedLinearRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/GeneralizedLinearRegression.scala
index f4d6e4e3562a..de10d1e7aed7 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/GeneralizedLinearRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/GeneralizedLinearRegression.scala
@@ -1032,22 +1032,22 @@ class GeneralizedLinearRegressionModel private[ml] (
   private lazy val familyAndLink = FamilyAndLink(this)
 
   override def predict(features: Vector): Double = {
-    predict(features, 0.0)
+    GeneralizedLinearRegressionModel.predict(features, 0.0, coefficients, intercept, familyAndLink)
   }
 
   /**
    * Calculates the predicted value when offset is set.
    */
   private def predict(features: Vector, offset: Double): Double = {
-    val eta = predictLink(features, offset)
-    familyAndLink.fitted(eta)
+    GeneralizedLinearRegressionModel.predict(
+      features, offset, coefficients, intercept, familyAndLink)
   }
 
   /**
    * Calculates the link prediction (linear predictor) of the given instance.
    */
   private def predictLink(features: Vector, offset: Double): Double = {
-    BLAS.dot(features, coefficients) + intercept + offset
+    GeneralizedLinearRegressionModel.predictLink(features, offset, coefficients, intercept)
   }
 
   override def transform(dataset: Dataset[_]): DataFrame = {
@@ -1061,9 +1061,14 @@ class GeneralizedLinearRegressionModel private[ml] (
     val offset = if (!hasOffsetCol) lit(0.0) else col($(offsetCol)).cast(DoubleType)
     var outputData = dataset
     var numColsOutput = 0
+    val localCoefficients = coefficients
+    val localIntercept = intercept
+    val localFamilyAndLink = familyAndLink
 
     if (hasLinkPredictionCol) {
-      val predLinkUDF = udf((features: Vector, offset: Double) => predictLink(features, offset))
+      val predLinkUDF = udf((features: Vector, offset: Double) =>
+        GeneralizedLinearRegressionModel.predictLink(
+          features, offset, localCoefficients, localIntercept))
       outputData = outputData
         .withColumn($(linkPredictionCol), predLinkUDF(col($(featuresCol)), offset),
           outputSchema($(linkPredictionCol)).metadata)
@@ -1076,7 +1081,9 @@ class GeneralizedLinearRegressionModel private[ml] (
         outputData = outputData.withColumn($(predictionCol), predUDF(col($(linkPredictionCol))),
           outputSchema($(predictionCol)).metadata)
       } else {
-        val predUDF = udf((features: Vector, offset: Double) => predict(features, offset))
+        val predUDF = udf((features: Vector, offset: Double) =>
+          GeneralizedLinearRegressionModel.predict(
+            features, offset, localCoefficients, localIntercept, localFamilyAndLink))
         outputData = outputData.withColumn($(predictionCol), predUDF(col($(featuresCol)), offset),
           outputSchema($(predictionCol)).metadata)
       }
@@ -1177,6 +1184,24 @@ class GeneralizedLinearRegressionModel private[ml] (
 object GeneralizedLinearRegressionModel extends MLReadable[GeneralizedLinearRegressionModel] {
   private[ml] case class Data(intercept: Double, coefficients: Vector)
 
+  private def predict(
+      features: Vector,
+      offset: Double,
+      coefficients: Vector,
+      intercept: Double,
+      familyAndLink: GeneralizedLinearRegression.FamilyAndLink): Double = {
+    val eta = predictLink(features, offset, coefficients, intercept)
+    familyAndLink.fitted(eta)
+  }
+
+  private def predictLink(
+      features: Vector,
+      offset: Double,
+      coefficients: Vector,
+      intercept: Double): Double = {
+    BLAS.dot(features, coefficients) + intercept + offset
+  }
+
   private[ml] def serializeData(data: Data, dos: DataOutputStream): Unit = {
     import ReadWriteUtils._
     dos.writeDouble(data.intercept)


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