Skip to content

Commit c1d8707

Browse files
committed
[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 <ruifengz@apache.org> Signed-off-by: Ruifeng Zheng <ruifengz@apache.org>
1 parent ce39ba8 commit c1d8707

2 files changed

Lines changed: 24 additions & 0 deletions

File tree

mllib/src/main/scala/org/apache/spark/ml/classification/OneVsRest.scala

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ import org.apache.spark.sql.functions._
4242
import org.apache.spark.sql.types._
4343
import org.apache.spark.storage.StorageLevel
4444
import org.apache.spark.util.ArrayImplicits._
45+
import org.apache.spark.util.SizeEstimator
4546
import org.apache.spark.util.ThreadUtils
4647

4748
private[ml] trait ClassifierTypeTrait {
@@ -148,6 +149,11 @@ final class OneVsRestModel private[ml] (
148149
@Since("2.4.0")
149150
val numFeatures: Int = models.head.numFeatures
150151

152+
private[spark] override def estimatedSize: Long = {
153+
estimateMatadataSize + SizeEstimator.estimate(labelMetadata) +
154+
models.iterator.map(_.estimatedSize).sum
155+
}
156+
151157
/** @group setParam */
152158
@Since("2.1.0")
153159
def setFeaturesCol(value: String): this.type = set(featuresCol, value)

mllib/src/test/scala/org/apache/spark/ml/classification/OneVsRestSuite.scala

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,24 @@ class OneVsRestSuite extends MLTest with DefaultReadWriteTest {
6868
ParamsSuite.checkParams(model)
6969
}
7070

71+
test("SPARK-58250: OneVsRestModel estimated size") {
72+
val trainingData = Seq(
73+
(0.0, Vectors.dense(0.0, 0.0)),
74+
(0.0, Vectors.dense(0.0, 1.0)),
75+
(1.0, Vectors.dense(1.0, 0.0)),
76+
(1.0, Vectors.dense(1.0, 1.0)),
77+
(2.0, Vectors.dense(2.0, 0.0)),
78+
(2.0, Vectors.dense(2.0, 1.0))).toDF("label", "features")
79+
80+
val model = new OneVsRest()
81+
.setClassifier(new LogisticRegression().setMaxIter(1))
82+
.fit(trainingData)
83+
84+
val maxSize = 32 * 1024
85+
assert(model.estimatedSize < maxSize,
86+
s"Estimation (${model.estimatedSize}) should be less than $maxSize")
87+
}
88+
7189
test("one-vs-rest: default params") {
7290
val numClasses = 3
7391
val ova = new OneVsRest()

0 commit comments

Comments
 (0)