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 0f9f33e1f745 [SPARK-59038][ML] Avoid unused broadcasts when transform
has no output columns
0f9f33e1f745 is described below
commit 0f9f33e1f745400ee1bb6ca9c89327af25d4d6ea
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Fri Aug 28 10:00:29 2026 +0800
[SPARK-59038][ML] Avoid unused broadcasts when transform has no output
columns
### What changes were proposed in this pull request?
Create the model broadcast in `RandomForestRegressionModel.transform` and
`GBTRegressionModel.transform` only when at least one of `predictionCol` or
`leafCol` is set.
### Why are the changes needed?
When both output columns are empty, these transforms return the input
unchanged. Previously they
still serialized and broadcast the entire model, consuming unnecessary
driver resources despite the
broadcast never being referenced.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
No tests were added because the output and existing execution branches are
unchanged. The patch was
checked with `git diff --check` and source line-length and non-ASCII scans.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (GPT-5)
Closes #58331 from zhengruifeng/ml-spark-59038-noop-broadcast-dev-5.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../apache/spark/ml/regression/GBTRegressor.scala | 37 +++++++++++-----------
.../ml/regression/RandomForestRegressor.scala | 37 +++++++++++-----------
2 files changed, 36 insertions(+), 38 deletions(-)
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 2119555d408e..0ccb4c6379c7 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
@@ -271,27 +271,26 @@ class GBTRegressionModel private[ml](
override def transform(dataset: Dataset[_]): DataFrame = {
val outputSchema = transformSchema(dataset.schema, logging = true)
- var predictionColNames = Seq.empty[String]
- var predictionColumns = Seq.empty[Column]
-
- val bcastModel = dataset.sparkSession.sparkContext.broadcast(this)
-
- if ($(predictionCol).nonEmpty) {
- val predictUDF = udf { features: Vector =>
bcastModel.value.predict(features) }
- predictionColNames :+= $(predictionCol)
- predictionColumns :+= predictUDF(col($(featuresCol)))
- .as($(predictionCol), outputSchema($(predictionCol)).metadata)
- }
+ if ($(predictionCol).nonEmpty || $(leafCol).nonEmpty) {
+ var predColNames = Seq.empty[String]
+ var predCols = Seq.empty[Column]
+ val bcModel = dataset.sparkSession.sparkContext.broadcast(this)
+
+ if ($(predictionCol).nonEmpty) {
+ val predUDF = udf { features: Vector =>
bcModel.value.predict(features) }
+ predColNames :+= $(predictionCol)
+ predCols :+= predUDF(col($(featuresCol)))
+ .as($(predictionCol), outputSchema($(predictionCol)).metadata)
+ }
- if ($(leafCol).nonEmpty) {
- val leafUDF = udf { features: Vector =>
bcastModel.value.predictLeaf(features) }
- predictionColNames :+= $(leafCol)
- predictionColumns :+= leafUDF(col($(featuresCol)))
- .as($(leafCol), outputSchema($(leafCol)).metadata)
- }
+ if ($(leafCol).nonEmpty) {
+ val leafUDF = udf { features: Vector =>
bcModel.value.predictLeaf(features) }
+ predColNames :+= $(leafCol)
+ predCols :+= leafUDF(col($(featuresCol)))
+ .as($(leafCol), outputSchema($(leafCol)).metadata)
+ }
- if (predictionColNames.nonEmpty) {
- dataset.withColumns(predictionColNames, predictionColumns)
+ dataset.withColumns(predColNames, predCols)
} else {
this.logWarning(log"${MDC(LogKeys.UUID, uid)}:
GBTRegressionModel.transform() " +
log"does nothing because no output columns were set.")
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 7826c7843504..9d42d601b1f1 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
@@ -238,27 +238,26 @@ class RandomForestRegressionModel private[ml] (
override def transform(dataset: Dataset[_]): DataFrame = {
val outputSchema = transformSchema(dataset.schema, logging = true)
- var predictionColNames = Seq.empty[String]
- var predictionColumns = Seq.empty[Column]
-
- val bcastModel = dataset.sparkSession.sparkContext.broadcast(this)
-
- if ($(predictionCol).nonEmpty) {
- val predictUDF = udf { features: Vector =>
bcastModel.value.predict(features) }
- predictionColNames :+= $(predictionCol)
- predictionColumns :+= predictUDF(col($(featuresCol)))
- .as($(predictionCol), outputSchema($(predictionCol)).metadata)
- }
+ if ($(predictionCol).nonEmpty || $(leafCol).nonEmpty) {
+ var predColNames = Seq.empty[String]
+ var predCols = Seq.empty[Column]
+ val bcModel = dataset.sparkSession.sparkContext.broadcast(this)
+
+ if ($(predictionCol).nonEmpty) {
+ val predUDF = udf { features: Vector =>
bcModel.value.predict(features) }
+ predColNames :+= $(predictionCol)
+ predCols :+= predUDF(col($(featuresCol)))
+ .as($(predictionCol), outputSchema($(predictionCol)).metadata)
+ }
- if ($(leafCol).nonEmpty) {
- val leafUDF = udf { features: Vector =>
bcastModel.value.predictLeaf(features) }
- predictionColNames :+= $(leafCol)
- predictionColumns :+= leafUDF(col($(featuresCol)))
- .as($(leafCol), outputSchema($(leafCol)).metadata)
- }
+ if ($(leafCol).nonEmpty) {
+ val leafUDF = udf { features: Vector =>
bcModel.value.predictLeaf(features) }
+ predColNames :+= $(leafCol)
+ predCols :+= leafUDF(col($(featuresCol)))
+ .as($(leafCol), outputSchema($(leafCol)).metadata)
+ }
- if (predictionColNames.nonEmpty) {
- dataset.withColumns(predictionColNames, predictionColumns)
+ dataset.withColumns(predColNames, predCols)
} else {
this.logWarning(log"${MDC(LogKeys.UUID, uid)}:
RandomForestRegressionModel.transform() " +
log"does nothing because no output columns were set.")
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
