Skip to content

[SPARK-58279][ML][CONNECT] Estimate size of TargetEncoder, VectorIndexer, CountVectorizer, and MinHashLSH models #57448

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
zhengruifeng wants to merge 4 commits into from
Closed

[SPARK-58279][ML][CONNECT] Estimate size of TargetEncoder, VectorIndexer, CountVectorizer, and MinHashLSH models #57448

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
aafd911
[ML][FEATURE] Estimate size of feature lookup models
zhengruifeng Jul 23, 2026
4f1ea22
[ML][FEATURE] Clarify lookup model size estimates
zhengruifeng Jul 23, 2026
fcc1dff
[ML][FEATURE] Use an accumulator for lookup model size estimates
zhengruifeng Jul 23, 2026
807b5af
[ML][FEATURE] Relax TargetEncoder size estimation bound
zhengruifeng Jul 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions mllib/src/main/scala/org/apache/spark/ml/feature/CountVectorizer.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}

/**
Expand Down Expand Up @@ -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)
Expand Down
8 changes: 8 additions & 0 deletions mllib/src/main/scala/org/apache/spark/ml/feature/MinHashLSH.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand Down
9 changes: 8 additions & 1 deletion mllib/src/main/scala/org/apache/spark/ml/feature/TargetEncoder.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -465,4 +473,3 @@ object TargetEncoderModel extends MLReadable[TargetEncoderModel] {
@Since("4.0.0")
override def load(path: String): TargetEncoderModel = super.load(path)
}

8 changes: 8 additions & 0 deletions mllib/src/main/scala/org/apache/spark/ml/feature/VectorIndexer.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 */
Expand Down Expand Up @@ -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 = {

@uros-b uros-b Jul 23, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note: VectorIndexerModel.estimatedSize counts only categoryMaps and omits partialFeatureAttributes, an eager private val (not lazy) built in the primary constructor as an Array[Attribute] of numFeatures allocated attribute objects. Because the old default estimated a parentless copy (new VectorIndexerModel(uid, numFeatures, categoryMaps), which re-materializes partialFeatureAttributes), that array was previously counted and is now dropped from the estimate. numFeatures is not bounded by categoryMaps.size, so a high-dimensional model (e.g. numFeatures=1000, few categorical) under-counts by ~1000 attribute objects. The merged sibling StringIndexerModel (SPARK-58251) sets the precedent by explicitly counting its derived eager field: SizeEstimator.estimate((labelsArray, labelsToIndexArray)). Fix is one line: estimateMatadataSize + SizeEstimator.estimate((categoryMaps, partialFeatureAttributes)).

@zhengruifeng zhengruifeng Jul 23, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point. partialFeatureAttributes is eagerly constructed derived state, so the current override can undercount it compared with the previous parentless-copy estimate. I will created a sub-task https://issues.apache.org/jira/browse/SPARK-58293 to defer its initialization and work on that follow-up soon. The follow-up will ensure size estimation accounts for the state when it is materialized.

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]] = {
Expand Down
8 changes: 8 additions & 0 deletions mllib/src/test/scala/org/apache/spark/ml/feature/CountVectorizerSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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") {
Expand Down
7 changes: 7 additions & 0 deletions mllib/src/test/scala/org/apache/spark/ml/feature/MinHashLSHSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
14 changes: 13 additions & 1 deletion mllib/src/test/scala/org/apache/spark/ml/feature/TargetEncoderSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -521,4 +533,4 @@ class TargetEncoderSuite extends MLTest with DefaultReadWriteTest {

}

}
}
7 changes: 7 additions & 0 deletions mllib/src/test/scala/org/apache/spark/ml/feature/VectorIndexerSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down