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/07 02:00:05 UTC
(spark) branch master updated: [SPARK-59247][ML] Optimize tree ensemble classification model transform closures
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 bc54e7a2c19a [SPARK-59247][ML] Optimize tree ensemble classification model transform closures
bc54e7a2c19a is described below
commit bc54e7a2c19a1806aea0f79471471e72ce41b1b5
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Mon Sep 7 09:59:51 2026 +0800
[SPARK-59247][ML] Optimize tree ensemble classification model transform closures
### What changes were proposed in this pull request?
This PR reduces the transform closure size of `RandomForestClassificationModel` and
`GBTClassificationModel`. It is the tree-ensemble counterpart to #58454.
The column-expression prediction hooks snapshot only the root nodes, tree weights, class count,
loss, and thresholds needed by each requested output column. Companion-object helpers perform raw
prediction and probability conversion without retaining the complete model. A shared
`TreeEnsembleModel` helper similarly detaches leaf prediction from the model.
### Why are the changes needed?
The default probabilistic classifier hooks and existing leaf UDF invoke bound model methods. Their
closures therefore retain the complete model and parameter graph even though scoring only needs
the tree nodes and a small amount of immutable prediction state. This adds avoidable driver memory
pressure for long-lived Spark Connect servers.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
The following check passed:
```
build/sbt mllib/compile
```
No new tests were added because the existing `testPredictMethods` coverage in
`RandomForestClassifierSuite` and `GBTClassifierSuite` exercises all combinations of raw
prediction, probability, and prediction columns.
A temporary local probe trained each model on the same deterministic dataset with 256 rows and
16 features. Both models used 20 trees; Random Forest used maximum depth 5 and GBT used maximum
depth 3. The probe extracted each `ScalaUDF.function` from the analyzed transform plan and
serialized it with Spark's closure serializer. "All" is the sum of the raw-prediction,
probability, prediction, and leaf UDF closures.
| Model/output | Before | After | Reduction |
|---|---:|---:|---:|
| Random Forest raw prediction | 75,171 B | 20,951 B | 72.1% |
| Random Forest probability | 75,204 B | 20,959 B | 72.1% |
| Random Forest prediction | 75,108 B | 20,925 B | 72.1% |
| Random Forest leaf | 75,223 B | 20,866 B | 72.3% |
| Random Forest all | 300,846 B | 46,059 B | 84.7% |
| GBT raw prediction | 71,108 B | 21,164 B | 70.2% |
| GBT probability | 71,141 B | 21,357 B | 70.0% |
| GBT prediction | 71,045 B | 21,138 B | 70.2% |
| GBT leaf | 71,142 B | 20,985 B | 70.5% |
| GBT all | 284,576 B | 46,814 B | 83.5% |
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #58519 from zhengruifeng/ml-spark-59247-tree-ensemble-closure-dev-5.
Authored-by: Ruifeng Zheng <ru...@apache.org>
Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
---
.../spark/ml/classification/GBTClassifier.scala | 95 +++++++++++++++++++++-
.../ml/classification/RandomForestClassifier.scala | 88 +++++++++++++++++++-
.../org/apache/spark/ml/tree/treeModels.scala | 10 +++
3 files changed, 191 insertions(+), 2 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 19e49f299cbe..f6fcab572498 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
@@ -32,6 +32,7 @@ import org.apache.spark.ml.util.DatasetUtils.extractInstances
import org.apache.spark.ml.util.DefaultParamsReader.Metadata
import org.apache.spark.ml.util.Instrumentation.instrumented
import org.apache.spark.mllib.tree.configuration.{Algo => OldAlgo}
+import org.apache.spark.mllib.tree.loss.{ClassificationLoss => OldClassificationLoss}
import org.apache.spark.mllib.tree.model.{GradientBoostedTreesModel => OldGBTModel}
import org.apache.spark.sql._
import org.apache.spark.sql.functions._
@@ -304,7 +305,10 @@ class GBTClassificationModel private[ml](
val outputData = super.transform(dataset)
if ($(leafCol).nonEmpty) {
- val leafUDF = udf { features: Vector => predictLeaf(features) }
+ val localRootNodes = _trees.map(_.rootNode)
+ val leafUDF = udf { features: Vector =>
+ TreeEnsembleModel.predictLeaf(features, localRootNodes)
+ }
outputData.withColumn($(leafCol), leafUDF(col($(featuresCol))),
outputSchema($(leafCol)).metadata)
} else {
@@ -312,6 +316,65 @@ class GBTClassificationModel private[ml](
}
}
+ override protected def predictRawColumn(features: Column): Column = {
+ val localRootNodes = _trees.map(_.rootNode)
+ val localTreeWeights = _treeWeights.clone()
+ udf((features: Vector) =>
+ GBTClassificationModel.predictRaw(features, localRootNodes, localTreeWeights)
+ ).apply(features)
+ }
+
+ override protected def raw2probabilityColumn(rawPrediction: Column): Column = {
+ val localLoss = loss
+ udf((rawPrediction: Vector) =>
+ GBTClassificationModel.raw2probability(rawPrediction, localLoss)
+ ).apply(rawPrediction)
+ }
+
+ override protected def predictProbabilityColumn(features: Column): Column = {
+ val localRootNodes = _trees.map(_.rootNode)
+ val localTreeWeights = _treeWeights.clone()
+ val localLoss = loss
+ udf((features: Vector) => {
+ val rawPrediction =
+ GBTClassificationModel.predictRaw(features, localRootNodes, localTreeWeights)
+ GBTClassificationModel.raw2probability(rawPrediction, localLoss)
+ }).apply(features)
+ }
+
+ override protected def raw2predictionColumn(rawPrediction: Column): Column = {
+ if (isDefined(thresholds)) {
+ val localThresholds = getThresholds.clone()
+ val localLoss = loss
+ udf((rawPrediction: Vector) => {
+ val probability = GBTClassificationModel.raw2probability(rawPrediction, localLoss)
+ ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+ }).apply(rawPrediction)
+ } else {
+ udf((rawPrediction: Vector) => rawPrediction.argmax.toDouble).apply(rawPrediction)
+ }
+ }
+
+ override protected def predictionColumn(features: Column): Column = {
+ val localRootNodes = _trees.map(_.rootNode)
+ val localTreeWeights = _treeWeights.clone()
+ if (isDefined(thresholds)) {
+ val localThresholds = getThresholds.clone()
+ val localLoss = loss
+ udf((features: Vector) => {
+ val rawPrediction =
+ GBTClassificationModel.predictRaw(features, localRootNodes, localTreeWeights)
+ val probability = GBTClassificationModel.raw2probability(rawPrediction, localLoss)
+ ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+ }).apply(features)
+ } else {
+ udf((features: Vector) => {
+ val margin = GBTClassificationModel.margin(features, localRootNodes, localTreeWeights)
+ if (margin > 0.0) 1.0 else 0.0
+ }).apply(features)
+ }
+ }
+
override def predict(features: Vector): Double = {
// If thresholds defined, use predictRaw to get probabilities, otherwise use optimization
if (isDefined(thresholds)) {
@@ -403,6 +466,36 @@ class GBTClassificationModel private[ml](
@Since("2.0.0")
object GBTClassificationModel extends MLReadable[GBTClassificationModel] {
+ private def margin(
+ features: Vector,
+ rootNodes: Array[Node],
+ treeWeights: Array[Double]): Double = {
+ var prediction = 0.0
+ var i = 0
+ while (i < rootNodes.length) {
+ prediction += rootNodes(i).predictImpl(features).prediction * treeWeights(i)
+ i += 1
+ }
+ prediction
+ }
+
+ private def predictRaw(
+ features: Vector,
+ rootNodes: Array[Node],
+ treeWeights: Array[Double]): Vector = {
+ val prediction = margin(features, rootNodes, treeWeights)
+ Vectors.dense(-prediction, prediction)
+ }
+
+ private def raw2probability(
+ rawPrediction: Vector,
+ loss: OldClassificationLoss): Vector = {
+ val probability = rawPrediction.copy.toDense
+ probability.values(0) = loss.computeProbability(probability.values(0))
+ probability.values(1) = 1.0 - probability.values(0)
+ probability
+ }
+
private val numFeaturesKey: String = "numFeatures"
private val numTreesKey: String = "numTrees"
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 41bcf3729e5a..cf990923bc7f 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
@@ -305,7 +305,10 @@ class RandomForestClassificationModel private[ml] (
val outputData = super.transform(dataset)
if ($(leafCol).nonEmpty) {
- val leafUDF = udf { features: Vector => predictLeaf(features) }
+ val localRootNodes = _trees.map(_.rootNode)
+ val leafUDF = udf { features: Vector =>
+ TreeEnsembleModel.predictLeaf(features, localRootNodes)
+ }
outputData.withColumn($(leafCol), leafUDF(col($(featuresCol))),
outputSchema($(leafCol)).metadata)
} else {
@@ -313,6 +316,61 @@ class RandomForestClassificationModel private[ml] (
}
}
+ override protected def predictRawColumn(features: Column): Column = {
+ val localRootNodes = _trees.map(_.rootNode)
+ val localNumClasses = numClasses
+ udf((features: Vector) =>
+ RandomForestClassificationModel.predictRaw(features, localRootNodes, localNumClasses)
+ ).apply(features)
+ }
+
+ override protected def raw2probabilityColumn(rawPrediction: Column): Column = {
+ udf((rawPrediction: Vector) =>
+ RandomForestClassificationModel.raw2probability(rawPrediction)
+ ).apply(rawPrediction)
+ }
+
+ override protected def predictProbabilityColumn(features: Column): Column = {
+ val localRootNodes = _trees.map(_.rootNode)
+ val localNumClasses = numClasses
+ udf((features: Vector) => {
+ val rawPrediction =
+ RandomForestClassificationModel.predictRaw(features, localRootNodes, localNumClasses)
+ RandomForestClassificationModel.raw2probability(rawPrediction)
+ }).apply(features)
+ }
+
+ override protected def raw2predictionColumn(rawPrediction: Column): Column = {
+ if (isDefined(thresholds)) {
+ val localThresholds = getThresholds.clone()
+ udf((rawPrediction: Vector) => {
+ val probability = RandomForestClassificationModel.raw2probability(rawPrediction)
+ ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+ }).apply(rawPrediction)
+ } else {
+ udf((rawPrediction: Vector) => rawPrediction.argmax.toDouble).apply(rawPrediction)
+ }
+ }
+
+ override protected def predictionColumn(features: Column): Column = {
+ val localRootNodes = _trees.map(_.rootNode)
+ val localNumClasses = numClasses
+ if (isDefined(thresholds)) {
+ val localThresholds = getThresholds.clone()
+ udf((features: Vector) => {
+ val rawPrediction =
+ RandomForestClassificationModel.predictRaw(features, localRootNodes, localNumClasses)
+ val probability = RandomForestClassificationModel.raw2probability(rawPrediction)
+ ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+ }).apply(features)
+ } else {
+ udf((features: Vector) =>
+ RandomForestClassificationModel.predictRaw(
+ features, localRootNodes, localNumClasses).argmax.toDouble
+ ).apply(features)
+ }
+ }
+
@Since("3.0.0")
override def predictRaw(features: Vector): Vector = {
// TODO: When we add a generic Bagging class, handle transform there: SPARK-7128
@@ -411,6 +469,34 @@ class RandomForestClassificationModel private[ml] (
@Since("2.0.0")
object RandomForestClassificationModel extends MLReadable[RandomForestClassificationModel] {
+ 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
+ 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 raw2probability(rawPrediction: Vector): Vector = {
+ val probability = rawPrediction.copy.toDense
+ ProbabilisticClassificationModel.normalizeToProbabilitiesInPlace(probability)
+ probability
+ }
+
@Since("2.0.0")
override def read: MLReader[RandomForestClassificationModel] =
new RandomForestClassificationModelReader
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 25d035e81946..99ed9ef8dc46 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,6 +169,16 @@ private[spark] trait TreeEnsembleModel[M <: DecisionTreeModel] {
private[ml] object TreeEnsembleModel {
+ private[ml] def predictLeaf(features: Vector, rootNodes: Array[Node]): Vector = {
+ val indices = Array.ofDim[Double](rootNodes.length)
+ var i = 0
+ while (i < rootNodes.length) {
+ indices(i) = DecisionTreeModel.predictLeaf(features, rootNodes(i))
+ i += 1
+ }
+ Vectors.dense(indices)
+ }
+
/**
* Given a tree ensemble model, compute the importance of each feature.
* This generalizes the idea of "Gini" importance to other losses,
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org