Skip to content

[SPARK-58281][ML][CONNECT] Avoid parent overcounting in PipelineModel size estimates #57451

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zhengruifeng wants to merge 6 commits into
base: master
Choose a base branch
from
+27 −1
Open

[SPARK-58281][ML][CONNECT] Avoid parent overcounting in PipelineModel size estimates #57451

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
524649a
[ML] Avoid parent overcounting in PipelineModel size estimates
zhengruifeng Jul 23, 2026
78c2666
[ML] Clarify PipelineModel size estimate threshold
zhengruifeng Jul 23, 2026
0a33c64
[ML] Exclude non-model pipeline stages from size estimate
zhengruifeng Jul 23, 2026
973f89e
[ML] Clarify non-model pipeline stage sizing
zhengruifeng Jul 23, 2026
16b7ab1
[ML] Strengthen PipelineModel size estimation test
zhengruifeng Jul 23, 2026
88e5126
[ML] Fix PipelineModel size estimation test input type
zhengruifeng Jul 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
18 changes: 17 additions & 1 deletion mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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._
Expand Down Expand Up @@ -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") {

@uros-b uros-b Jul 23, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new regression test covers only a single-stage pipeline whose one stage is a Model (StringIndexer). The case _ => 0L branch; the novel, behavior-changing part of this PR (the old default walk counted non-model stage bytes; the new code skips them) is never exercised. The assertion is also upper-bound-only (estimatedSize < 16 KiB), so a zero-returning implementation would pass it. Consider adding a pipeline that mixes a Model stage with a non-model Transformer stage, plus a lower-bound assertion and a hasParent-preserved check, mirroring the SPARK-57521 tests in ModelSuite (ModelSuite.scala:36-47).

@zhengruifeng zhengruifeng Jul 23, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

make sense

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]
Expand Down