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 2d2d457ee9e6 [SPARK-59052][ML] Optimize ClassificationModel by
reducing closure size and accelerating transform
2d2d457ee9e6 is described below
commit 2d2d457ee9e6514c2fad5289e1318808fae23d44
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Fri Aug 28 16:11:39 2026 +0800
[SPARK-59052][ML] Optimize ClassificationModel by reducing closure size and
accelerating transform
### What changes were proposed in this pull request?
This PR adds protected column-expression hooks to `ClassificationModel` for
producing raw
predictions, converting raw predictions to labels, and predicting labels
directly from features.
`ClassificationModel.transform` uses these hooks while retaining the
existing UDF-based defaults.
`LinearSVCModel` overrides the hooks to snapshot only coefficients,
intercept, and threshold in
its UDF closures. Raw-prediction-to-label conversion uses Spark column
expressions, and the
model-capturing `margin` function field is removed.
The corresponding `ProbabilisticClassificationModel` integration will be
handled separately in
#58336.
### Why are the changes needed?
The previous transform implementation built UDFs from bound model methods.
This serialized the
entire model into each closure and performed parameter lookup for every
input row. The hooks let
classification models provide compact closures or native Spark expressions
without changing
prediction behavior.
An in-process comparison used 4,096-feature models for closure
serialization and 3,000,000 cached
rows with 32 features for 20 alternating execution measurements:
| Transform path | Previous closure | New closure | Previous median | New
median | Speedup |
|---|---:|---:|---:|---:|---:|
| Raw prediction | 38,268 B | 34,627 B (-9.5%) | 553.5 ms | 527.1 ms
(-4.8%) | 1.05x |
| Prediction only | 38,206 B | 34,580 B (-9.5%) | 426.3 ms | 229.2 ms
(-46.2%) | 1.86x |
| Raw + prediction | 76,552 B (2 UDFs) | 34,627 B (1 UDF) (-54.8%) | 818.8
ms | 586.7 ms (-28.3%) | 1.40x |
Across these paths, the new approach is up to 1.86x faster.
The temporary benchmark probe was removed after measurement.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
- `build/sbt mllib/scalastyle`
- `build/sbt 'mllib/testOnly
org.apache.spark.ml.classification.ClassifierSuite'`
(2 tests passed)
- `build/sbt 'mllib/testOnly
org.apache.spark.ml.classification.LinearSVCSuite'`
(16 tests passed; 1 pre-existing test ignored)
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #58341 from zhengruifeng/ml-spark-59052-linearsvc-closure-dev-4.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../spark/ml/classification/Classifier.scala | 55 ++++++++++++++++------
.../apache/spark/ml/classification/LinearSVC.scala | 34 +++++++++++--
.../spark/ml/classification/ClassifierSuite.scala | 53 +++++++++++++++++++--
3 files changed, 120 insertions(+), 22 deletions(-)
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 2772bdb8c3f6..46fdb93ad9bd 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
@@ -24,7 +24,7 @@ import org.apache.spark.ml.linalg.{SQLDataTypes, Vector}
import org.apache.spark.ml.param.ParamMap
import org.apache.spark.ml.param.shared.HasRawPredictionCol
import org.apache.spark.ml.util._
-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}
@@ -92,6 +92,38 @@ abstract class Classifier[
abstract class ClassificationModel[FeaturesType, M <:
ClassificationModel[FeaturesType, M]]
extends PredictionModel[FeaturesType, M] with ClassifierParams {
+ /**
+ * Returns an expression that produces a raw-prediction vector from a
features column.
+ * The default wraps [[predictRaw]] in a UDF. Models may override this with
a native expression or
+ * a UDF that snapshots prediction state.
+ *
+ * @param features input features column
+ * @return raw-prediction column of type `Vector`
+ */
+ protected def predictRawColumn(features: Column): Column = {
+ udf { value: Any => predictRaw(value.asInstanceOf[FeaturesType])
}.apply(features)
+ }
+
+ /**
+ * Returns an expression that produces a predicted label from a
raw-prediction vector column.
+ *
+ * @param rawPrediction input raw-prediction column
+ * @return prediction column of type `Double`
+ */
+ protected def raw2predictionColumn(rawPrediction: Column): Column = {
+ udf(raw2prediction _).apply(rawPrediction)
+ }
+
+ /**
+ * Returns an expression that produces a predicted label directly from a
features column.
+ *
+ * @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]
@@ -127,24 +159,19 @@ abstract class ClassificationModel[FeaturesType, M <:
ClassificationModel[Featur
// This is a bit complicated since it tries to avoid repeated computation.
var outputData = dataset
var numColsOutput = 0
- if (getRawPredictionCol != "") {
- val predictRawUDF = udf { features: Any =>
- predictRaw(features.asInstanceOf[FeaturesType])
- }
- outputData = outputData.withColumn(getRawPredictionCol,
predictRawUDF(col(getFeaturesCol)),
+ if ($(rawPredictionCol).nonEmpty) {
+ outputData = outputData.withColumn($(rawPredictionCol),
+ predictRawColumn(col($(featuresCol))),
outputSchema($(rawPredictionCol)).metadata)
numColsOutput += 1
}
- if (getPredictionCol != "") {
- val predCol = if (getRawPredictionCol != "") {
- udf(raw2prediction _).apply(col(getRawPredictionCol))
+ if ($(predictionCol).nonEmpty) {
+ val predCol = if ($(rawPredictionCol).nonEmpty) {
+ raw2predictionColumn(col($(rawPredictionCol)))
} else {
- val predictUDF = udf { features: Any =>
- predict(features.asInstanceOf[FeaturesType])
- }
- predictUDF(col(getFeaturesCol))
+ predictionColumn(col($(featuresCol)))
}
- outputData = outputData.withColumn(getPredictionCol, predCol,
+ outputData = outputData.withColumn($(predictionCol), predCol,
outputSchema($(predictionCol)).metadata)
numColsOutput += 1
}
diff --git
a/mllib/src/main/scala/org/apache/spark/ml/classification/LinearSVC.scala
b/mllib/src/main/scala/org/apache/spark/ml/classification/LinearSVC.scala
index 89f1f7792c07..d9b335325643 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/LinearSVC.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/LinearSVC.scala
@@ -29,6 +29,7 @@ import org.apache.spark.SparkException
import org.apache.spark.annotation.Since
import org.apache.spark.internal.Logging
import org.apache.spark.internal.LogKeys.{COUNT, RANGE}
+import org.apache.spark.ml.{functions => MLFunctions}
import org.apache.spark.ml.feature._
import org.apache.spark.ml.linalg._
import org.apache.spark.ml.optim.aggregator._
@@ -42,6 +43,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._
import org.apache.spark.storage.StorageLevel
/** Params for linear SVM Classifier. */
@@ -377,8 +379,29 @@ class LinearSVCModel private[classification] (
@Since("2.2.0")
def setThreshold(value: Double): this.type = set(threshold, value)
- private val margin: Vector => Double = (features) => {
- BLAS.dot(features, coefficients) + intercept
+ override protected def predictRawColumn(features: Column): Column = {
+ val localCoefficients = coefficients
+ val localIntercept = intercept
+ udf((features: Vector) => {
+ val margin = BLAS.dot(features, localCoefficients) + localIntercept
+ Vectors.dense(-margin, margin)
+ }).apply(features)
+ }
+
+ override protected def raw2predictionColumn(rawPrediction: Column): Column =
{
+ val localThreshold = getThreshold
+ val rawScore = MLFunctions.vector_get(rawPrediction, lit(1))
+ when(!isnan(rawScore) && rawScore > localThreshold, 1.0).otherwise(0.0)
+ }
+
+ override protected def predictionColumn(features: Column): Column = {
+ val localCoefficients = coefficients
+ val localIntercept = intercept
+ val localThreshold = getThreshold
+ udf((features: Vector) => {
+ val margin = BLAS.dot(features, localCoefficients) + localIntercept
+ if (margin > localThreshold) 1.0 else 0.0
+ }).apply(features)
}
/**
@@ -403,13 +426,14 @@ class LinearSVCModel private[classification] (
}
override def predict(features: Vector): Double = {
- if (margin(features) > $(threshold)) 1.0 else 0.0
+ val margin = BLAS.dot(features, coefficients) + intercept
+ if (margin > $(threshold)) 1.0 else 0.0
}
@Since("3.0.0")
override def predictRaw(features: Vector): Vector = {
- val m = margin(features)
- Vectors.dense(-m, m)
+ val margin = BLAS.dot(features, coefficients) + intercept
+ Vectors.dense(-margin, margin)
}
override protected def raw2prediction(rawPrediction: Vector): Double = {
diff --git
a/mllib/src/test/scala/org/apache/spark/ml/classification/ClassifierSuite.scala
b/mllib/src/test/scala/org/apache/spark/ml/classification/ClassifierSuite.scala
index 57cd99ecced1..cb5fce195627 100644
---
a/mllib/src/test/scala/org/apache/spark/ml/classification/ClassifierSuite.scala
+++
b/mllib/src/test/scala/org/apache/spark/ml/classification/ClassifierSuite.scala
@@ -18,13 +18,15 @@
package org.apache.spark.ml.classification
import org.apache.spark.{SparkException, SparkFunSuite}
-import org.apache.spark.ml.classification.ClassifierSuite.MockClassifier
+import org.apache.spark.ml.classification.ClassifierSuite.{
+ MockClassificationModelWithColumnFunctions, MockClassifier}
import org.apache.spark.ml.feature.LabeledPoint
import org.apache.spark.ml.linalg.{Vector, Vectors}
import org.apache.spark.ml.param.ParamMap
import org.apache.spark.ml.util._
import org.apache.spark.mllib.util.MLlibTestSparkContext
-import org.apache.spark.sql.{DataFrame, Dataset}
+import org.apache.spark.sql.{Column, DataFrame, Dataset}
+import org.apache.spark.sql.functions.udf
class ClassifierSuite extends SparkFunSuite with MLlibTestSparkContext {
@@ -62,6 +64,20 @@ class ClassifierSuite extends SparkFunSuite with
MLlibTestSparkContext {
assert(e.getMessage == "ML algorithm was given empty dataset.")
}
}
+
+ test("transform uses column prediction hooks") {
+ val expectedRawPrediction = Vectors.dense(-1.0, 1.0)
+ val data = Seq(Tuple1(expectedRawPrediction)).toDF("features")
+
+ val model = new MockClassificationModelWithColumnFunctions
+ val bothOutputs = model.transform(data).select("rawPrediction",
"prediction").head()
+ assert(bothOutputs.getAs[Vector](0) === expectedRawPrediction)
+ assert(bothOutputs.getDouble(1) === 1.0)
+
+ val predictionOnly = new
MockClassificationModelWithColumnFunctions().setRawPredictionCol("")
+ .transform(data).select("prediction").head()
+ assert(predictionOnly.getDouble(0) === 1.0)
+ }
}
object ClassifierSuite {
@@ -95,7 +111,7 @@ object ClassifierSuite {
def this() = this(Identifiable.randomUID("mockclassificationmodel"))
- def predictRaw(features: Vector): Vector = throw new
UnsupportedOperationException()
+ override def predictRaw(features: Vector): Vector = throw new
UnsupportedOperationException()
override def copy(extra: ParamMap): MockClassificationModel =
throw new UnsupportedOperationException()
@@ -103,4 +119,35 @@ object ClassifierSuite {
override def numClasses: Int = throw new UnsupportedOperationException()
}
+ class MockClassificationModelWithColumnFunctions(override val uid: String)
+ extends ClassificationModel[Vector,
MockClassificationModelWithColumnFunctions] {
+
+ def this() =
this(Identifiable.randomUID("mockclassificationmodelwithcolumnfunctions"))
+
+ override protected def predictRawColumn(features: Column): Column = {
+ udf((value: Vector) => value).apply(features)
+ }
+
+ override protected def raw2predictionColumn(rawPrediction: Column): Column
= {
+ udf((value: Vector) => value.argmax.toDouble).apply(rawPrediction)
+ }
+
+ override protected def predictionColumn(features: Column): Column = {
+ udf((value: Vector) => value.argmax.toDouble).apply(features)
+ }
+
+ override def predict(features: Vector): Double = throw new
UnsupportedOperationException()
+
+ override def predictRaw(features: Vector): Vector = throw new
UnsupportedOperationException()
+
+ override protected def raw2prediction(rawPrediction: Vector): Double = {
+ throw new UnsupportedOperationException()
+ }
+
+ override def copy(extra: ParamMap):
MockClassificationModelWithColumnFunctions =
+ throw new UnsupportedOperationException()
+
+ override def numClasses: Int = 2
+ }
+
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
