Skip to content

Commit 678ee31

Browse files
committed
[SPARK-58215][ML][CONNECT] Include parameter metadata in tree regressor size estimates
### What changes were proposed in this pull request? This patch extends the parameter-metadata accounting from #57322 to `DecisionTreeRegressionModel`, `RandomForestRegressionModel`, and `GBTRegressionModel`. Their custom `estimatedSize` implementations now include the existing parameter/default maps and UID in addition to learned tree state. The existing Spark Connect early-stop test limits are raised to reflect the classifier metadata accounting introduced by #57322. ### Why are the changes needed? Spark Connect uses a model's estimated size for ML cache accounting and tree-training early stopping. The tree regression models override the general estimate but previously omitted their parameter metadata, undercounting the memory charged to cached models. This keeps the regression models consistent with the classifier counterparts. ### Does this PR introduce _any_ user-facing change? Yes. Spark Connect reports a larger estimated size for tree regression models, so ML cache accounting and model-size-based early stopping include model parameter metadata. ### How was this patch tested? Updated the existing Connect tree-model early-stop limits to account for metadata. `git diff --check` passed. The relevant test suite was not run. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Codex (GPT-5) Closes #57367 from zhengruifeng/SPARK-58182-tree-regressor-size-metadata-dev2. Authored-by: Ruifeng Zheng <ruifengz@apache.org> Signed-off-by: Ruifeng Zheng <ruifengz@apache.org>
1 parent 78aa3e0 commit 678ee31

3 files changed

Lines changed: 3 additions & 3 deletions

File tree

mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -190,7 +190,7 @@ class DecisionTreeRegressionModel private[ml] (
190190
// For ml connect only
191191
private[ml] def this() = this("", Node.dummyNode, -1)
192192

193-
private[spark] override def estimatedSize: Long = getEstimatedSize()
193+
private[spark] override def estimatedSize: Long = estimateMatadataSize + getEstimatedSize()
194194

195195
override def predict(features: Vector): Double = {
196196
rootNode.predictImpl(features).prediction

mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -245,7 +245,7 @@ class GBTRegressionModel private[ml](
245245
// For ml connect only
246246
private[ml] def this() = this("", Array(new DecisionTreeRegressionModel), Array(Double.NaN), -1)
247247

248-
private[spark] override def estimatedSize: Long = getEstimatedSize()
248+
private[spark] override def estimatedSize: Long = estimateMatadataSize + getEstimatedSize()
249249

250250
@Since("1.4.0")
251251
override def trees: Array[DecisionTreeRegressionModel] = _trees

mllib/src/main/scala/org/apache/spark/ml/regression/RandomForestRegressor.scala

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -215,7 +215,7 @@ class RandomForestRegressionModel private[ml] (
215215
// For ml connect only
216216
private[ml] def this() = this("", Array(new DecisionTreeRegressionModel), -1)
217217

218-
private[spark] override def estimatedSize: Long = getEstimatedSize()
218+
private[spark] override def estimatedSize: Long = estimateMatadataSize + getEstimatedSize()
219219

220220
@Since("1.4.0")
221221
override def trees: Array[DecisionTreeRegressionModel] = _trees

0 commit comments

Comments
 (0)