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 01:53:49 UTC
(spark) branch master updated: [SPARK-59338][ML] Move predictionColumn to PredictionModel
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 b25686edb87c [SPARK-59338][ML] Move predictionColumn to PredictionModel
b25686edb87c is described below
commit b25686edb87c3fb2d3536054aa7ab3b59381ada9
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Wed Sep 9 09:53:37 2026 +0800
[SPARK-59338][ML] Move predictionColumn to PredictionModel
### What changes were proposed in this pull request?
Move the `predictionColumn` hook from `ClassificationModel` to `PredictionModel`, and use it in
`PredictionModel.transformImpl`.
The existing classification-specific overrides continue to use the same hook. A focused test
verifies that the generic `PredictionModel` transform path invokes an overridden
`predictionColumn`.
### Why are the changes needed?
Prediction-column expression customization is useful to prediction models generally, not only to
classification models. Defining the hook in `PredictionModel` lets other model families provide a
native expression or a UDF that snapshots only the prediction state they need.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Added a unit test and ran:
```
JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64 build/sbt \
'mllib/testOnly org.apache.spark.ml.PredictorSuite' \
'mllib/testOnly org.apache.spark.ml.classification.ClassifierSuite'
```
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (GPT-5)
Closes #58625 from zhengruifeng/ml-spark-59338-prediction-column-dev-7.
Authored-by: Ruifeng Zheng <ru...@apache.org>
Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
---
.../main/scala/org/apache/spark/ml/Predictor.scala | 19 ++++++++++++++-----
.../apache/spark/ml/classification/Classifier.scala | 12 ------------
.../scala/org/apache/spark/ml/PredictorSuite.scala | 12 +++++++++++-
3 files changed, 25 insertions(+), 18 deletions(-)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/Predictor.scala b/mllib/src/main/scala/org/apache/spark/ml/Predictor.scala
index 38d4c43bc154..8328b044e510 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/Predictor.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/Predictor.scala
@@ -23,7 +23,7 @@ import org.apache.spark.ml.linalg.SQLDataTypes
import org.apache.spark.ml.param._
import org.apache.spark.ml.param.shared._
import org.apache.spark.ml.util.SchemaUtils
-import org.apache.spark.sql.{DataFrame, Dataset}
+import org.apache.spark.sql.{Column, DataFrame, Dataset}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types.{DataType, DoubleType, StructType}
@@ -181,6 +181,18 @@ abstract class PredictionModel[FeaturesType, M <: PredictionModel[FeaturesType,
outputSchema
}
+ /**
+ * Returns an expression that produces a predicted label directly from a features column.
+ * The default wraps [[predict]] in a UDF. Models may override this with a native expression or
+ * a UDF that snapshots prediction state.
+ *
+ * @param features input features column
+ * @return prediction column of type `Double`
+ */
+ protected def predictionColumn(features: Column): Column = {
+ udf { value: Any => predict(value.asInstanceOf[FeaturesType]) }.apply(features)
+ }
+
/**
* Transforms dataset by reading from [[featuresCol]], calling `predict`, and storing
* the predictions as a new column [[predictionCol]].
@@ -201,10 +213,7 @@ abstract class PredictionModel[FeaturesType, M <: PredictionModel[FeaturesType,
protected def transformImpl(dataset: Dataset[_]): DataFrame = {
val outputSchema = transformSchema(dataset.schema, logging = true)
- val predictUDF = udf { features: Any =>
- predict(features.asInstanceOf[FeaturesType])
- }
- dataset.withColumn($(predictionCol), predictUDF(col($(featuresCol))),
+ dataset.withColumn($(predictionCol), predictionColumn(col($(featuresCol))),
outputSchema($(predictionCol)).metadata)
}
diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/Classifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/Classifier.scala
index 5e29d58a905c..7a53331cad78 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/Classifier.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/Classifier.scala
@@ -151,18 +151,6 @@ abstract class ClassificationModel[FeaturesType, M <: ClassificationModel[Featur
udf(raw2prediction _).apply(rawPrediction)
}
- /**
- * Returns an expression that produces a predicted label directly from a features column.
- *
- * `transform` uses this when [[predictionCol]] is set and [[rawPredictionCol]] is not set.
- *
- * @param features input features column
- * @return prediction column of type `Double`
- */
- protected def predictionColumn(features: Column): Column = {
- udf { value: Any => predict(value.asInstanceOf[FeaturesType]) }.apply(features)
- }
-
/** @group setParam */
def setRawPredictionCol(value: String): M = set(rawPredictionCol, value).asInstanceOf[M]
diff --git a/mllib/src/test/scala/org/apache/spark/ml/PredictorSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/PredictorSuite.scala
index dff00eade620..a91aed76f787 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/PredictorSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/PredictorSuite.scala
@@ -23,7 +23,7 @@ import org.apache.spark.ml.param.ParamMap
import org.apache.spark.ml.param.shared.HasWeightCol
import org.apache.spark.ml.util._
import org.apache.spark.mllib.util.MLlibTestSparkContext
-import org.apache.spark.sql.Dataset
+import org.apache.spark.sql.{Column, Dataset}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types._
@@ -55,6 +55,14 @@ class PredictorSuite extends SparkFunSuite with MLlibTestSparkContext {
predictor.fit(df.select(col("label"), col("weight").cast(StringType), col("features")))
}
}
+
+ test("PredictionModel transform should use predictionColumn") {
+ val df = spark.createDataFrame(Seq(Tuple1(Vectors.dense(1.0)))).toDF("features")
+ val prediction = new MockPredictionModel()
+ .transform(df).select("prediction").head().getDouble(0)
+
+ assert(prediction === 1.0)
+ }
}
object PredictorSuite {
@@ -84,6 +92,8 @@ object PredictorSuite {
override def predict(features: Vector): Double =
throw new UnsupportedOperationException()
+ override protected def predictionColumn(features: Column): Column = lit(1.0)
+
override def copy(extra: ParamMap): MockPredictionModel =
throw new UnsupportedOperationException()
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org