diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala index c7fc4ce6898bd..a85e9236517f7 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala @@ -33,6 +33,7 @@ import org.apache.spark.sql.functions._ import org.apache.spark.sql.types._ import org.apache.spark.storage.StorageLevel import org.apache.spark.util.ArrayImplicits._ +import org.apache.spark.util.SizeEstimator import org.apache.spark.util.collection.{OpenHashMap, Utils} /** @@ -285,6 +286,13 @@ class CountVectorizerModel( // For ml connect only private[ml] def this() = this("", Array.empty) + private[spark] override def estimatedSize: Long = { + var size = estimateMatadataSize + // vocabulary: Array[String] + size += SizeEstimator.estimate(vocabulary) + size + } + @Since("1.5.0") def this(vocabulary: Array[String]) = { this(Identifiable.randomUID("cntVecModel"), vocabulary) diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/MinHashLSH.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/MinHashLSH.scala index 90aaaea7e36e6..4b0e5ca4fb31a 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/MinHashLSH.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/MinHashLSH.scala @@ -29,6 +29,7 @@ import org.apache.spark.ml.param.ParamMap import org.apache.spark.ml.param.shared.HasSeed import org.apache.spark.ml.util._ import org.apache.spark.sql.types.StructType +import org.apache.spark.util.SizeEstimator /** * Model produced by [[MinHashLSH]], where multiple hash functions are stored. Each hash function @@ -53,6 +54,13 @@ class MinHashLSHModel private[ml]( // For ml connect only private[ml] def this() = this("", Array.empty) + private[spark] override def estimatedSize: Long = { + var size = estimateMatadataSize + // randCoefficients: Array[(Int, Int)] + size += SizeEstimator.estimate(randCoefficients) + size + } + /** @group setParam */ @Since("2.4.0") override def setInputCol(value: String): this.type = super.set(inputCol, value) diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/TargetEncoder.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/TargetEncoder.scala index 736aa4fc5356b..edbcb7fca7428 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/TargetEncoder.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/TargetEncoder.scala @@ -31,6 +31,7 @@ import org.apache.spark.ml.util._ import org.apache.spark.sql.{Column, DataFrame, Dataset, Row} import org.apache.spark.sql.functions._ import org.apache.spark.sql.types._ +import org.apache.spark.util.SizeEstimator /** Private trait for params and common methods for TargetEncoder and TargetEncoderModel */ private[ml] trait TargetEncoderBase extends Params with HasLabelCol @@ -291,6 +292,13 @@ class TargetEncoderModel private[ml] ( // For ml connect only private[ml] def this() = this("", Array.empty) + private[spark] override def estimatedSize: Long = { + var size = estimateMatadataSize + // stats: Array[Map[Double, (Double, Double)]] + size += SizeEstimator.estimate(stats) + size + } + /** @group setParam */ @Since("4.0.0") def setInputCol(value: String): this.type = set(inputCol, value) @@ -465,4 +473,3 @@ object TargetEncoderModel extends MLReadable[TargetEncoderModel] { @Since("4.0.0") override def load(path: String): TargetEncoderModel = super.load(path) } - diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorIndexer.scala b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorIndexer.scala index 290e5fe05fcab..7fce6d51dcfcb 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/feature/VectorIndexer.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/feature/VectorIndexer.scala @@ -37,6 +37,7 @@ import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession} import org.apache.spark.sql.functions.udf import org.apache.spark.sql.types.{StructField, StructType} import org.apache.spark.util.ArrayImplicits._ +import org.apache.spark.util.SizeEstimator import org.apache.spark.util.collection.{OpenHashSet, Utils} /** Private trait for params for VectorIndexer and VectorIndexerModel */ @@ -302,6 +303,13 @@ class VectorIndexerModel private[ml] ( // For ml connect only private[ml] def this() = this("", -1, Map.empty) + private[spark] override def estimatedSize: Long = { + var size = estimateMatadataSize + // categoryMaps: Map[Int, Map[Double, Int]] + size += SizeEstimator.estimate(categoryMaps) + size + } + /** Java-friendly version of [[categoryMaps]] */ @Since("1.4.0") def javaCategoryMaps: JMap[JInt, JMap[JDouble, JInt]] = { diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/CountVectorizerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/CountVectorizerSuite.scala index ecde4e92a34f7..295e96bcfe6a0 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/feature/CountVectorizerSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/feature/CountVectorizerSuite.scala @@ -32,6 +32,14 @@ class CountVectorizerSuite extends MLTest with DefaultReadWriteTest { ParamsSuite.checkParams(new CountVectorizerModel(Array("empty"))) } + test("model estimated size") { + val df = Seq(Seq("a", "b"), Seq("a", "c")).toDF("words") + val model = new CountVectorizer().setInputCol("words").setOutputCol("features").fit(df) + val maxSize = 1024 * 4 + assert(model.estimatedSize < maxSize, + s"Estimation (${model.estimatedSize}) should be less than $maxSize") + } + private def split(s: String): Seq[String] = s.split("\\s+").toImmutableArraySeq test("CountVectorizerModel common cases") { diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/MinHashLSHSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/MinHashLSHSuite.scala index c99e0fa3f8623..151135a459a81 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/feature/MinHashLSHSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/feature/MinHashLSHSuite.scala @@ -42,6 +42,13 @@ class MinHashLSHSuite extends MLTest with DefaultReadWriteTest { ParamsSuite.checkParams(model) } + test("model estimated size") { + val model = new MinHashLSH().setInputCol("keys").setOutputCol("values").fit(dataset) + val maxSize = 1024 * 4 + assert(model.estimatedSize < maxSize, + s"Estimation (${model.estimatedSize}) should be less than $maxSize") + } + test("setters") { val model = new MinHashLSHModel("mh", randCoefficients = Array((1, 0))) .setInputCol("testkeys") diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/TargetEncoderSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/TargetEncoderSuite.scala index 6bb3ce224a2e7..7b2d0ef005c57 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/feature/TargetEncoderSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/feature/TargetEncoderSuite.scala @@ -90,6 +90,18 @@ class TargetEncoderSuite extends MLTest with DefaultReadWriteTest { ParamsSuite.checkParams(new TargetEncoder) } + test("model estimated size") { + val df = spark.createDataFrame(sc.parallelize(data_binary), schema) + val model = new TargetEncoder() + .setLabelCol("label") + .setInputCols(Array("input1", "input2", "input3")) + .setOutputCols(Array("output1", "output2", "output3")) + .fit(df) + val maxSize = 1024 * 6 + assert(model.estimatedSize < maxSize, + s"Estimation (${model.estimatedSize}) should be less than $maxSize") + } + test("TargetEncoder - binary target") { val df = spark.createDataFrame(sc.parallelize(data_binary), schema) @@ -521,4 +533,4 @@ class TargetEncoderSuite extends MLTest with DefaultReadWriteTest { } -} \ No newline at end of file +} diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/VectorIndexerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/VectorIndexerSuite.scala index 0c97e48079599..1529b4c661cec 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/feature/VectorIndexerSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/feature/VectorIndexerSuite.scala @@ -111,6 +111,13 @@ class VectorIndexerSuite extends MLTest with DefaultReadWriteTest with Logging { ParamsSuite.checkParams(model) } + test("model estimated size") { + val model = getIndexer.fit(densePoints1) + val maxSize = 1024 * 4 + assert(model.estimatedSize < maxSize, + s"Estimation (${model.estimatedSize}) should be less than $maxSize") + } + test("Cannot fit an empty DataFrame") { val rdd = Array.empty[Vector].map(FeatureData).toSeq.toDF() val vectorIndexer = getIndexer