You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ru...@apache.org on 2026/09/11 00:13:00 UTC
(spark) branch master updated: [SPARK-59400][ML] Reuse TreeEnsembleModel prediction helper
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 619f4f9a60c3 [SPARK-59400][ML] Reuse TreeEnsembleModel prediction helper
619f4f9a60c3 is described below
commit 619f4f9a60c32af2ee872784fe8c15abc37d4542
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Fri Sep 11 08:12:46 2026 +0800
[SPARK-59400][ML] Reuse TreeEnsembleModel prediction helper
### What changes were proposed in this pull request?
This pull request centralizes tree ensemble prediction logic in companion objects:
- Add weighted and unweighted `TreeEnsembleModel.predictRaw` overloads that accept tree models.
- Rename the existing root-node helper to `predictRaw` for consistent naming.
- Reuse these helpers in GBT classification, GBT regression, and random forest regression.
- Move random forest classification's tree-based `predictRaw` implementation to its companion
object and delegate to it from the model.
### Why are the changes needed?
The model classes duplicated tree traversal and aggregation loops. Keeping the implementations in
companion-object helpers makes the direct model prediction paths consistent with the optimized
column prediction paths. The unweighted overload also avoids allocating an intermediate array when
random forest regression sums tree predictions.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
No new tests were added because this is a behavior-preserving refactor of prediction paths already
covered by the existing GBT and random forest suites.
The patch passed `git diff --check`, changed-file line-length checks, and changed-file non-ASCII
checks. Unit tests were not run locally.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (GPT-5)
Closes #58691 from zhengruifeng/ml-reuse-tree-ensemble-predict-dev-6.
Authored-by: Ruifeng Zheng <ru...@apache.org>
Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
---
.../spark/ml/classification/GBTClassifier.scala | 19 +++-------
.../ml/classification/RandomForestClassifier.scala | 42 +++++++++++-----------
.../apache/spark/ml/regression/GBTRegressor.scala | 15 ++------
.../ml/regression/RandomForestRegressor.scala | 8 ++---
.../org/apache/spark/ml/tree/treeModels.scala | 27 +++++++++++++-
5 files changed, 56 insertions(+), 55 deletions(-)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala
index ff7e449653c0..05e3d0ab2665 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala
@@ -369,7 +369,7 @@ class GBTClassificationModel private[ml](
}).apply(features)
} else {
udf((features: Vector) => {
- val margin = TreeEnsembleModel.predict(features, localRootNodes, localTreeWeights)
+ val margin = TreeEnsembleModel.predictRaw(features, localRootNodes, localTreeWeights)
if (margin > 0.0) 1.0 else 0.0
}).apply(features)
}
@@ -380,13 +380,13 @@ class GBTClassificationModel private[ml](
if (isDefined(thresholds)) {
super.predict(features)
} else {
- if (margin(features) > 0.0) 1.0 else 0.0
+ if (TreeEnsembleModel.predictRaw(features, _trees, _treeWeights) > 0.0) 1.0 else 0.0
}
}
@Since("3.0.0")
override def predictRaw(features: Vector): Vector = {
- val prediction: Double = margin(features)
+ val prediction = TreeEnsembleModel.predictRaw(features, _trees, _treeWeights)
Vectors.dense(Array(-prediction, prediction))
}
@@ -428,17 +428,6 @@ class GBTClassificationModel private[ml](
lazy val featureImportances: Vector =
TreeEnsembleModel.featureImportances(trees, numFeatures, perTreeNormalization = false)
- /** Raw prediction for the positive class. */
- private def margin(features: Vector): Double = {
- var prediction = 0.0
- var i = 0
- while (i < _trees.length) {
- prediction += _trees(i).rootNode.predictImpl(features).prediction * _treeWeights(i)
- i += 1
- }
- prediction
- }
-
/** (private[ml]) Convert to a model in the old API */
private[ml] def toOld: OldGBTModel = {
new OldGBTModel(OldAlgo.Classification, _trees.map(_.toOld), _treeWeights)
@@ -470,7 +459,7 @@ object GBTClassificationModel extends MLReadable[GBTClassificationModel] {
features: Vector,
rootNodes: Array[Node],
treeWeights: Array[Double]): Vector = {
- val prediction = TreeEnsembleModel.predict(features, rootNodes, treeWeights)
+ val prediction = TreeEnsembleModel.predictRaw(features, rootNodes, treeWeights)
Vectors.dense(-prediction, prediction)
}
diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala
index cf990923bc7f..26c5e2fc856d 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala
@@ -372,24 +372,8 @@ class RandomForestClassificationModel private[ml] (
}
@Since("3.0.0")
- override def predictRaw(features: Vector): Vector = {
- // TODO: When we add a generic Bagging class, handle transform there: SPARK-7128
- // Classifies using majority votes.
- // Ignore the tree weights since all are 1.0 for now.
- val votes = Array.ofDim[Double](numClasses)
- _trees.foreach { tree =>
- val classCounts = tree.rootNode.predictImpl(features).impurityStats.stats
- val total = classCounts.sum
- if (total != 0) {
- var i = 0
- while (i < numClasses) {
- votes(i) += classCounts(i) / total
- i += 1
- }
- }
- }
- Vectors.dense(votes)
- }
+ override def predictRaw(features: Vector): Vector =
+ RandomForestClassificationModel.predictRaw(features, _trees, numClasses)
override protected def raw2probabilityInPlace(rawPrediction: Vector): Vector = {
rawPrediction match {
@@ -469,13 +453,29 @@ class RandomForestClassificationModel private[ml] (
@Since("2.0.0")
object RandomForestClassificationModel extends MLReadable[RandomForestClassificationModel] {
+ private def predictRaw(
+ features: Vector,
+ trees: Array[DecisionTreeClassificationModel],
+ numClasses: Int): Vector = {
+ val votes = Array.ofDim[Double](numClasses)
+ trees.foreach { tree =>
+ val classCounts = tree.rootNode.predictImpl(features).impurityStats.stats
+ val total = classCounts.sum
+ if (total != 0) {
+ var i = 0
+ while (i < numClasses) {
+ votes(i) += classCounts(i) / total
+ i += 1
+ }
+ }
+ }
+ Vectors.dense(votes)
+ }
+
private def predictRaw(
features: Vector,
rootNodes: Array[Node],
numClasses: Int): Vector = {
- // TODO: When we add a generic Bagging class, handle transform there: SPARK-7128
- // Classifies using majority votes.
- // Ignore the tree weights since all are 1.0 for now.
val votes = Array.ofDim[Double](numClasses)
rootNodes.foreach { rootNode =>
val classCounts = rootNode.predictImpl(features).impurityStats.stats
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala
index 069cf39f19ad..c33ce2552b09 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala
@@ -280,7 +280,7 @@ class GBTRegressionModel private[ml](
if ($(predictionCol).nonEmpty) {
val predUDF = udf { features: Vector =>
val (rootNodes, treeWeights) = bcTreeData.value
- TreeEnsembleModel.predict(features, rootNodes, treeWeights)
+ TreeEnsembleModel.predictRaw(features, rootNodes, treeWeights)
}
predColNames :+= $(predictionCol)
predCols :+= predUDF(col($(featuresCol)))
@@ -304,17 +304,8 @@ class GBTRegressionModel private[ml](
}
}
- override def predict(features: Vector): Double = {
- // TODO: When we add a generic Boosting class, handle transform there? SPARK-7129
- // Classifies by thresholding sum of weighted tree predictions
- var prediction = 0.0
- var i = 0
- while (i < _trees.length) {
- prediction += _trees(i).rootNode.predictImpl(features).prediction * _treeWeights(i)
- i += 1
- }
- prediction
- }
+ override def predict(features: Vector): Double =
+ TreeEnsembleModel.predictRaw(features, _trees, _treeWeights)
@Since("1.4.0")
override def copy(extra: ParamMap): GBTRegressionModel = {
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/RandomForestRegressor.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/RandomForestRegressor.scala
index 4d9a039ecaca..d0c342b45020 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/RandomForestRegressor.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/RandomForestRegressor.scala
@@ -262,12 +262,8 @@ class RandomForestRegressionModel private[ml] (
}
}
- override def predict(features: Vector): Double = {
- // TODO: When we add a generic Bagging class, handle transform there. SPARK-7128
- // Predict average of tree predictions.
- // Ignore the weights since all are 1.0 for now.
- _trees.map(_.rootNode.predictImpl(features).prediction).sum / getNumTrees
- }
+ override def predict(features: Vector): Double =
+ TreeEnsembleModel.predictRaw(features, _trees) / getNumTrees
@Since("1.4.0")
override def copy(extra: ParamMap): RandomForestRegressionModel = {
diff --git a/mllib/src/main/scala/org/apache/spark/ml/tree/treeModels.scala b/mllib/src/main/scala/org/apache/spark/ml/tree/treeModels.scala
index d436fc896382..f970726760b2 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/tree/treeModels.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/tree/treeModels.scala
@@ -169,7 +169,32 @@ private[spark] trait TreeEnsembleModel[M <: DecisionTreeModel] {
private[ml] object TreeEnsembleModel {
- private[ml] def predict(
+ private[ml] def predictRaw[M <: DecisionTreeModel](
+ features: Vector,
+ trees: Array[M],
+ treeWeights: Array[Double]): Double = {
+ var prediction = 0.0
+ var i = 0
+ while (i < trees.length) {
+ prediction += trees(i).rootNode.predictImpl(features).prediction * treeWeights(i)
+ i += 1
+ }
+ prediction
+ }
+
+ private[ml] def predictRaw[M <: DecisionTreeModel](
+ features: Vector,
+ trees: Array[M]): Double = {
+ var prediction = 0.0
+ var i = 0
+ while (i < trees.length) {
+ prediction += trees(i).rootNode.predictImpl(features).prediction
+ i += 1
+ }
+ prediction
+ }
+
+ private[ml] def predictRaw(
features: Vector,
rootNodes: Array[Node],
treeWeights: Array[Double]): Double = {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org