This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new c1d87077e6e3 [SPARK-58250][ML][CONNECT] Add OneVsRest size estimate
c1d87077e6e3 is described below
commit c1d87077e6e3f6d5b90d16d7e07a17b5e81f7a3a
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Wed Jul 22 19:46:35 2026 +0800
[SPARK-58250][ML][CONNECT] Add OneVsRest size estimate
### What changes were proposed in this pull request?
This patch adds a model-owned `estimatedSize` implementation for
`OneVsRestModel`. The estimate includes the model's parameter metadata, label
metadata, and the estimates of its binary classification submodels.
### Why are the changes needed?
Spark Connect uses `Model.estimatedSize` to enforce ML model-cache limits.
`OneVsRestModel` was the only concrete classification model that inherited the
default graph traversal. Its copied submodels retain their parent estimators,
so that traversal can include shared Spark session/context state rather than
only the cached model's data. The explicit estimate follows the established
classifier-model pattern and accounts only for state owned by the One-vs-Rest
model.
### Does this PR introduce _any_ user-facing change?
Yes. Spark Connect cache accounting for One-vs-Rest models becomes more
accurate, preventing shared runtime state from inflating the estimated model
size. There is no public API change.
### How was this patch tested?
Added a regression test that verifies the estimate is the sum of
One-vs-Rest parameter metadata, label metadata, and its contained models'
estimates.
Ran:
```
JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64 build/sbt 'mllib/testOnly
*OneVsRestSuite'
```
All 17 tests passed.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #57420 from zhengruifeng/SPARK-58181-one-vs-rest-estimated-size-dev3.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../org/apache/spark/ml/classification/OneVsRest.scala | 6 ++++++
.../spark/ml/classification/OneVsRestSuite.scala | 18 ++++++++++++++++++
2 files changed, 24 insertions(+)
diff --git
a/mllib/src/main/scala/org/apache/spark/ml/classification/OneVsRest.scala
b/mllib/src/main/scala/org/apache/spark/ml/classification/OneVsRest.scala
index f62b951fd650..38c68fd7000c 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/OneVsRest.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/OneVsRest.scala
@@ -42,6 +42,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.ThreadUtils
private[ml] trait ClassifierTypeTrait {
@@ -148,6 +149,11 @@ final class OneVsRestModel private[ml] (
@Since("2.4.0")
val numFeatures: Int = models.head.numFeatures
+ private[spark] override def estimatedSize: Long = {
+ estimateMatadataSize + SizeEstimator.estimate(labelMetadata) +
+ models.iterator.map(_.estimatedSize).sum
+ }
+
/** @group setParam */
@Since("2.1.0")
def setFeaturesCol(value: String): this.type = set(featuresCol, value)
diff --git
a/mllib/src/test/scala/org/apache/spark/ml/classification/OneVsRestSuite.scala
b/mllib/src/test/scala/org/apache/spark/ml/classification/OneVsRestSuite.scala
index 61b5a50fea69..70408731a20f 100644
---
a/mllib/src/test/scala/org/apache/spark/ml/classification/OneVsRestSuite.scala
+++
b/mllib/src/test/scala/org/apache/spark/ml/classification/OneVsRestSuite.scala
@@ -68,6 +68,24 @@ class OneVsRestSuite extends MLTest with
DefaultReadWriteTest {
ParamsSuite.checkParams(model)
}
+ test("SPARK-58250: OneVsRestModel estimated size") {
+ val trainingData = Seq(
+ (0.0, Vectors.dense(0.0, 0.0)),
+ (0.0, Vectors.dense(0.0, 1.0)),
+ (1.0, Vectors.dense(1.0, 0.0)),
+ (1.0, Vectors.dense(1.0, 1.0)),
+ (2.0, Vectors.dense(2.0, 0.0)),
+ (2.0, Vectors.dense(2.0, 1.0))).toDF("label", "features")
+
+ val model = new OneVsRest()
+ .setClassifier(new LogisticRegression().setMaxIter(1))
+ .fit(trainingData)
+
+ val maxSize = 32 * 1024
+ assert(model.estimatedSize < maxSize,
+ s"Estimation (${model.estimatedSize}) should be less than $maxSize")
+ }
+
test("one-vs-rest: default params") {
val numClasses = 3
val ova = new OneVsRest()
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
