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 fdccb58a07c7 [SPARK-58251][ML][CONNECT] Add size estimates for
OneHotEncoderModel, PCAModel, StringIndexerModel, and IDFModel
fdccb58a07c7 is described below
commit fdccb58a07c70a6bd891d98fcc3f5a727327e843
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Thu Jul 23 10:05:03 2026 +0800
[SPARK-58251][ML][CONNECT] Add size estimates for OneHotEncoderModel,
PCAModel, StringIndexerModel, and IDFModel
### What changes were proposed in this pull request?
Add explicit `estimatedSize` implementations for OneHotEncoderModel,
PCAModel, StringIndexerModel, and IDFModel. Each directly accounts for model
metadata and model-owned learned state; vectors and matrices use
`getSizeInBytes`.
### Why are the changes needed?
SPARK-57521 already makes the default `Model.estimatedSize` estimate a
parentless copy, avoiding parent-estimator and SparkSession traversal. These
explicit implementations make size accounting cheaper and more precise for
feature models, and align them with models that already provide specialized
estimates.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
- Added fitted-model size bounds to the corresponding feature algorithm
suites.
- `JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64 build/sbt
'mllib/Test/compile'`\n- Targeted suites were not run.\n\n### Was this patch
authored or co-authored using generative AI tooling?\n\nGenerated-by: OpenAI
Codex (GPT-5)
Closes #57422 from zhengruifeng/feature-model-size-estimates-master-dev3.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit 36d311f3c62652e0a95ca6dab4ad1f3f3d78fecd)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala | 14 ++++++++++++++
.../scala/org/apache/spark/ml/feature/OneHotEncoder.scala | 4 ++++
mllib/src/main/scala/org/apache/spark/ml/feature/PCA.scala | 11 +++++++++++
.../scala/org/apache/spark/ml/feature/StringIndexer.scala | 4 ++++
.../test/scala/org/apache/spark/ml/feature/IDFSuite.scala | 13 +++++++++++++
.../org/apache/spark/ml/feature/OneHotEncoderSuite.scala | 11 +++++++++++
.../test/scala/org/apache/spark/ml/feature/PCASuite.scala | 14 ++++++++++++++
.../org/apache/spark/ml/feature/StringIndexerSuite.scala | 11 +++++++++++
8 files changed, 82 insertions(+)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala
b/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala
index 289217e1b817..3ef916bc200c 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/IDF.scala
@@ -34,6 +34,7 @@ import org.apache.spark.rdd.RDD
import org.apache.spark.sql._
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.StructType
+import org.apache.spark.util.SizeEstimator
import org.apache.spark.util.VersionUtils.majorVersion
/**
@@ -126,6 +127,19 @@ class IDFModel private[ml] (
// For ml connect only
private[ml] def this() = this("", null)
+ private[spark] override def estimatedSize: Long = {
+ var size = estimateMatadataSize
+ if (idfModel != null) {
+ if (idfModel.idf != null) {
+ size += idfModel.idf.asML.getSizeInBytes
+ }
+ if (idfModel.docFreq != null) {
+ size += SizeEstimator.estimate(idfModel.docFreq)
+ }
+ }
+ size
+ }
+
/** @group setParam */
@Since("1.4.0")
def setInputCol(value: String): this.type = set(inputCol, value)
diff --git
a/mllib/src/main/scala/org/apache/spark/ml/feature/OneHotEncoder.scala
b/mllib/src/main/scala/org/apache/spark/ml/feature/OneHotEncoder.scala
index b38dfcb4ed5d..c5120d026d47 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/OneHotEncoder.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/OneHotEncoder.scala
@@ -34,6 +34,7 @@ import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.functions.{col, lit, udf}
import org.apache.spark.sql.types.{DoubleType, StructField, StructType}
import org.apache.spark.util.ArrayImplicits._
+import org.apache.spark.util.SizeEstimator
/** Private trait for params and common methods for OneHotEncoder and
OneHotEncoderModel */
private[ml] trait OneHotEncoderBase extends Params with HasHandleInvalid
@@ -239,6 +240,9 @@ class OneHotEncoderModel private[ml] (
// For ml connect only
private[ml] def this() = this("", Array.emptyIntArray)
+ private[spark] override def estimatedSize: Long =
+ estimateMatadataSize + SizeEstimator.estimate(categorySizes)
+
// Returns the category size for each index with `dropLast` and
`handleInvalid`
// taken into account.
private def getConfigedCategorySizes: Array[Int] = {
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/PCA.scala
b/mllib/src/main/scala/org/apache/spark/ml/feature/PCA.scala
index d09d471788cf..a731b6750573 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/PCA.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/PCA.scala
@@ -132,6 +132,17 @@ class PCAModel private[ml] (
// For ml connect only
private[ml] def this() = this("", Matrices.empty, Vectors.empty)
+ private[spark] override def estimatedSize: Long = {
+ var size = estimateMatadataSize
+ if (pc != null) {
+ size += pc.getSizeInBytes
+ }
+ if (explainedVariance != null) {
+ size += explainedVariance.getSizeInBytes
+ }
+ size
+ }
+
/** @group setParam */
@Since("1.5.0")
def setInputCol(value: String): this.type = set(inputCol, value)
diff --git
a/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
b/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
index 6ba57f583ae3..8bb07d3f0260 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/StringIndexer.scala
@@ -33,6 +33,7 @@ import org.apache.spark.sql.{AnalysisException, Column,
DataFrame, Dataset}
import org.apache.spark.sql.functions.{get => fget, printf => fprintf, _}
import org.apache.spark.sql.types._
import org.apache.spark.util.ArrayImplicits._
+import org.apache.spark.util.SizeEstimator
import org.apache.spark.util.VersionUtils.majorMinorVersion
import org.apache.spark.util.collection.OpenHashMap
@@ -327,6 +328,9 @@ class StringIndexerModel (
}
}
+ private[spark] override def estimatedSize: Long =
+ estimateMatadataSize + SizeEstimator.estimate((labelsArray,
labelsToIndexArray))
+
/** @group setParam */
@Since("1.6.0")
def setHandleInvalid(value: String): this.type = set(handleInvalid, value)
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/IDFSuite.scala
b/mllib/src/test/scala/org/apache/spark/ml/feature/IDFSuite.scala
index 4f5c8b25de77..330b714555af 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/IDFSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/IDFSuite.scala
@@ -49,6 +49,19 @@ class IDFSuite extends MLTest with DefaultReadWriteTest {
ParamsSuite.checkParams(model)
}
+ test("IDFModel estimated size") {
+ val model = new IDF()
+ .setInputCol("features")
+ .setOutputCol("idf")
+ .fit(Seq(
+ Tuple1(Vectors.sparse(3, Array(0), Array(1.0))),
+ Tuple1(Vectors.sparse(3, Array(1), Array(1.0)))).toDF("features"))
+
+ val maxSize = 2 * 1024
+ assert(model.estimatedSize < maxSize,
+ s"Estimation (${model.estimatedSize}) should be less than $maxSize")
+ }
+
test("compute IDF with default parameter") {
val numOfFeatures = 4
val data = Array(
diff --git
a/mllib/src/test/scala/org/apache/spark/ml/feature/OneHotEncoderSuite.scala
b/mllib/src/test/scala/org/apache/spark/ml/feature/OneHotEncoderSuite.scala
index 49ba4b0233ed..a3003d3e0174 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/OneHotEncoderSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/OneHotEncoderSuite.scala
@@ -35,6 +35,17 @@ class OneHotEncoderSuite extends MLTest with
DefaultReadWriteTest {
ParamsSuite.checkParams(new OneHotEncoder)
}
+ test("OneHotEncoderModel estimated size") {
+ val model = new OneHotEncoder()
+ .setInputCol("input")
+ .setOutputCol("output")
+ .fit(Seq(0.0, 1.0).toDF("input"))
+
+ val maxSize = 4 * 1024
+ assert(model.estimatedSize < maxSize,
+ s"Estimation (${model.estimatedSize}) should be less than $maxSize")
+ }
+
test("OneHotEncoder dropLast = false") {
val data = Seq(
Row(0.0, Vectors.sparse(3, Seq((0, 1.0)))),
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/PCASuite.scala
b/mllib/src/test/scala/org/apache/spark/ml/feature/PCASuite.scala
index 576065ce3be8..ecdab2872f91 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/PCASuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/PCASuite.scala
@@ -38,6 +38,20 @@ class PCASuite extends MLTest with DefaultReadWriteTest {
ParamsSuite.checkParams(model)
}
+ test("PCAModel estimated size") {
+ val model = new PCA()
+ .setInputCol("features")
+ .setOutputCol("pca")
+ .setK(1)
+ .fit(Seq(
+ Tuple1(Vectors.dense(1.0, 0.0)),
+ Tuple1(Vectors.dense(0.0, 1.0))).toDF("features"))
+
+ val maxSize = 2 * 1024
+ assert(model.estimatedSize < maxSize,
+ s"Estimation (${model.estimatedSize}) should be less than $maxSize")
+ }
+
test("pca") {
val data = Array(
Vectors.sparse(5, Seq((1, 1.0), (3, 7.0))),
diff --git
a/mllib/src/test/scala/org/apache/spark/ml/feature/StringIndexerSuite.scala
b/mllib/src/test/scala/org/apache/spark/ml/feature/StringIndexerSuite.scala
index 8f3750959d2b..18005687f565 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/StringIndexerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/StringIndexerSuite.scala
@@ -37,6 +37,17 @@ class StringIndexerSuite extends MLTest with
DefaultReadWriteTest {
ParamsSuite.checkParams(modelWithoutUid)
}
+ test("StringIndexerModel estimated size") {
+ val model = new StringIndexer()
+ .setInputCol("input")
+ .setOutputCol("output")
+ .fit(Seq("a", "b").toDF("input"))
+
+ val maxSize = 4 * 1024
+ assert(model.estimatedSize < maxSize,
+ s"Estimation (${model.estimatedSize}) should be less than $maxSize")
+ }
+
test("params: input/output columns") {
val stringIndexerSingleCol = new StringIndexer()
.setInputCol("in").setOutputCol("out")
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
