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 f7944cb728e8 [SPARK-58281][ML][CONNECT] Avoid parent overcounting in
PipelineModel size estimates
f7944cb728e8 is described below
commit f7944cb728e82c05bf55ab8bd4d3eb6ed0e6fc70
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Fri Jul 24 14:51:22 2026 +0800
[SPARK-58281][ML][CONNECT] Avoid parent overcounting in PipelineModel size
estimates
### What changes were proposed in this pull request?
Adds a `PipelineModel.estimatedSize` override that sums model-stage
estimates individually and directly estimates non-model transformer stages.
This prevents a reflection walk of the complete pipeline object graph.
Adds a regression test that fits a pipeline containing `StringIndexer` and
asserts its size estimate remains below 16 KiB.
### Why are the changes needed?
SPARK-57521 clears the top-level model parent before the default size walk.
A copied `PipelineModel` still contains copied model stages whose estimator
parents are retained, so a reflection walk can still reach shared Spark-session
state. Estimating nested model stages independently applies their
parentless-copy behavior at every model boundary.
### Does this PR introduce _any_ user-facing change?
Yes. Pipeline model cache-size estimates no longer include shared state
reachable through nested stage parents, preventing unnecessary cache
overcounting.
### How was this patch tested?
- Added a fitted `StringIndexer` pipeline size-estimation regression test.
- `JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64 build/sbt
'mllib/Test/compile'`\n\n### Was this patch authored or co-authored using
generative AI tooling?\n\nGenerated-by: Codex GPT-5
Closes #57451 from zhengruifeng/pipeline-model-size-estimate-dev3.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit f8836939c6a89bd788ad648ec30df757cde204ac)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../src/main/scala/org/apache/spark/ml/Pipeline.scala | 10 ++++++++++
.../test/scala/org/apache/spark/ml/PipelineSuite.scala | 18 +++++++++++++++++-
2 files changed, 27 insertions(+), 1 deletion(-)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala
b/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala
index 518e00af306c..92b1e8801fdc 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala
@@ -322,6 +322,16 @@ class PipelineModel private[ml] (
@Since("1.4.0") val stages: Array[Transformer])
extends Model[PipelineModel] with MLWritable with Logging {
+ private[spark] override def estimatedSize: Long = {
+ estimateMatadataSize + stages.iterator.map {
+ case model: Model[_] => model.estimatedSize
+ // Non-model transformers are skipped because:
+ // - ML Connect cache accounting counts only models.
+ // - They can retain incidental references, such as SparkSession.
+ case _ => 0L
+ }.sum
+ }
+
/** A Java/Python-friendly auxiliary constructor. */
private[ml] def this(uid: String, stages: ju.List[Transformer]) = {
this(uid, stages.asScala.toArray)
diff --git a/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala
b/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala
index 3bf339f1fc99..d2ad51766f04 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala
@@ -26,7 +26,7 @@ import org.scalatestplus.mockito.MockitoSugar.mock
import org.apache.spark.SparkFunSuite
import org.apache.spark.ml.Pipeline.SharedReadWrite
-import org.apache.spark.ml.feature.{HashingTF, MinMaxScaler}
+import org.apache.spark.ml.feature.{HashingTF, MinMaxScaler, StringIndexer,
Tokenizer}
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.ml.param.{IntParam, ParamMap}
import org.apache.spark.ml.util._
@@ -128,6 +128,22 @@ class PipelineSuite extends SparkFunSuite with
MLlibTestSparkContext with Defaul
"copy should create an instance with the same parent")
}
+ test("PipelineModel estimated size") {
+ val dataset = Seq("a", "b", "c").toDF("input")
+ val pipeline = new Pipeline().setStages(Array(
+ new Tokenizer().setInputCol("input").setOutputCol("tokens"),
+ new StringIndexer().setInputCol("input").setOutputCol("indexed")))
+ val model = pipeline.fit(dataset)
+ val minSize = 1024
+ val maxSize = 1024 * 16
+ val estimatedSize = model.estimatedSize
+ assert(estimatedSize > minSize,
+ s"Estimation ($estimatedSize) should be greater than $minSize")
+ assert(estimatedSize < maxSize,
+ s"Estimation ($estimatedSize) should be less than $maxSize")
+ assert(model.hasParent, "parent should be preserved after estimatedSize
call")
+ }
+
test("pipeline model constructors") {
val transform0 = mock[Transformer]
val model1 = mock[MyModel]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
