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/01 01:03:09 UTC
(spark) branch branch-4.x updated: [SPARK-59048][ML] Optimize ProbabilisticClassificationModel by reducing closure size and accelerating transform
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 da266c169b0c [SPARK-59048][ML] Optimize ProbabilisticClassificationModel by reducing closure size and accelerating transform
da266c169b0c is described below
commit da266c169b0c7a37a07c9f1896991979c5e1989a
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Tue Sep 1 09:02:30 2026 +0800
[SPARK-59048][ML] Optimize ProbabilisticClassificationModel by reducing closure size and accelerating transform
### What changes were proposed in this pull request?
This PR adds protected column-expression hooks to `ClassificationModel` and
`ProbabilisticClassificationModel`. The shared transform implementations use these hooks to build
raw prediction, probability, and prediction columns. Their default implementations preserve the
existing row-based UDF behavior, while individual models can provide native Catalyst expressions or
UDFs backed by compact prediction state.
`LogisticRegressionModel` uses the new hooks to snapshot only coefficients, intercepts, mode, and
thresholds needed by each expression. Binary raw/probability thresholding uses native column
expressions, while the remaining UDFs call companion-object helpers without retaining the model.
### Why are the changes needed?
The previous classifier transform UDFs invoked bound model methods. As a result, their closures
retained the full model, including its parameter graph, even when a prediction stage needed only a
small subset of immutable scoring state.
The column-expression hooks provide a common design for classifiers to reduce those closures and to
use native Spark expressions where possible. This lowers driver memory pressure from retained query
plans on long-lived Spark Connect servers.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
The following checks were run locally:
```
build/sbt mllib/compile
build/sbt mllib/scalastyle
build/sbt 'mllib/testOnly org.apache.spark.ml.classification.ClassifierSuite org.apache.spark.ml.classification.LogisticRegressionSuite'
build/sbt 'mllib/testOnly org.apache.spark.ml.classification.ProbabilisticClassifierSuite org.apache.spark.ml.classification.LogisticRegressionSuite'
```
The focused test run completed 66 tests successfully, with one existing ignored test.
A temporary local probe compared the new implementation with a model subclass that recreated the
previous bound-method UDF expressions and cached binary prediction state. It serialized each
`ScalaUDF.function` with Spark's closure serializer for binary models with 64 and 4,096 features:
| Features | Output columns | Before | After | Delta |
|---:|---|---:|---:|---:|
| 64 | raw prediction | 9,869 B | 2,877 B | -70.8% |
| 64 | probability | 9,877 B | 2,885 B | -70.8% |
| 64 | prediction | 9,665 B | 2,866 B | -70.3% |
| 64 | raw + probability + prediction | 29,418 B (3 UDFs) | 4,969 B (2 UDFs) | -83.1% |
| 4,096 | raw prediction | 42,125 B | 35,133 B | -16.6% |
| 4,096 | probability | 42,133 B | 35,141 B | -16.6% |
| 4,096 | prediction | 41,921 B | 35,122 B | -16.2% |
| 4,096 | raw + probability + prediction | 126,186 B (3 UDFs) | 37,225 B (2 UDFs) | -70.5% |
The same in-process probe excluded SparkSession setup, ran 3 alternating warm-up pairs, and then
measured 10 alternating before/after action pairs for each workload. Both workloads used 3,000,000
rows. The input was left uncached because caching the 4,096-feature workload would materialize
about 98 GB of vector values. Mean execution times were:
| Features | Transform path | Before | After | Change |
|---:|---|---:|---:|---:|
| 64 | prediction only | 278.127 ms | 241.504 ms | 1.15x faster |
| 64 | raw prediction to prediction | 629.430 ms | 599.644 ms | 1.05x faster |
| 64 | probability to prediction | 595.218 ms | 562.201 ms | 1.06x faster |
| 4,096 | prediction only | 7,551.752 ms | 7,515.061 ms | 0.5% faster |
| 4,096 | raw prediction to prediction | 8,191.263 ms | 7,988.399 ms | 1.03x faster |
| 4,096 | probability to prediction | 8,167.511 ms | 8,026.213 ms | 1.02x faster |
Feature dimensionality changes where transform time is spent. With 4,096 features, vector
computation dominates the total duration, so avoiding per-row parameter lookup has a smaller
relative effect. With 64 features, vector computation is cheaper and parameter lookup contributes
more to the total duration, making the optimization more visible.
The probe was used only for measurement and is not included in this PR.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #58336 from zhengruifeng/ml-spark-59048-classifier-closure-dev-4.
Authored-by: Ruifeng Zheng <ru...@apache.org>
Signed-off-by: Ruifeng Zheng <ru...@apache.org>
(cherry picked from commit 2948ffa5589f65ad4885113031a8d2c1d681f2bd)
Signed-off-by: Ruifeng Zheng <ru...@apache.org>
---
.../ml/classification/LogisticRegression.scala | 231 +++++++++++++++------
.../classification/ProbabilisticClassifier.scala | 189 +++++++++++++----
.../ProbabilisticClassifierSuite.scala | 100 ++++++++-
3 files changed, 419 insertions(+), 101 deletions(-)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
index b602632e68c0..d91261cb304b 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/LogisticRegression.scala
@@ -44,6 +44,7 @@ import org.apache.spark.ml.util.Instrumentation.instrumented
import org.apache.spark.mllib.util.MLUtils
import org.apache.spark.rdd.RDD
import org.apache.spark.sql._
+import org.apache.spark.sql.functions.{get => fget, _}
import org.apache.spark.sql.types.{DataType, StructType}
import org.apache.spark.storage.StorageLevel
import org.apache.spark.util._
@@ -1094,11 +1095,6 @@ class LogisticRegressionModel private[spark] (
throw new SparkException("Multinomial models contain a matrix of coefficients, use " +
"coefficientMatrix instead.")
} else {
- _coefficients
- }
-
- // convert to appropriate vector representation without replicating data
- private lazy val _coefficients: Vector = {
require(coefficientMatrix.isTransposed,
"LogisticRegressionModel coefficients should be row major for binomial model.")
coefficientMatrix match {
@@ -1118,29 +1114,7 @@ class LogisticRegressionModel private[spark] (
throw new SparkException("Multinomial models contain a vector of intercepts, use " +
"interceptVector instead.")
} else {
- _intercept
- }
-
- private val _interceptVector = if (isMultinomial) interceptVector.toDense else null
- private val _intercept = if (!isMultinomial) interceptVector(0) else Double.NaN
- // Array(0.5, 0.0) is the value for default threshold (0.5) and thresholds (unset)
- private var _binaryThresholds: Array[Double] = if (!isMultinomial) Array(0.5, 0.0) else null
-
- private[ml] override def onParamChange(param: Param[_]): Unit = {
- if (!isMultinomial && (param.name == "threshold" || param.name == "thresholds")) {
- if (isDefined(threshold) || isDefined(thresholds)) {
- val _threshold = getThreshold
- if (_threshold == 0.0) {
- _binaryThresholds = Array(_threshold, Double.NegativeInfinity)
- } else if (_threshold == 1.0) {
- _binaryThresholds = Array(_threshold, Double.PositiveInfinity)
- } else {
- _binaryThresholds = Array(_threshold, math.log(_threshold / (1.0 - _threshold)))
- }
- } else {
- _binaryThresholds = null
- }
- }
+ interceptVector(0)
}
@Since("1.5.0")
@@ -1155,22 +1129,117 @@ class LogisticRegressionModel private[spark] (
@Since("1.5.0")
override def getThresholds: Array[Double] = super.getThresholds
- /** Margin (rawPrediction) for class label 1. For binary classification only. */
- private val margin: Vector => Double = (features) => {
- BLAS.dot(features, _coefficients) + _intercept
+ override protected def predictRawColumn(features: Column): Column = if (isMultinomial) {
+ val localCoefficientMatrix = coefficientMatrix
+ val localInterceptVector = interceptVector.toDense
+ udf((features: Vector) => LogisticRegressionModel.predictRaw(
+ features, localCoefficientMatrix, localInterceptVector)).apply(features)
+ } else {
+ val localCoefficients = coefficients
+ val localIntercept = interceptVector(0)
+ udf((features: Vector) => LogisticRegressionModel.predictRaw(
+ features, localCoefficients, localIntercept)).apply(features)
+ }
+
+ override protected def raw2probabilityColumn(rawPrediction: Column): Column = {
+ if (isMultinomial) {
+ udf((rawPrediction: Vector) =>
+ LogisticRegressionModel.raw2probabilityInPlaceMultinomial(rawPrediction)
+ ).apply(rawPrediction)
+ } else {
+ udf((rawPrediction: Vector) =>
+ LogisticRegressionModel.raw2probabilityInPlaceBinary(rawPrediction)
+ ).apply(rawPrediction)
+ }
+ }
+
+ override protected def predictProbabilityColumn(features: Column): Column = {
+ if (isMultinomial) {
+ val localCoefficientMatrix = coefficientMatrix
+ val localInterceptVector = interceptVector.toDense
+ udf((features: Vector) => {
+ val rawPrediction = LogisticRegressionModel.predictRaw(
+ features, localCoefficientMatrix, localInterceptVector)
+ LogisticRegressionModel.raw2probabilityInPlaceMultinomial(rawPrediction)
+ }).apply(features)
+ } else {
+ val localCoefficients = coefficients
+ val localIntercept = interceptVector(0)
+ udf((features: Vector) => {
+ val rawPrediction = LogisticRegressionModel.predictRaw(
+ features, localCoefficients, localIntercept)
+ LogisticRegressionModel.raw2probabilityInPlaceBinary(rawPrediction)
+ }).apply(features)
+ }
}
- /** Margin (rawPrediction) for each class label. */
- private val margins: Vector => Vector = (features) => {
- val m = _interceptVector.copy
- BLAS.gemv(1.0, coefficientMatrix, features, 1.0, m)
- m
+ override protected def raw2predictionColumn(rawPrediction: Column): Column = if (isMultinomial) {
+ if (isDefined(thresholds)) {
+ val localThresholds = getThresholds.clone()
+ udf((rawPrediction: Vector) => {
+ val probability = LogisticRegressionModel.raw2probabilityInPlaceMultinomial(
+ rawPrediction.copy)
+ ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+ }).apply(rawPrediction)
+ } else {
+ udf((rawPrediction: Vector) => rawPrediction.argmax.toDouble).apply(rawPrediction)
+ }
+ } else {
+ val localRawThreshold = LogisticRegressionModel.rawThreshold(getThreshold)
+ val rawScore = fget(unwrap_udt(rawPrediction).getField("values"), lit(1))
+ when(rawScore > localRawThreshold && !isnan(rawScore), 1.0).otherwise(0.0)
}
- /** Score (probability) for class label 1. For binary classification only. */
- private val score: Vector => Double = (features) => {
- val m = margin(features)
- 1.0 / (1.0 + math.exp(-m))
+ override protected def probability2predictionColumn(
+ probability: Column): Column = if (isMultinomial) {
+ if (isDefined(thresholds)) {
+ val localThresholds = getThresholds.clone()
+ probability2predictionColumn(probability, localThresholds)
+ } else {
+ udf((probability: Vector) => probability.argmax.toDouble).apply(probability)
+ }
+ } else {
+ val localProbabilityThreshold = getThreshold
+ val probabilityScore = fget(unwrap_udt(probability).getField("values"), lit(1))
+ when(probabilityScore > localProbabilityThreshold && !isnan(probabilityScore), 1.0)
+ .otherwise(0.0)
+ }
+
+ private def probability2predictionColumn(
+ probability: Column,
+ localThresholds: Array[Double]): Column = {
+ udf((probability: Vector) =>
+ ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+ ).apply(probability)
+ }
+
+ override protected def predictionColumn(features: Column): Column = if (isMultinomial) {
+ val localCoefficientMatrix = coefficientMatrix
+ val localInterceptVector = interceptVector.toDense
+ if (isDefined(thresholds)) {
+ val localThresholds = getThresholds.clone()
+ udf((features: Vector) => {
+ val rawPrediction = LogisticRegressionModel.predictRaw(
+ features, localCoefficientMatrix, localInterceptVector)
+ val probability =
+ LogisticRegressionModel.raw2probabilityInPlaceMultinomial(rawPrediction)
+ ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+ }).apply(features)
+ } else {
+ udf((features: Vector) => {
+ val rawPrediction = LogisticRegressionModel.predictRaw(
+ features, localCoefficientMatrix, localInterceptVector)
+ rawPrediction.argmax.toDouble
+ }).apply(features)
+ }
+ } else {
+ val localCoefficients = coefficients
+ val localIntercept = interceptVector(0)
+ val localRawThreshold = LogisticRegressionModel.rawThreshold(getThreshold)
+ udf((features: Vector) => {
+ val margin = BLAS.dot(features, localCoefficients) + localIntercept
+ if (margin > localRawThreshold) 1.0 else 0.0
+ }).apply(features)
}
@Since("1.6.0")
@@ -1221,34 +1290,26 @@ class LogisticRegressionModel private[spark] (
override def predict(features: Vector): Double = if (isMultinomial) {
super.predict(features)
} else {
- // Note: We should use _threshold instead of $(threshold) since getThreshold is overridden.
- if (score(features) > _binaryThresholds(0)) 1 else 0
+ // Note: We should use getThreshold instead of $(threshold) since getThreshold is overridden.
+ val localRawThreshold = LogisticRegressionModel.rawThreshold(getThreshold)
+ val margin = BLAS.dot(features, coefficients) + interceptVector(0)
+ if (margin > localRawThreshold) 1 else 0
}
override protected def raw2probabilityInPlace(rawPrediction: Vector): Vector = {
- rawPrediction match {
- case dv: DenseVector =>
- val values = dv.values
- if (isMultinomial) {
- Utils.softmax(values)
- } else {
- values(0) = 1.0 / (1.0 + math.exp(-values(0)))
- values(1) = 1.0 - values(0)
- }
- dv
- case sv: SparseVector =>
- throw new RuntimeException("Unexpected error in LogisticRegressionModel:" +
- " raw2probabilitiesInPlace encountered SparseVector")
+ if (isMultinomial) {
+ LogisticRegressionModel.raw2probabilityInPlaceMultinomial(rawPrediction)
+ } else {
+ LogisticRegressionModel.raw2probabilityInPlaceBinary(rawPrediction)
}
}
@Since("3.0.0")
override def predictRaw(features: Vector): Vector = {
if (isMultinomial) {
- margins(features)
+ LogisticRegressionModel.predictRaw(features, coefficientMatrix, interceptVector.toDense)
} else {
- val m = margin(features)
- Vectors.dense(-m, m)
+ LogisticRegressionModel.predictRaw(features, coefficients, interceptVector(0))
}
}
@@ -1274,8 +1335,9 @@ class LogisticRegressionModel private[spark] (
if (isMultinomial) {
super.raw2prediction(rawPrediction)
} else {
- // Note: We should use _threshold instead of $(threshold) since getThreshold is overridden.
- if (rawPrediction(1) > _binaryThresholds(1)) 1.0 else 0.0
+ // Note: We should use getThreshold instead of $(threshold) since getThreshold is overridden.
+ val localRawThreshold = LogisticRegressionModel.rawThreshold(getThreshold)
+ if (rawPrediction(1) > localRawThreshold) 1.0 else 0.0
}
}
@@ -1283,8 +1345,8 @@ class LogisticRegressionModel private[spark] (
if (isMultinomial) {
super.probability2prediction(probability)
} else {
- // Note: We should use _threshold instead of $(threshold) since getThreshold is overridden.
- if (probability(1) > _binaryThresholds(0)) 1.0 else 0.0
+ // Note: We should use getThreshold instead of $(threshold) since getThreshold is overridden.
+ if (probability(1) > getThreshold) 1.0 else 0.0
}
}
@@ -1354,6 +1416,51 @@ class LogisticRegressionModel private[spark] (
@Since("1.6.0")
object LogisticRegressionModel extends MLReadable[LogisticRegressionModel] {
+ /**
+ * Converts a binary probability threshold to the equivalent raw margin threshold.
+ * Probability thresholds 0 and 1 map to negative and positive infinity, respectively,
+ * where the log-odds formula is undefined.
+ */
+ private def rawThreshold(threshold: Double): Double = {
+ if (threshold == 0.0) {
+ Double.NegativeInfinity
+ } else if (threshold == 1.0) {
+ Double.PositiveInfinity
+ } else {
+ math.log(threshold / (1.0 - threshold))
+ }
+ }
+
+ private def predictRaw(
+ features: Vector,
+ coefficients: Vector,
+ intercept: Double): Vector = {
+ val margin = BLAS.dot(features, coefficients) + intercept
+ Vectors.dense(-margin, margin)
+ }
+
+ private def predictRaw(
+ features: Vector,
+ coefficientMatrix: Matrix,
+ interceptVector: DenseVector): Vector = {
+ val margins = interceptVector.copy
+ BLAS.gemv(1.0, coefficientMatrix, features, 1.0, margins)
+ margins
+ }
+
+ private def raw2probabilityInPlaceBinary(rawPrediction: Vector): Vector = {
+ val values = rawPrediction.toArray
+ values(0) = 1.0 / (1.0 + math.exp(-values(0)))
+ values(1) = 1.0 - values(0)
+ Vectors.dense(values)
+ }
+
+ private def raw2probabilityInPlaceMultinomial(rawPrediction: Vector): Vector = {
+ val values = rawPrediction.toArray
+ Utils.softmax(values)
+ Vectors.dense(values)
+ }
+
private[ml] case class Data(
numClasses: Int,
numFeatures: Int,
diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/ProbabilisticClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/ProbabilisticClassifier.scala
index a6648f4fab7b..2e3e681df939 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/ProbabilisticClassifier.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/ProbabilisticClassifier.scala
@@ -23,7 +23,7 @@ import org.apache.spark.ml.linalg.{DenseVector, SQLDataTypes, Vector}
import org.apache.spark.ml.param.ParamMap
import org.apache.spark.ml.param.shared._
import org.apache.spark.ml.util.{Identifiable, 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, StructType}
@@ -67,6 +67,68 @@ abstract class ProbabilisticClassifier[
* Model produced by a [[ProbabilisticClassifier]].
* Classes are indexed {0, 1, ..., numClasses - 1}.
*
+ * `transform` selects column-producing methods based on which output columns are set:
+ *
+ * <table>
+ * <caption>Column-producing methods by requested output columns</caption>
+ * <tr>
+ * <th>Raw prediction</th>
+ * <th>Probability</th>
+ * <th>Prediction</th>
+ * <th>Column-producing methods</th>
+ * </tr>
+ * <tr>
+ * <td>Not set</td>
+ * <td>Not set</td>
+ * <td>Not set</td>
+ * <td>None</td>
+ * </tr>
+ * <tr>
+ * <td>Set</td>
+ * <td>Not set</td>
+ * <td>Not set</td>
+ * <td><code>predictRawColumn</code></td>
+ * </tr>
+ * <tr>
+ * <td>Not set</td>
+ * <td>Set</td>
+ * <td>Not set</td>
+ * <td><code>predictProbabilityColumn</code></td>
+ * </tr>
+ * <tr>
+ * <td>Not set</td>
+ * <td>Not set</td>
+ * <td>Set</td>
+ * <td><code>predictionColumn</code></td>
+ * </tr>
+ * <tr>
+ * <td>Set</td>
+ * <td>Set</td>
+ * <td>Not set</td>
+ * <td><code>predictRawColumn</code> => <code>raw2probabilityColumn</code></td>
+ * </tr>
+ * <tr>
+ * <td>Set</td>
+ * <td>Not set</td>
+ * <td>Set</td>
+ * <td><code>predictRawColumn</code> => <code>raw2predictionColumn</code></td>
+ * </tr>
+ * <tr>
+ * <td>Not set</td>
+ * <td>Set</td>
+ * <td>Set</td>
+ * <td><code>predictProbabilityColumn</code> =>
+ * <code>probability2predictionColumn</code></td>
+ * </tr>
+ * <tr>
+ * <td>Set</td>
+ * <td>Set</td>
+ * <td>Set</td>
+ * <td><code>predictRawColumn</code> => <code>raw2probabilityColumn</code> =>
+ * <code>raw2predictionColumn</code></td>
+ * </tr>
+ * </table>
+ *
* @tparam FeaturesType Type of input features. E.g., `Vector`
* @tparam M Concrete Model type
*/
@@ -75,6 +137,66 @@ abstract class ProbabilisticClassificationModel[
M <: ProbabilisticClassificationModel[FeaturesType, M]]
extends ClassificationModel[FeaturesType, M] with ProbabilisticClassifierParams {
+ /**
+ * Returns an expression that converts a raw-prediction vector column to a probability vector.
+ *
+ * `transform` uses this when both [[rawPredictionCol]] and
+ * [[probabilityCol]] are set. It consumes the output of [[predictRawColumn]].
+ *
+ * @param rawPrediction input raw-prediction column
+ * @return probability column of type `Vector`
+ * @note Thresholds do not affect this conversion; they are applied only when producing a
+ * prediction.
+ */
+ protected def raw2probabilityColumn(rawPrediction: Column): Column = {
+ udf(raw2probability _).apply(rawPrediction)
+ }
+
+ /**
+ * Returns an expression that produces a probability vector directly from a features column.
+ *
+ * `transform` uses this when [[probabilityCol]] is set and
+ * [[rawPredictionCol]] is not set. When [[predictionCol]] is also set, it is used
+ * together with [[probability2predictionColumn]].
+ *
+ * @param features input features column
+ * @return probability column of type `Vector`
+ */
+ protected def predictProbabilityColumn(features: Column): Column = {
+ udf { value: Any => predictProbability(value.asInstanceOf[FeaturesType]) }.apply(features)
+ }
+
+ /**
+ * Returns an expression that produces a predicted label from a raw-prediction vector column.
+ *
+ * `transform` uses this when both [[rawPredictionCol]] and
+ * [[predictionCol]] are set. It consumes the output of [[predictRawColumn]].
+ *
+ * @param rawPrediction input raw-prediction column
+ * @return prediction column of type `Double`
+ * @note Unlike the default classification behavior, probabilistic classification honors
+ * [[thresholds]] by converting raw predictions to probabilities before selecting a
+ * prediction.
+ */
+ override protected def raw2predictionColumn(rawPrediction: Column): Column = {
+ udf(raw2prediction _).apply(rawPrediction)
+ }
+
+ /**
+ * Returns an expression that produces a predicted label from a probability vector column.
+ *
+ * `transform` uses this when [[predictionCol]] and
+ * [[probabilityCol]] are set and [[rawPredictionCol]] is not set. It consumes the output of
+ * [[predictProbabilityColumn]].
+ *
+ * @param probability input probability column
+ * @return prediction column of type `Double`
+ * @note This method honors [[thresholds]] when they are set.
+ */
+ protected def probability2predictionColumn(probability: Column): Column = {
+ udf(probability2prediction _).apply(probability)
+ }
+
/** @group setParam */
def setProbabilityCol(value: String): M = set(probabilityCol, value).asInstanceOf[M]
@@ -118,21 +240,16 @@ abstract class ProbabilisticClassificationModel[
var outputData = dataset
var numColsOutput = 0
if ($(rawPredictionCol).nonEmpty) {
- val predictRawUDF = udf { features: Any =>
- predictRaw(features.asInstanceOf[FeaturesType])
- }
- outputData = outputData.withColumn(getRawPredictionCol, predictRawUDF(col(getFeaturesCol)),
+ outputData = outputData.withColumn(getRawPredictionCol,
+ predictRawColumn(col(getFeaturesCol)),
outputSchema($(rawPredictionCol)).metadata)
numColsOutput += 1
}
if ($(probabilityCol).nonEmpty) {
val probCol = if ($(rawPredictionCol).nonEmpty) {
- udf(raw2probability _).apply(col($(rawPredictionCol)))
+ raw2probabilityColumn(col($(rawPredictionCol)))
} else {
- val probabilityUDF = udf { features: Any =>
- predictProbability(features.asInstanceOf[FeaturesType])
- }
- probabilityUDF(col($(featuresCol)))
+ predictProbabilityColumn(col($(featuresCol)))
}
outputData = outputData.withColumn($(probabilityCol), probCol,
outputSchema($(probabilityCol)).metadata)
@@ -140,14 +257,11 @@ abstract class ProbabilisticClassificationModel[
}
if ($(predictionCol).nonEmpty) {
val predCol = if ($(rawPredictionCol).nonEmpty) {
- udf(raw2prediction _).apply(col($(rawPredictionCol)))
+ raw2predictionColumn(col($(rawPredictionCol)))
} else if ($(probabilityCol).nonEmpty) {
- udf(probability2prediction _).apply(col($(probabilityCol)))
+ probability2predictionColumn(col($(probabilityCol)))
} else {
- val predictUDF = udf { features: Any =>
- predict(features.asInstanceOf[FeaturesType])
- }
- predictUDF(col($(featuresCol)))
+ predictionColumn(col($(featuresCol)))
}
outputData = outputData.withColumn($(predictionCol), predCol,
outputSchema($(predictionCol)).metadata)
@@ -208,27 +322,10 @@ abstract class ProbabilisticClassificationModel[
* @return predicted label
*/
protected def probability2prediction(probability: Vector): Double = {
- if (!isDefined(thresholds)) {
- probability.argmax
+ if (isDefined(thresholds)) {
+ ProbabilisticClassificationModel.probability2prediction(probability, getThresholds)
} else {
- val thresholds = getThresholds
- var argMax = 0
- var max = Double.NegativeInfinity
- var i = 0
- val probabilitySize = probability.size
- while (i < probabilitySize) {
- // Thresholds are all > 0, excepting that at most one may be 0.
- // The single class whose threshold is 0, if any, will always be predicted
- // ('scaled' = +Infinity). However in the case that this class also has
- // 0 probability, the class will not be selected ('scaled' is NaN).
- val scaled = probability(i) / thresholds(i)
- if (scaled > max) {
- max = scaled
- argMax = i
- }
- i += 1
- }
- argMax
+ probability.argmax
}
}
@@ -256,6 +353,26 @@ abstract class ProbabilisticClassificationModel[
private[ml] object ProbabilisticClassificationModel {
+ def probability2prediction(probability: Vector, thresholds: Array[Double]): Double = {
+ var argMax = 0
+ var max = Double.NegativeInfinity
+ var i = 0
+ val probabilitySize = probability.size
+ while (i < probabilitySize) {
+ // Thresholds are all > 0, excepting that at most one may be 0.
+ // The single class whose threshold is 0, if any, will always be predicted
+ // ('scaled' = +Infinity). However in the case that this class also has
+ // 0 probability, the class will not be selected ('scaled' is NaN).
+ val scaled = probability(i) / thresholds(i)
+ if (scaled > max) {
+ max = scaled
+ argMax = i
+ }
+ i += 1
+ }
+ argMax
+ }
+
/**
* Normalize a vector of raw predictions to be a multinomial probability vector, in place.
*
diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/ProbabilisticClassifierSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/ProbabilisticClassifierSuite.scala
index dbdc6a27bc3b..bb2a0b82d241 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/classification/ProbabilisticClassifierSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/classification/ProbabilisticClassifierSuite.scala
@@ -22,9 +22,11 @@ import org.scalatest.Assertions._
import org.apache.spark.SparkFunSuite
import org.apache.spark.ml.linalg.{Vector, Vectors}
import org.apache.spark.ml.param.ParamMap
-import org.apache.spark.ml.util.MLTest
+import org.apache.spark.ml.util.{Identifiable, MLTest}
import org.apache.spark.ml.util.TestingUtils._
-import org.apache.spark.sql.{Dataset, Row}
+import org.apache.spark.mllib.util.MLlibTestSparkContext
+import org.apache.spark.sql.{Column, Dataset, Row}
+import org.apache.spark.sql.functions.udf
final class TestProbabilisticClassificationModel(
override val uid: String,
@@ -47,8 +49,71 @@ final class TestProbabilisticClassificationModel(
}
}
+final class MockProbabilisticClassificationModelWithColumnFunctions(override val uid: String)
+ extends ProbabilisticClassificationModel[
+ Vector, MockProbabilisticClassificationModelWithColumnFunctions] {
-class ProbabilisticClassifierSuite extends SparkFunSuite {
+ def this() = this(Identifiable.randomUID("mockprobabilisticmodelwithcolumnfunctions"))
+
+ private def vectorColumn(input: Column): Column = udf((value: Vector) => value).apply(input)
+
+ private def predictionFromVector(input: Column): Column = {
+ udf((value: Vector) => value.argmax.toDouble).apply(input)
+ }
+
+ override protected def predictRawColumn(features: Column): Column = vectorColumn(features)
+
+ override protected def raw2probabilityColumn(rawPrediction: Column): Column = {
+ vectorColumn(rawPrediction)
+ }
+
+ override protected def predictProbabilityColumn(features: Column): Column = vectorColumn(features)
+
+ override protected def raw2predictionColumn(rawPrediction: Column): Column = {
+ predictionFromVector(rawPrediction)
+ }
+
+ override protected def probability2predictionColumn(probability: Column): Column = {
+ predictionFromVector(probability)
+ }
+
+ override protected def predictionColumn(features: Column): Column = predictionFromVector(features)
+
+ override def predict(features: Vector): Double = throw new UnsupportedOperationException()
+
+ override def predictRaw(features: Vector): Vector = throw new UnsupportedOperationException()
+
+ override protected def raw2probabilityInPlace(rawPrediction: Vector): Vector = {
+ throw new UnsupportedOperationException()
+ }
+
+ override protected def raw2probability(rawPrediction: Vector): Vector = {
+ throw new UnsupportedOperationException()
+ }
+
+ override def predictProbability(features: Vector): Vector = {
+ throw new UnsupportedOperationException()
+ }
+
+ override protected def raw2prediction(rawPrediction: Vector): Double = {
+ throw new UnsupportedOperationException()
+ }
+
+ override protected def probability2prediction(probability: Vector): Double = {
+ throw new UnsupportedOperationException()
+ }
+
+ override def copy(extra: ParamMap): MockProbabilisticClassificationModelWithColumnFunctions = {
+ throw new UnsupportedOperationException()
+ }
+
+ override def numClasses: Int = 2
+}
+
+
+class ProbabilisticClassifierSuite extends SparkFunSuite with MLlibTestSparkContext {
+
+ import testImplicits._
test("test thresholding") {
val testModel = new TestProbabilisticClassificationModel("myuid", 2, 2)
@@ -101,6 +166,35 @@ class ProbabilisticClassifierSuite extends SparkFunSuite {
ProbabilisticClassificationModel.normalizeToProbabilitiesInPlace(vec3)
}
}
+
+ test("transform uses probabilistic column prediction hooks") {
+ val expected = Vectors.dense(0.25, 0.75)
+ val data = Seq(Tuple1(expected)).toDF("features")
+
+ val allOutputs = new MockProbabilisticClassificationModelWithColumnFunctions()
+ .transform(data)
+ .select("rawPrediction", "probability", "prediction")
+ .head()
+ assert(allOutputs.getAs[Vector](0) === expected)
+ assert(allOutputs.getAs[Vector](1) === expected)
+ assert(allOutputs.getDouble(2) === 1.0)
+
+ val probabilityAndPrediction = new MockProbabilisticClassificationModelWithColumnFunctions()
+ .setRawPredictionCol("")
+ .transform(data)
+ .select("probability", "prediction")
+ .head()
+ assert(probabilityAndPrediction.getAs[Vector](0) === expected)
+ assert(probabilityAndPrediction.getDouble(1) === 1.0)
+
+ val predictionOnly = new MockProbabilisticClassificationModelWithColumnFunctions()
+ .setRawPredictionCol("")
+ .setProbabilityCol("")
+ .transform(data)
+ .select("prediction")
+ .head()
+ assert(predictionOnly.getDouble(0) === 1.0)
+ }
}
object ProbabilisticClassifierSuite {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org