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]