Skip to content

[SPARK-58360][ML][CONNECT] Avoid nested parent overcount in RFormulaModel size estimate #57552

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 7 commits into from
Closed

[SPARK-58360][ML][CONNECT] Avoid nested parent overcount in RFormulaModel size estimate #57552

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
6932b5f
[WIP][ML] Avoid nested parent overcount in RFormulaModel size estimate
zhengruifeng Jul 27, 2026
de64138
[WIP][ML] Document RFormulaModel size estimate
zhengruifeng Jul 27, 2026
92f9d5a
[WIP][ML] Clarify RFormulaModel size estimate
zhengruifeng Jul 27, 2026
a81091c
[MINOR][ML] Add lower bound to RFormulaModel size test
zhengruifeng Jul 28, 2026
843b2e0
[MINOR][ML] Tighten RFormulaModel size test bound
zhengruifeng Jul 28, 2026
31b2a31
Revert "[MINOR][ML] Tighten RFormulaModel size test bound"
zhengruifeng Jul 28, 2026
390fd4d
Revert "[MINOR][ML] Add lower bound to RFormulaModel size test"
zhengruifeng Jul 28, 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
10 changes: 10 additions & 0 deletions mllib/src/main/scala/org/apache/spark/ml/feature/RFormula.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ import org.apache.spark.ml.util._
import org.apache.spark.sql.{DataFrame, Dataset}
import org.apache.spark.sql.functions.col
import org.apache.spark.sql.types._
import org.apache.spark.util.SizeEstimator

/**
* Base trait for [[RFormula]] and [[RFormulaModel]].
Expand Down Expand Up @@ -354,6 +355,15 @@ class RFormulaModel private[feature](
// For ml connect only
private[ml] def this() = this("", null, null)

private[spark] override def estimatedSize: Long = {
var size = estimateMatadataSize
// ResolvedRFormula(label: String, terms: Seq[Seq[String]], hasIntercept: Boolean)
size += SizeEstimator.estimate(resolvedFormula)
// pipelineModel: PipelineModel
size += pipelineModel.estimatedSize
size
}

@Since("2.0.0")
override def transform(dataset: Dataset[_]): DataFrame = {
checkCanTransform(dataset.schema)
Expand Down
9 changes: 9 additions & 0 deletions mllib/src/test/scala/org/apache/spark/ml/feature/RFormulaSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,15 @@ class RFormulaSuite extends MLTest with DefaultReadWriteTest {
ParamsSuite.checkParams(new RFormula())
}

test("RFormulaModel estimated size") {
val dataset = Seq(("a", 1.0), ("b", 2.0), ("c", 3.0)).toDF("category", "label")
val model = new RFormula().setFormula("label ~ category").fit(dataset)
val maxSize = 1024 * 16
assert(
model.estimatedSize < maxSize,
s"Estimation (${model.estimatedSize}) should be less than $maxSize")
}

@uros-b uros-b Jul 27, 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.

The new test asserts only an upper bound (estimatedSize < 16KiB). It genuinely fails without the fix (the default walk blows past 16KiB via SparkSession), so it is a real regression guard, but a buggy implementation returning 0 or a too-small value would also pass. Please add a lower bound to match that standard from predecessor PRs.


test("transform numeric data") {
val formula = new RFormula().setFormula("id ~ v1 + v2")
val original = Seq((0, 1.0, 3.0), (2, 2.0, 5.0)).toDF("id", "v1", "v2")
Expand Down