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 6b77b817e390 [SPARK-58278][ML][CONNECT] Estimate size of Univariate,
ChiSq, Variance selector, and Imputer models
6b77b817e390 is described below
commit 6b77b817e3904a887826d48d5d29ff730342e1af
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Thu Jul 23 20:44:36 2026 +0800
[SPARK-58278][ML][CONNECT] Estimate size of Univariate, ChiSq, Variance
selector, and Imputer models
### What changes were proposed in this pull request?
Adds explicit `estimatedSize` implementations for
`UnivariateFeatureSelectorModel`, `ChiSqSelectorModel`,
`VarianceThresholdSelectorModel`, and `ImputerModel`. The selector models
account for their selected-feature arrays; the Imputer model accounts for its
column names and surrogate values.
Each model is covered by a fitted-model regression test that asserts the
estimate remains below 2 KiB.
### 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 #57447 from zhengruifeng/feature-selector-stats-size-estimates-dev3.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit 8f710a6544f6479ad6a9412a4ace15df3b7f5128)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
mllib/src/main/scala/org/apache/spark/ml/feature/Imputer.scala | 9 +++++++++
mllib/src/main/scala/org/apache/spark/ml/feature/Selector.scala | 7 ++++++-
.../org/apache/spark/ml/feature/UnivariateFeatureSelector.scala | 6 ++++++
.../org/apache/spark/ml/feature/VarianceThresholdSelector.scala | 6 ++++++
.../scala/org/apache/spark/ml/feature/ChiSqSelectorSuite.scala | 8 ++++++++
.../test/scala/org/apache/spark/ml/feature/ImputerSuite.scala | 9 +++++++++
.../apache/spark/ml/feature/UnivariateFeatureSelectorSuite.scala | 8 ++++++++
.../apache/spark/ml/feature/VarianceThresholdSelectorSuite.scala | 7 +++++++
8 files changed, 59 insertions(+), 1 deletion(-)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Imputer.scala
b/mllib/src/main/scala/org/apache/spark/ml/feature/Imputer.scala
index a5b1a3cfde82..693a67088ac3 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/Imputer.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Imputer.scala
@@ -29,6 +29,7 @@ import org.apache.spark.sql.{DataFrame, Dataset, Row,
SparkSession}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types._
import org.apache.spark.util.ArrayImplicits._
+import org.apache.spark.util.SizeEstimator
/**
* Params for [[Imputer]] and [[ImputerModel]].
@@ -246,6 +247,14 @@ class ImputerModel private[ml] (
// For ml connect only
private[ml] def this() = this("", Array.empty, Array.emptyDoubleArray)
+ private[spark] override def estimatedSize: Long = {
+ // columnNames: Array[String]
+ var size = estimateMatadataSize + SizeEstimator.estimate(columnNames)
+ // surrogates: Array[Double]
+ size += SizeEstimator.estimate(surrogates)
+ size
+ }
+
@Since("2.2.0")
def surrogateDF: DataFrame = {
val spark = SparkSession.builder().getOrCreate()
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/Selector.scala
b/mllib/src/main/scala/org/apache/spark/ml/feature/Selector.scala
index dde1068c5b92..1914a98014da 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/Selector.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/Selector.scala
@@ -29,6 +29,7 @@ import org.apache.spark.ml.util._
import org.apache.spark.sql.{DataFrame, Dataset}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.{StructField, StructType}
+import org.apache.spark.util.SizeEstimator
/**
@@ -275,6 +276,11 @@ private[ml] abstract class SelectorModel[T <:
SelectorModel[T]] (
extends Model[T] with SelectorParams with MLWritable {
self: T =>
+ private[spark] override def estimatedSize: Long = {
+ // selectedFeatures: Array[Int]
+ estimateMatadataSize + SizeEstimator.estimate(selectedFeatures)
+ }
+
/** @group setParam */
@Since("3.1.0")
def setFeaturesCol(value: String): this.type = set(featuresCol, value)
@@ -385,4 +391,3 @@ private[feature] object SelectorModel {
(newIndices.result(), newValues.result())
}
}
-
diff --git
a/mllib/src/main/scala/org/apache/spark/ml/feature/UnivariateFeatureSelector.scala
b/mllib/src/main/scala/org/apache/spark/ml/feature/UnivariateFeatureSelector.scala
index ca7e9eb826fc..6eb5696520cd 100644
---
a/mllib/src/main/scala/org/apache/spark/ml/feature/UnivariateFeatureSelector.scala
+++
b/mllib/src/main/scala/org/apache/spark/ml/feature/UnivariateFeatureSelector.scala
@@ -35,6 +35,7 @@ import org.apache.spark.sql.{DataFrame, Dataset, SparkSession}
import org.apache.spark.sql.functions.{col, udf}
import org.apache.spark.sql.types.{StructField, StructType}
import org.apache.spark.util.ArrayImplicits._
+import org.apache.spark.util.SizeEstimator
/**
@@ -294,6 +295,11 @@ class UnivariateFeatureSelectorModel private[ml](
// For ml connect only
private[ml] def this() = this("", Array.emptyIntArray)
+ private[spark] override def estimatedSize: Long = {
+ // selectedFeatures: Array[Int]
+ estimateMatadataSize + SizeEstimator.estimate(selectedFeatures)
+ }
+
/** @group setParam */
@Since("3.1.1")
def setFeaturesCol(value: String): this.type = set(featuresCol, value)
diff --git
a/mllib/src/main/scala/org/apache/spark/ml/feature/VarianceThresholdSelector.scala
b/mllib/src/main/scala/org/apache/spark/ml/feature/VarianceThresholdSelector.scala
index b794e6df3944..cdbdf122dc05 100644
---
a/mllib/src/main/scala/org/apache/spark/ml/feature/VarianceThresholdSelector.scala
+++
b/mllib/src/main/scala/org/apache/spark/ml/feature/VarianceThresholdSelector.scala
@@ -32,6 +32,7 @@ import org.apache.spark.sql._
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.StructType
import org.apache.spark.util.ArrayImplicits._
+import org.apache.spark.util.SizeEstimator
/**
* Params for [[VarianceThresholdSelector]] and
[[VarianceThresholdSelectorModel]].
@@ -131,6 +132,11 @@ class VarianceThresholdSelectorModel private[ml](
// For ml connect only
private[ml] def this() = this("", Array.emptyIntArray)
+ private[spark] override def estimatedSize: Long = {
+ // selectedFeatures: Array[Int]
+ estimateMatadataSize + SizeEstimator.estimate(selectedFeatures)
+ }
+
if (selectedFeatures.length >= 2) {
require(selectedFeatures.sliding(2).forall(l => l(0) < l(1)),
"Index should be strictly increasing.")
diff --git
a/mllib/src/test/scala/org/apache/spark/ml/feature/ChiSqSelectorSuite.scala
b/mllib/src/test/scala/org/apache/spark/ml/feature/ChiSqSelectorSuite.scala
index b48ef0fad158..8ae59fb3802a 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/ChiSqSelectorSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/ChiSqSelectorSuite.scala
@@ -118,6 +118,14 @@ class ChiSqSelectorSuite extends MLTest with
DefaultReadWriteTest {
ParamsSuite.checkParams(model)
}
+ test("model estimated size") {
+ val model = new
ChiSqSelector().setSelectorType("numTopFeatures").setNumTopFeatures(1)
+ .fit(dataset)
+ val maxSize = 1024 * 4
+ assert(model.estimatedSize < maxSize,
+ s"Estimation (${model.estimatedSize}) should be less than $maxSize")
+ }
+
test("Test Chi-Square selector: numTopFeatures") {
val selector = new ChiSqSelector()
.setOutputCol("filtered").setSelectorType("numTopFeatures").setNumTopFeatures(1)
diff --git
a/mllib/src/test/scala/org/apache/spark/ml/feature/ImputerSuite.scala
b/mllib/src/test/scala/org/apache/spark/ml/feature/ImputerSuite.scala
index a2fc333535fc..2328763c0dd5 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/feature/ImputerSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/feature/ImputerSuite.scala
@@ -27,6 +27,15 @@ import org.apache.spark.sql.types._
class ImputerSuite extends MLTest with DefaultReadWriteTest {
+ test("model estimated size") {
+ val df = spark.createDataFrame(Seq((1.0, 2.0), (3.0, 4.0))).toDF("a", "b")
+ val model = new Imputer().setInputCols(Array("a",
"b")).setOutputCols(Array("x", "y"))
+ .fit(df)
+ val maxSize = 1024 * 4
+ assert(model.estimatedSize < maxSize,
+ s"Estimation (${model.estimatedSize}) should be less than $maxSize")
+ }
+
test("Imputer for Double with default missing Value NaN") {
val df = spark.createDataFrame(Seq(
(0, 1.0, 4.0, 1.0, 1.0, 1.0, 4.0, 4.0, 4.0),
diff --git
a/mllib/src/test/scala/org/apache/spark/ml/feature/UnivariateFeatureSelectorSuite.scala
b/mllib/src/test/scala/org/apache/spark/ml/feature/UnivariateFeatureSelectorSuite.scala
index 2f9aa69e413d..62324f1465a8 100644
---
a/mllib/src/test/scala/org/apache/spark/ml/feature/UnivariateFeatureSelectorSuite.scala
+++
b/mllib/src/test/scala/org/apache/spark/ml/feature/UnivariateFeatureSelectorSuite.scala
@@ -343,6 +343,14 @@ class UnivariateFeatureSelectorSuite extends MLTest with
DefaultReadWriteTest {
ParamsSuite.checkParams(new UnivariateFeatureSelector())
}
+ test("model estimated size") {
+ val model =
selector3.setSelectionMode("numTopFeatures").setSelectionThreshold(1)
+ .fit(datasetChi2)
+ val maxSize = 1024 * 4
+ assert(model.estimatedSize < maxSize,
+ s"Estimation (${model.estimatedSize}) should be less than $maxSize")
+ }
+
test("Test numTopFeatures") {
val testParams: Seq[(UnivariateFeatureSelector, Dataset[_])] = Seq(
(selector1.setSelectionMode("numTopFeatures").setSelectionThreshold(1),
datasetAnova),
diff --git
a/mllib/src/test/scala/org/apache/spark/ml/feature/VarianceThresholdSelectorSuite.scala
b/mllib/src/test/scala/org/apache/spark/ml/feature/VarianceThresholdSelectorSuite.scala
index 77b819d1c999..6a019f5f07d3 100644
---
a/mllib/src/test/scala/org/apache/spark/ml/feature/VarianceThresholdSelectorSuite.scala
+++
b/mllib/src/test/scala/org/apache/spark/ml/feature/VarianceThresholdSelectorSuite.scala
@@ -54,6 +54,13 @@ class VarianceThresholdSelectorSuite extends MLTest with
DefaultReadWriteTest {
ParamsSuite.checkParams(new VarianceThresholdSelector)
}
+ test("model estimated size") {
+ val model = new VarianceThresholdSelector().fit(dataset)
+ val maxSize = 1024 * 2
+ assert(model.estimatedSize < maxSize,
+ s"Estimation (${model.estimatedSize}) should be less than $maxSize")
+ }
+
test("Test VarianceThresholdSelector: varianceThreshold not set") {
val selector = new VarianceThresholdSelector().setOutputCol("filtered")
val model = testSelector(selector, dataset)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
