This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch branch-4.1
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.1 by this push:
new 743271a26fc4 [SPARK-58994][ML] Fix prediction column metadata in
GBTRegressionModel
743271a26fc4 is described below
commit 743271a26fc496e004d778943f9b7fd80660ba60
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Wed Aug 26 08:36:02 2026 +0800
[SPARK-58994][ML] Fix prediction column metadata in GBTRegressionModel
### What changes were proposed in this pull request?
This PR fixes the prediction column alias in `GBTRegressionModel.transform`
to use
`predictionCol` and its output metadata. It also adds a regression test
that verifies the
transformed prediction metadata matches `transformSchema`.
### Why are the changes needed?
The prediction expression was aliased with `featuresCol` and the feature
column metadata.
Although `withColumns` assigned the final prediction column name, its outer
alias inherited the
feature metadata from the inner alias. As a result, the transformed
prediction column could carry
vector feature metadata instead of its expected numeric prediction metadata.
### Does this PR introduce _any_ user-facing change?
Yes. `GBTRegressionModel.transform` now attaches the correct numeric
metadata to the prediction
column. Prediction values, data types, and column names are unchanged.
### How was this patch tested?
Added a regression test and ran:
`./build/sbt 'mllib/testOnly
org.apache.spark.ml.regression.GBTRegressorSuite'`
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (GPT-5)
Closes #58277 from zhengruifeng/ml-fix-gbt-prediction-metadata-dev-5.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit 0ab7b7b1468dbb50b8718a1b421042a150ec8708)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../scala/org/apache/spark/ml/regression/GBTRegressor.scala | 2 +-
.../org/apache/spark/ml/regression/GBTRegressorSuite.scala | 11 +++++++++++
2 files changed, 12 insertions(+), 1 deletion(-)
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 71436036d1ea..23d86ac277c7 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](
val predictUDF = udf { features: Vector =>
bcastModel.value.predict(features) }
predictionColNames :+= $(predictionCol)
predictionColumns :+= predictUDF(col($(featuresCol)))
- .as($(featuresCol), outputSchema($(featuresCol)).metadata)
+ .as($(predictionCol), outputSchema($(predictionCol)).metadata)
}
if ($(leafCol).nonEmpty) {
diff --git
a/mllib/src/test/scala/org/apache/spark/ml/regression/GBTRegressorSuite.scala
b/mllib/src/test/scala/org/apache/spark/ml/regression/GBTRegressorSuite.scala
index d7f15dc2cfe9..07a11fa92caf 100644
---
a/mllib/src/test/scala/org/apache/spark/ml/regression/GBTRegressorSuite.scala
+++
b/mllib/src/test/scala/org/apache/spark/ml/regression/GBTRegressorSuite.scala
@@ -119,6 +119,17 @@ class GBTRegressorSuite extends MLTest with
DefaultReadWriteTest {
testPredictionModelSinglePrediction(model, validationData.toDF())
}
+ test("prediction column has correct metadata") {
+ val tree = new DecisionTreeRegressionModel("dtr", TreeTests.root0, 3)
+ val model = new GBTRegressionModel("gbtr", Array(tree), Array(1.0), 3)
+ val df = sc.parallelize(TreeTests.getTwoTreesLeafData.toImmutableArraySeq,
1)
+ .toDF("leafId", "features")
+
+ val expectedMetadata =
model.transformSchema(df.schema)(model.getPredictionCol).metadata
+ val actualMetadata =
model.transform(df).schema(model.getPredictionCol).metadata
+ assert(actualMetadata === expectedMetadata)
+ }
+
test("Checkpointing") {
val tempDir = Utils.createTempDir()
val path = tempDir.toURI.toString
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
