This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new a27ffada2d8d [SPARK-58279][ML][CONNECT] Estimate size of
TargetEncoder, VectorIndexer, CountVectorizer, and MinHashLSH models
a27ffada2d8d is described below
commit a27ffada2d8d77068f8a0c9329bb524dc3ab90f9
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Fri Jul 24 09:13:37 2026 +0800
[SPARK-58279][ML][CONNECT] Estimate size of TargetEncoder, VectorIndexer,
CountVectorizer, and MinHashLSH models
### What changes were proposed in this pull request?
Adds explicit `estimatedSize` implementations for `TargetEncoderModel`,
`VectorIndexerModel`, `CountVectorizerModel`, and `MinHashLSHModel`. Each
override accounts for model metadata and the model-owned lookup state: encoder
statistics, category maps, vocabulary, or random hash coefficients.
Adds fitted-model regression tests in the owning algorithm suites with a 4
KiB upper bound.
### Why are the changes needed?
The default reflection-based estimate can traverse substantially more state
than the fitted model owns. Direct accounting keeps model-size estimation
small, predictable, and focused on the learned state.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
- `JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64 build/sbt
'mllib/Test/compile'`\n\n### Was this patch authored or co-authored using
generative AI tooling?\n\nGenerated-by: Codex GPT-5
Closes #57448 from zhengruifeng/feature-lookup-size-estimates-dev3.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit b3a45442dd0a8bb56c5730bc075c4c8b65eaf997)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../org/apache/spark/ml/feature/CountVectorizer.scala | 8 ++++++++
.../scala/org/apache/spark/ml/feature/MinHashLSH.scala | 8 ++++++++
.../scala/org/apache/spark/ml/feature/TargetEncoder.scala | 9 ++++++++-
.../scala/org/apache/spark/ml/feature/VectorIndexer.scala | 8 ++++++++
.../org/apache/spark/ml/feature/CountVectorizerSuite.scala | 8 ++++++++
.../org/apache/spark/ml/feature/MinHashLSHSuite.scala | 7 +++++++
.../org/apache/spark/ml/feature/TargetEncoderSuite.scala | 14 +++++++++++++-
.../org/apache/spark/ml/feature/VectorIndexerSuite.scala | 7 +++++++
8 files changed, 67 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 c7fc4ce6898b..a85e9236517f 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 90aaaea7e36e..4b0e5ca4fb31 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 736aa4fc5356..edbcb7fca742 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 290e5fe05fca..7fce6d51dcfc 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 ecde4e92a34f..295e96bcfe6a 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 c99e0fa3f862..151135a459a8 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 6bb3ce224a2e..7b2d0ef005c5 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 0c97e4807959..1529b4c661ce 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
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
