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/09 05:17:47 UTC
(spark) branch branch-4.x updated: [SPARK-59335][ML] Reduce GBTRegressionModel broadcast size
This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new f34c4ee2b0fa [SPARK-59335][ML] Reduce GBTRegressionModel broadcast size
f34c4ee2b0fa is described below
commit f34c4ee2b0fac776b7fa7997e5f3517c4183a8ec
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Wed Sep 9 13:17:20 2026 +0800
[SPARK-59335][ML] Reduce GBTRegressionModel broadcast size
### What changes were proposed in this pull request?
This PR reduces the broadcast payload created by `GBTRegressionModel.transform`.
Instead of broadcasting the complete model, the transform broadcasts only the tree root nodes and
a cloned tree-weight array. Prediction uses a companion-object helper over that state, while leaf
prediction uses the existing `TreeEnsembleModel.predictLeaf` helper.
### Why are the changes needed?
Prediction and leaf traversal need the root nodes, and prediction additionally needs the tree
weights. Broadcasting the complete model also serializes unrelated state such as the model and
tree-model parameter graphs, UIDs, and metadata. Avoiding that state reduces driver and executor
memory pressure for long-lived Spark Connect servers.
A temporary local probe trained a 20-tree model with 30 total nodes and serialized each broadcast
payload using Spark's configured serializer. The full model required 53,494 bytes, while the root
nodes and tree weights required 3,284 bytes, a 93.9% reduction.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
The following checks passed:
```
build/sbt mllib/compile
build/sbt 'mllib/testOnly org.apache.spark.ml.regression.GBTRegressorSuite'
```
`GBTRegressorSuite` ran 16 tests covering transform prediction and leaf-index output. No new test
was added because the change only narrows the serialized state used by those existing code paths.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (GPT-5)
Closes #58623 from zhengruifeng/ml-spark-59335-gbt-broadcast-dev-7.
Authored-by: Ruifeng Zheng <ru...@apache.org>
Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
(cherry picked from commit 0e2f610cc834116fbc17be8795e313e6643e87e8)
Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
---
.../apache/spark/ml/classification/GBTClassifier.scala | 17 ++---------------
.../org/apache/spark/ml/regression/GBTRegressor.scala | 12 +++++++++---
.../scala/org/apache/spark/ml/tree/treeModels.scala | 13 +++++++++++++
3 files changed, 24 insertions(+), 18 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 f6fcab572498..ff7e449653c0 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 = GBTClassificationModel.margin(features, localRootNodes, localTreeWeights)
+ val margin = TreeEnsembleModel.predict(features, localRootNodes, localTreeWeights)
if (margin > 0.0) 1.0 else 0.0
}).apply(features)
}
@@ -466,24 +466,11 @@ 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)
+ val prediction = TreeEnsembleModel.predict(features, rootNodes, treeWeights)
Vectors.dense(-prediction, prediction)
}
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 d60e6dfc7c2c..069cf39f19ad 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
@@ -274,17 +274,23 @@ class GBTRegressionModel private[ml](
if ($(predictionCol).nonEmpty || $(leafCol).nonEmpty) {
var predColNames = Seq.empty[String]
var predCols = Seq.empty[Column]
- val bcModel = dataset.sparkSession.sparkContext.broadcast(this)
+ val bcTreeData = dataset.sparkSession.sparkContext.broadcast(
+ (_trees.map(_.rootNode), _treeWeights.clone()))
if ($(predictionCol).nonEmpty) {
- val predUDF = udf { features: Vector => bcModel.value.predict(features) }
+ val predUDF = udf { features: Vector =>
+ val (rootNodes, treeWeights) = bcTreeData.value
+ TreeEnsembleModel.predict(features, rootNodes, treeWeights)
+ }
predColNames :+= $(predictionCol)
predCols :+= predUDF(col($(featuresCol)))
.as($(predictionCol), outputSchema($(predictionCol)).metadata)
}
if ($(leafCol).nonEmpty) {
- val leafUDF = udf { features: Vector => bcModel.value.predictLeaf(features) }
+ val leafUDF = udf { features: Vector =>
+ TreeEnsembleModel.predictLeaf(features, bcTreeData.value._1)
+ }
predColNames :+= $(leafCol)
predCols :+= leafUDF(col($(featuresCol)))
.as($(leafCol), outputSchema($(leafCol)).metadata)
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 99ed9ef8dc46..d436fc896382 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,19 @@ private[spark] trait TreeEnsembleModel[M <: DecisionTreeModel] {
private[ml] object TreeEnsembleModel {
+ private[ml] def predict(
+ 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[ml] def predictLeaf(features: Vector, rootNodes: Array[Node]): Vector = {
val indices = Array.ofDim[Double](rootNodes.length)
var i = 0
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org