From 524649a1c3ad7695c3891950ce44906c16fdec61 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Thu, 23 Jul 2026 02:02:11 +0000 Subject: [PATCH 1/6] [ML] Avoid parent overcounting in PipelineModel size estimates --- .../main/scala/org/apache/spark/ml/Pipeline.scala | 8 ++++++++ .../scala/org/apache/spark/ml/PipelineSuite.scala | 12 +++++++++++- 2 files changed, 19 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 518e00af306c8..80c0973d2cb1b 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala @@ -35,6 +35,7 @@ import org.apache.spark.ml.util.Instrumentation.instrumented import org.apache.spark.sql.{DataFrame, Dataset, SparkSession} import org.apache.spark.sql.types.StructType import org.apache.spark.util.ArrayImplicits._ +import org.apache.spark.util.SizeEstimator /** * A stage in a pipeline, either an [[Estimator]] or a [[Transformer]]. @@ -322,6 +323,13 @@ 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 + case stage => SizeEstimator.estimate(stage) + }.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 3bf339f1fc994..58b8c86d420f9 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} import org.apache.spark.ml.linalg.Vectors import org.apache.spark.ml.param.{IntParam, ParamMap} import org.apache.spark.ml.util._ @@ -128,6 +128,16 @@ 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 StringIndexer().setInputCol("input").setOutputCol("indexed"))) + val model = pipeline.fit(dataset) + val maxSize = 16384 + assert(model.estimatedSize < maxSize, + s"Estimation (${model.estimatedSize}) should be less than $maxSize") + } + test("pipeline model constructors") { val transform0 = mock[Transformer] val model1 = mock[MyModel] From 78c26666e096452692bc4bc8c7a6ee0b2678302d Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Thu, 23 Jul 2026 02:03:52 +0000 Subject: [PATCH 2/6] [ML] Clarify PipelineModel size estimate threshold --- mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 58b8c86d420f9..e4877b1e0851e 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala @@ -133,7 +133,7 @@ class PipelineSuite extends SparkFunSuite with MLlibTestSparkContext with Defaul val pipeline = new Pipeline().setStages(Array( new StringIndexer().setInputCol("input").setOutputCol("indexed"))) val model = pipeline.fit(dataset) - val maxSize = 16384 + val maxSize = 1024 * 16 assert(model.estimatedSize < maxSize, s"Estimation (${model.estimatedSize}) should be less than $maxSize") } From 0a33c643cb79598a7d2e7211227d477746211e2d Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Thu, 23 Jul 2026 02:49:00 +0000 Subject: [PATCH 3/6] [ML] Exclude non-model pipeline stages from size estimate --- mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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 80c0973d2cb1b..d61d696afd7bf 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala @@ -35,7 +35,6 @@ import org.apache.spark.ml.util.Instrumentation.instrumented import org.apache.spark.sql.{DataFrame, Dataset, SparkSession} import org.apache.spark.sql.types.StructType import org.apache.spark.util.ArrayImplicits._ -import org.apache.spark.util.SizeEstimator /** * A stage in a pipeline, either an [[Estimator]] or a [[Transformer]]. @@ -326,7 +325,9 @@ class PipelineModel private[ml] ( private[spark] override def estimatedSize: Long = { estimateMatadataSize + stages.iterator.map { case model: Model[_] => model.estimatedSize - case stage => SizeEstimator.estimate(stage) + // ML Connect cache accounting only counts models. A non-model transformer may retain + // incidental references to objects such as SparkSession, so it has zero size here. + case _ => 0L }.sum } From 973f89eb35992e900cf99ae1767fbadad5e77191 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Thu, 23 Jul 2026 03:02:59 +0000 Subject: [PATCH 4/6] [ML] Clarify non-model pipeline stage sizing --- mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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 d61d696afd7bf..92b1e8801fdc4 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/Pipeline.scala @@ -325,8 +325,9 @@ class PipelineModel private[ml] ( private[spark] override def estimatedSize: Long = { estimateMatadataSize + stages.iterator.map { case model: Model[_] => model.estimatedSize - // ML Connect cache accounting only counts models. A non-model transformer may retain - // incidental references to objects such as SparkSession, so it has zero size here. + // Non-model transformers are skipped because: + // - ML Connect cache accounting counts only models. + // - They can retain incidental references, such as SparkSession. case _ => 0L }.sum } From 16b7ab1790adf9af434c4387ce3d9d8d2d9eeb85 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Thu, 23 Jul 2026 12:47:22 +0000 Subject: [PATCH 5/6] [ML] Strengthen PipelineModel size estimation test --- .../test/scala/org/apache/spark/ml/PipelineSuite.scala | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) 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 e4877b1e0851e..6065d9bb9e1cb 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala @@ -131,11 +131,17 @@ class PipelineSuite extends SparkFunSuite with MLlibTestSparkContext with Defaul test("PipelineModel estimated size") { val dataset = Seq("a", "b", "c").toDF("input") val pipeline = new Pipeline().setStages(Array( + new HashingTF().setInputCol("input").setOutputCol("features"), new StringIndexer().setInputCol("input").setOutputCol("indexed"))) val model = pipeline.fit(dataset) + val minSize = 1024 val maxSize = 1024 * 16 - assert(model.estimatedSize < maxSize, - s"Estimation (${model.estimatedSize}) should be less than $maxSize") + 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") { From 88e51269a503d51dd839f347db6fdfca4585bbf6 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Fri, 24 Jul 2026 00:52:09 +0000 Subject: [PATCH 6/6] [ML] Fix PipelineModel size estimation test input type --- mllib/src/test/scala/org/apache/spark/ml/PipelineSuite.scala | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 6065d9bb9e1cb..d2ad51766f049 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, StringIndexer} +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._ @@ -131,7 +131,7 @@ class PipelineSuite extends SparkFunSuite with MLlibTestSparkContext with Defaul test("PipelineModel estimated size") { val dataset = Seq("a", "b", "c").toDF("input") val pipeline = new Pipeline().setStages(Array( - new HashingTF().setInputCol("input").setOutputCol("features"), + new Tokenizer().setInputCol("input").setOutputCol("tokens"), new StringIndexer().setInputCol("input").setOutputCol("indexed"))) val model = pipeline.fit(dataset) val minSize = 1024