(spark) branch branch-4.x updated: [SPARK-58471][ML][CONNECT] Improve TrainValidationSplitModel size estimation

[ru...@apache.org]

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 9e94ef0adc33 [SPARK-58471][ML][CONNECT] Improve TrainValidationSplitModel size estimation
9e94ef0adc33 is described below

commit 9e94ef0adc33f783fb824c8f22b02648069b73d4
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Mon Aug 3 16:09:24 2026 +0800

    [SPARK-58471][ML][CONNECT] Improve TrainValidationSplitModel size estimation
    
    ### What changes were proposed in this pull request?
    
    This subtask of [SPARK-58181](https://issues.apache.org/jira/browse/SPARK-58181), and follow-up to #57665, applies parameter-aware size estimation to `TrainValidationSplitModel`.
    
    It excludes the `estimator`, `estimatorParamMaps`, and `evaluator` params from metadata sizing, then explicitly accounts for the best model, validation metrics, and collected submodels. A regression test confirms that logger-initialized estimator state is not counted.
    
    ### Why are the changes needed?
    
    TrainValidationSplitModel retains the same kinds of tuning params as CrossValidatorModel. Their params and values may retain shared runtime state, such as SparkSession, which must not contribute to Spark Connect ML model-cache accounting.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    Static checks: `git diff --check`, non-ASCII scan, and changed-Scala line-length scan.
    
    The targeted Scala suite was not run locally.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: Codex GPT-5
    
    Closes #57674 from zhengruifeng/train-validation-split-model-size-dev3.
    
    Authored-by: Ruifeng Zheng <ru...@apache.org>
    Signed-off-by: Ruifeng Zheng <ru...@apache.org>
    (cherry picked from commit 84122de296b13f146d24bca21e29301702af13ec)
    Signed-off-by: Ruifeng Zheng <ru...@apache.org>
---
 .../spark/ml/tuning/TrainValidationSplit.scala     | 25 +++++++++++++++++++++-
 .../ml/tuning/TrainValidationSplitSuite.scala      | 15 +++++++++++++
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/tuning/TrainValidationSplit.scala b/mllib/src/main/scala/org/apache/spark/ml/tuning/TrainValidationSplit.scala
index 6ee64ef99a66..5b0b77abb2dd 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/tuning/TrainValidationSplit.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/tuning/TrainValidationSplit.scala
@@ -38,8 +38,8 @@ import org.apache.spark.ml.util._
 import org.apache.spark.ml.util.Instrumentation.instrumented
 import org.apache.spark.sql.{DataFrame, Dataset}
 import org.apache.spark.sql.types.StructType
+import org.apache.spark.util.{SizeEstimator, ThreadUtils}
 import org.apache.spark.util.ArrayImplicits._
-import org.apache.spark.util.ThreadUtils
 
 /**
  * Params for [[TrainValidationSplit]] and [[TrainValidationSplitModel]].
@@ -293,6 +293,29 @@ class TrainValidationSplitModel private[ml] (
   @Since("2.3.0")
   def hasSubModels: Boolean = _subModels.isDefined
 
+  private[spark] override def estimatedSize: Long = {
+    var size = estimateMatadataSize(excluded = Seq(
+      // estimator: Param[Estimator[_]]
+      estimator,
+      // estimatorParamMaps: Param[Array[ParamMap]]
+      estimatorParamMaps,
+      // evaluator: Param[Evaluator]
+      evaluator))
+    // bestModel: Model[_]
+    size += bestModel.estimatedSize
+    // validationMetrics: Array[Double]
+    size += SizeEstimator.estimate(validationMetrics)
+    // _subModels: Option[Array[Model[_]]]
+    _subModels.foreach { modelArray =>
+      modelArray.foreach { model =>
+        if (model != null) {
+          size += model.estimatedSize
+        }
+      }
+    }
+    size
+  }
+
   @Since("2.0.0")
   override def transform(dataset: Dataset[_]): DataFrame = {
     transformSchema(dataset.schema, logging = true)
diff --git a/mllib/src/test/scala/org/apache/spark/ml/tuning/TrainValidationSplitSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/tuning/TrainValidationSplitSuite.scala
index 20ba69a5adbb..0b4b022f7e6c 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/tuning/TrainValidationSplitSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/tuning/TrainValidationSplitSuite.scala
@@ -73,6 +73,21 @@ class TrainValidationSplitSuite
     }
   }
 
+  test("TrainValidationSplitModel estimated size") {
+    val estimator = new LogisticRegression().setMaxIter(1)
+    // Initialize the estimator's logger before retaining it in the TrainValidationSplitModel.
+    estimator.fit(dataset)
+    val model = new TrainValidationSplit()
+      .setEstimator(estimator)
+      .setEstimatorParamMaps(Array(ParamMap.empty))
+      .setEvaluator(new BinaryClassificationEvaluator())
+      .fit(dataset)
+
+    val maxSize = 16 * 1024
+    assert(model.estimatedSize < maxSize,
+      s"Estimation (${model.estimatedSize}) should not include shared runtime state")
+  }
+
   test("train validation with linear regression") {
     val dataset = sc.parallelize(
       LinearDataGenerator.generateLinearInput(


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