From aafd9113fb75e0724e5e611ea001b25176b2b401 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Thu, 23 Jul 2026 01:04:41 +0000 Subject: [PATCH 1/4] [ML][FEATURE] Estimate size of feature lookup models --- .../apache/spark/ml/feature/CountVectorizer.scala | 5 +++++ .../org/apache/spark/ml/feature/MinHashLSH.scala | 5 +++++ .../apache/spark/ml/feature/TargetEncoder.scala | 6 +++++- .../apache/spark/ml/feature/VectorIndexer.scala | 5 +++++ .../spark/ml/feature/CountVectorizerSuite.scala | 8 ++++++++ .../apache/spark/ml/feature/MinHashLSHSuite.scala | 7 +++++++ .../spark/ml/feature/TargetEncoderSuite.scala | 14 +++++++++++++- .../spark/ml/feature/VectorIndexerSuite.scala | 7 +++++++ 8 files changed, 55 insertions(+), 2 deletions(-) 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..44d1981d48341 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,10 @@ class CountVectorizerModel( // For ml connect only private[ml] def this() = this("", Array.empty) + private[spark] override def estimatedSize: Long = { + estimateMatadataSize + SizeEstimator.estimate(vocabulary) + } + @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..7fa687fc35ce4 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,10 @@ class MinHashLSHModel private[ml]( // For ml connect only private[ml] def this() = this("", Array.empty) + private[spark] override def estimatedSize: Long = { + estimateMatadataSize + SizeEstimator.estimate(randCoefficients) + } + /** @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..b6f3f5401fadb 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,10 @@ class TargetEncoderModel private[ml] ( // For ml connect only private[ml] def this() = this("", Array.empty) + private[spark] override def estimatedSize: Long = { + estimateMatadataSize + SizeEstimator.estimate(stats) + } + /** @group setParam */ @Since("4.0.0") def setInputCol(value: String): this.type = set(inputCol, value) @@ -465,4 +470,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..953ffd2119296 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,10 @@ class VectorIndexerModel private[ml] ( // For ml connect only private[ml] def this() = this("", -1, Map.empty) + private[spark] override def estimatedSize: Long = { + estimateMatadataSize + SizeEstimator.estimate(categoryMaps) + } + /** 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..45b2e5db23173 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 = 4096 + 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..1e9d0f528c5d8 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 = 4096 + 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..14d8322c292e2 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 = 4096 + 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..8c6551c4ea27b 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 = 4096 + 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 From 4f1ea22fa8b8b20343c9a2c8facbb2ba9224c7f1 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Thu, 23 Jul 2026 02:33:20 +0000 Subject: [PATCH 2/4] [ML][FEATURE] Clarify lookup model size estimates --- .../scala/org/apache/spark/ml/feature/CountVectorizer.scala | 1 + .../src/main/scala/org/apache/spark/ml/feature/MinHashLSH.scala | 1 + .../main/scala/org/apache/spark/ml/feature/TargetEncoder.scala | 1 + .../main/scala/org/apache/spark/ml/feature/VectorIndexer.scala | 1 + .../org/apache/spark/ml/feature/CountVectorizerSuite.scala | 2 +- .../scala/org/apache/spark/ml/feature/MinHashLSHSuite.scala | 2 +- .../scala/org/apache/spark/ml/feature/TargetEncoderSuite.scala | 2 +- .../scala/org/apache/spark/ml/feature/VectorIndexerSuite.scala | 2 +- 8 files changed, 8 insertions(+), 4 deletions(-) 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 44d1981d48341..78c650ff233d2 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 @@ -287,6 +287,7 @@ class CountVectorizerModel( private[ml] def this() = this("", Array.empty) private[spark] override def estimatedSize: Long = { + // vocabulary: Array[String] estimateMatadataSize + SizeEstimator.estimate(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 7fa687fc35ce4..077133baaf8bf 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 @@ -55,6 +55,7 @@ class MinHashLSHModel private[ml]( private[ml] def this() = this("", Array.empty) private[spark] override def estimatedSize: Long = { + // randCoefficients: Array[(Int, Int)] estimateMatadataSize + SizeEstimator.estimate(randCoefficients) } 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 b6f3f5401fadb..c289d814cccfd 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 @@ -293,6 +293,7 @@ class TargetEncoderModel private[ml] ( private[ml] def this() = this("", Array.empty) private[spark] override def estimatedSize: Long = { + // stats: Array[Map[Double, (Double, Double)]] estimateMatadataSize + SizeEstimator.estimate(stats) } 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 953ffd2119296..90c7482b50ef8 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 @@ -304,6 +304,7 @@ class VectorIndexerModel private[ml] ( private[ml] def this() = this("", -1, Map.empty) private[spark] override def estimatedSize: Long = { + // categoryMaps: Map[Int, Map[Double, Int]] estimateMatadataSize + SizeEstimator.estimate(categoryMaps) } 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 45b2e5db23173..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 @@ -35,7 +35,7 @@ class CountVectorizerSuite extends MLTest with DefaultReadWriteTest { 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 = 4096 + val maxSize = 1024 * 4 assert(model.estimatedSize < maxSize, s"Estimation (${model.estimatedSize}) should be less than $maxSize") } 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 1e9d0f528c5d8..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 @@ -44,7 +44,7 @@ class MinHashLSHSuite extends MLTest with DefaultReadWriteTest { test("model estimated size") { val model = new MinHashLSH().setInputCol("keys").setOutputCol("values").fit(dataset) - val maxSize = 4096 + val maxSize = 1024 * 4 assert(model.estimatedSize < maxSize, s"Estimation (${model.estimatedSize}) should be less than $maxSize") } 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 14d8322c292e2..f10c5153c0788 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 @@ -97,7 +97,7 @@ class TargetEncoderSuite extends MLTest with DefaultReadWriteTest { .setInputCols(Array("input1", "input2", "input3")) .setOutputCols(Array("output1", "output2", "output3")) .fit(df) - val maxSize = 4096 + val maxSize = 1024 * 4 assert(model.estimatedSize < maxSize, s"Estimation (${model.estimatedSize}) should be less than $maxSize") } 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 8c6551c4ea27b..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 @@ -113,7 +113,7 @@ class VectorIndexerSuite extends MLTest with DefaultReadWriteTest with Logging { test("model estimated size") { val model = getIndexer.fit(densePoints1) - val maxSize = 4096 + val maxSize = 1024 * 4 assert(model.estimatedSize < maxSize, s"Estimation (${model.estimatedSize}) should be less than $maxSize") } From fcc1dffc12033c70021d26a454453ebcaef921e0 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Thu, 23 Jul 2026 03:01:08 +0000 Subject: [PATCH 3/4] [ML][FEATURE] Use an accumulator for lookup model size estimates --- .../scala/org/apache/spark/ml/feature/CountVectorizer.scala | 4 +++- .../main/scala/org/apache/spark/ml/feature/MinHashLSH.scala | 4 +++- .../scala/org/apache/spark/ml/feature/TargetEncoder.scala | 4 +++- .../scala/org/apache/spark/ml/feature/VectorIndexer.scala | 4 +++- 4 files changed, 12 insertions(+), 4 deletions(-) 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 78c650ff233d2..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 @@ -287,8 +287,10 @@ class CountVectorizerModel( private[ml] def this() = this("", Array.empty) private[spark] override def estimatedSize: Long = { + var size = estimateMatadataSize // vocabulary: Array[String] - estimateMatadataSize + SizeEstimator.estimate(vocabulary) + size += SizeEstimator.estimate(vocabulary) + size } @Since("1.5.0") 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 077133baaf8bf..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 @@ -55,8 +55,10 @@ class MinHashLSHModel private[ml]( private[ml] def this() = this("", Array.empty) private[spark] override def estimatedSize: Long = { + var size = estimateMatadataSize // randCoefficients: Array[(Int, Int)] - estimateMatadataSize + SizeEstimator.estimate(randCoefficients) + size += SizeEstimator.estimate(randCoefficients) + size } /** @group setParam */ 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 c289d814cccfd..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 @@ -293,8 +293,10 @@ class TargetEncoderModel private[ml] ( private[ml] def this() = this("", Array.empty) private[spark] override def estimatedSize: Long = { + var size = estimateMatadataSize // stats: Array[Map[Double, (Double, Double)]] - estimateMatadataSize + SizeEstimator.estimate(stats) + size += SizeEstimator.estimate(stats) + size } /** @group setParam */ 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 90c7482b50ef8..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 @@ -304,8 +304,10 @@ class VectorIndexerModel private[ml] ( private[ml] def this() = this("", -1, Map.empty) private[spark] override def estimatedSize: Long = { + var size = estimateMatadataSize // categoryMaps: Map[Int, Map[Double, Int]] - estimateMatadataSize + SizeEstimator.estimate(categoryMaps) + size += SizeEstimator.estimate(categoryMaps) + size } /** Java-friendly version of [[categoryMaps]] */ From 807b5afffe152e983adc5d0da7e981f0c13e60e8 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Thu, 23 Jul 2026 04:49:32 +0000 Subject: [PATCH 4/4] [ML][FEATURE] Relax TargetEncoder size estimation bound --- .../scala/org/apache/spark/ml/feature/TargetEncoderSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 f10c5153c0788..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 @@ -97,7 +97,7 @@ class TargetEncoderSuite extends MLTest with DefaultReadWriteTest { .setInputCols(Array("input1", "input2", "input3")) .setOutputCols(Array("output1", "output2", "output3")) .fit(df) - val maxSize = 1024 * 4 + val maxSize = 1024 * 6 assert(model.estimatedSize < maxSize, s"Estimation (${model.estimatedSize}) should be less than $maxSize") }