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 03:30:07 UTC
(spark) branch master updated: [SPARK-59286][ML] Reduce NaiveBayesModel transform closure size
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 6fdf20178f13 [SPARK-59286][ML] Reduce NaiveBayesModel transform closure size
6fdf20178f13 is described below
commit 6fdf20178f1345d2159d78e7a364e3314473c27b
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Wed Sep 9 11:29:53 2026 +0800
[SPARK-59286][ML] Reduce NaiveBayesModel transform closure size
### What changes were proposed in this pull request?
This PR reduces the transform closure size of `NaiveBayesModel` for multinomial, complement,
Bernoulli, and Gaussian models.
The model keeps its existing transient lazy prediction state for single-instance prediction. The
column-expression prediction hooks build a local prediction function from the fitted parameters,
constructing Bernoulli or Gaussian auxiliary state once for the selected UDF. Optional thresholds
are also snapshotted, so the UDFs do not retain the complete model and its parameter graph.
Companion-object helpers perform the state construction and stateless prediction calculations.
This retries #58557 with a smaller implementation after its revert in #58610.
### Why are the changes needed?
The default probabilistic classifier hooks invoke bound model methods. Their UDF closures therefore
retain the complete `NaiveBayesModel`, even though scoring only needs its fitted vectors, matrices,
model-specific derived state, and optional thresholds. This adds avoidable driver memory pressure
for long-lived Spark Connect servers.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
The following checks passed:
```
build/sbt mllib/compile
build/sbt 'mllib/testOnly org.apache.spark.ml.classification.NaiveBayesSuite'
```
`NaiveBayesSuite` ran 17 tests covering all four model types.
A temporary local probe used a deterministic three-class model with 1,024 features for each model
type. It extracted each `ScalaUDF.function` from the analyzed transform plan and serialized it with
Spark's closure serializer. The "before" implementation recreated the bound-model hooks. "All"
serializes the raw-prediction, probability, and prediction functions together.
| Model/output | Before | After | Reduction |
|---|---:|---:|---:|
| Multinomial raw prediction | 29,957 B | 26,853 B | 10.4% |
| Multinomial probability | 29,964 B | 26,861 B | 10.4% |
| Multinomial prediction | 29,894 B | 26,897 B | 10.0% |
| Multinomial all | 30,819 B | 27,496 B | 10.8% |
| Complement raw prediction | 29,947 B | 26,715 B | 10.8% |
| Complement probability | 29,954 B | 26,723 B | 10.8% |
| Complement prediction | 29,884 B | 26,759 B | 10.5% |
| Complement all | 30,809 B | 27,358 B | 11.2% |
| Bernoulli raw prediction | 29,937 B | 26,858 B | 10.3% |
| Bernoulli probability | 29,944 B | 26,866 B | 10.3% |
| Bernoulli prediction | 29,874 B | 26,902 B | 9.9% |
| Bernoulli all | 30,799 B | 27,501 B | 10.7% |
| Gaussian raw prediction | 54,503 B | 51,525 B | 5.5% |
| Gaussian probability | 54,510 B | 51,533 B | 5.5% |
| Gaussian prediction | 54,440 B | 51,569 B | 5.3% |
| Gaussian all | 55,365 B | 52,168 B | 5.8% |
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (GPT-5)
Closes #58611 from zhengruifeng/ml-spark-59286-naive-bayes-udfs-dev-7.
Authored-by: Ruifeng Zheng <ru...@apache.org>
Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
---
.../spark/ml/classification/NaiveBayes.scala | 232 +++++++++++++++------
1 file changed, 165 insertions(+), 67 deletions(-)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/NaiveBayes.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/NaiveBayes.scala
index aaa1ef4a4764..bc297909adaa 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/NaiveBayes.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/NaiveBayes.scala
@@ -442,7 +442,7 @@ class NaiveBayesModel private[ml] (
*/
@transient private lazy val thetaMinusNegTheta = $(modelType) match {
case Bernoulli =>
- theta.map(value => value - math.log1p(-math.exp(value)))
+ NaiveBayesModel.bernoulliThetaMinusNegTheta(theta)
case _ =>
// This should never happen.
throw new IllegalArgumentException(s"Invalid modelType: ${$(modelType)}. " +
@@ -451,11 +451,7 @@ class NaiveBayesModel private[ml] (
@transient private lazy val piMinusThetaSum = $(modelType) match {
case Bernoulli =>
- val negTheta = theta.map(value => math.log1p(-math.exp(value)))
- val ones = new DenseVector(Array.fill(theta.numCols)(1.0))
- val piMinusThetaSum = pi.toDense.copy
- BLAS.gemv(1.0, negTheta, ones, 1.0, piMinusThetaSum)
- piMinusThetaSum
+ NaiveBayesModel.bernoulliPiMinusThetaSum(pi, theta)
case _ =>
// This should never happen.
throw new IllegalArgumentException(s"Invalid modelType: ${$(modelType)}. " +
@@ -469,11 +465,7 @@ class NaiveBayesModel private[ml] (
*/
@transient private lazy val logVarSum = $(modelType) match {
case Gaussian =>
- Array.tabulate(numClasses) { i =>
- Iterator.range(0, numFeatures).map { j =>
- math.log(sigma(i, j))
- }.sum
- }
+ NaiveBayesModel.gaussianLogVarSum(sigma)
case _ =>
// This should never happen.
throw new IllegalArgumentException(s"Invalid modelType: ${$(modelType)}. " +
@@ -486,14 +478,139 @@ class NaiveBayesModel private[ml] (
@Since("1.5.0")
override val numClasses: Int = pi.size
- private def multinomialCalculation(features: Vector) = {
+ override protected def predictRawColumn(features: Column): Column = {
+ val localPredictRaw = NaiveBayesModel.predictRawFunction(
+ $(modelType), pi, theta, sigma)
+ udf((features: Vector) => localPredictRaw(features)).apply(features)
+ }
+
+ override protected def raw2probabilityColumn(rawPrediction: Column): Column = {
+ udf((rawPrediction: Vector) =>
+ NaiveBayesModel.raw2probabilityInPlace(rawPrediction.copy)
+ ).apply(rawPrediction)
+ }
+
+ override protected def predictProbabilityColumn(features: Column): Column = {
+ val localPredictRaw = NaiveBayesModel.predictRawFunction(
+ $(modelType), pi, theta, sigma)
+ udf((features: Vector) => {
+ val rawPrediction = localPredictRaw(features)
+ NaiveBayesModel.raw2probabilityInPlace(rawPrediction)
+ }).apply(features)
+ }
+
+ override protected def raw2predictionColumn(rawPrediction: Column): Column = {
+ if (isDefined(thresholds)) {
+ val localThresholds = getThresholds.clone()
+ udf((rawPrediction: Vector) => {
+ val probability = NaiveBayesModel.raw2probabilityInPlace(rawPrediction.copy)
+ ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+ }).apply(rawPrediction)
+ } else {
+ udf((rawPrediction: Vector) => rawPrediction.argmax.toDouble).apply(rawPrediction)
+ }
+ }
+
+ override protected def predictionColumn(features: Column): Column = {
+ val localPredictRaw = NaiveBayesModel.predictRawFunction(
+ $(modelType), pi, theta, sigma)
+ if (isDefined(thresholds)) {
+ val localThresholds = getThresholds.clone()
+ udf((features: Vector) => {
+ val rawPrediction = localPredictRaw(features)
+ val probability = NaiveBayesModel.raw2probabilityInPlace(rawPrediction)
+ ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+ }).apply(features)
+ } else {
+ udf((features: Vector) => localPredictRaw(features).argmax.toDouble).apply(features)
+ }
+ }
+
+ @Since("3.0.0")
+ override def predictRaw(features: Vector): Vector = {
+ $(modelType) match {
+ case Multinomial =>
+ NaiveBayesModel.multinomialCalculation(features, pi, theta)
+ case Complement =>
+ NaiveBayesModel.complementCalculation(features, theta)
+ case Bernoulli =>
+ NaiveBayesModel.bernoulliCalculation(
+ features, piMinusThetaSum, thetaMinusNegTheta)
+ case Gaussian =>
+ NaiveBayesModel.gaussianCalculation(features, pi, theta, sigma, logVarSum)
+ }
+ }
+
+ override protected def raw2probabilityInPlace(rawPrediction: Vector): Vector = {
+ NaiveBayesModel.raw2probabilityInPlace(rawPrediction)
+ }
+
+ private[spark] override def estimatedSize: Long = {
+ var size = this.estimateMatadataSize
+ if (this.pi != null) {
+ size += this.pi.getSizeInBytes
+ }
+ if (this.theta != null) {
+ size += this.theta.getSizeInBytes
+ }
+ if (this.sigma != null) {
+ size += this.sigma.getSizeInBytes
+ }
+ size
+ }
+
+ @Since("1.5.0")
+ override def copy(extra: ParamMap): NaiveBayesModel = {
+ copyValues(new NaiveBayesModel(uid, pi, theta, sigma).setParent(this.parent), extra)
+ }
+
+ @Since("1.5.0")
+ override def toString: String = {
+ s"NaiveBayesModel: uid=$uid, modelType=${$(modelType)}, numClasses=$numClasses, " +
+ s"numFeatures=$numFeatures"
+ }
+
+ @Since("1.6.0")
+ override def write: MLWriter = new NaiveBayesModel.NaiveBayesModelWriter(this)
+}
+
+@Since("1.6.0")
+object NaiveBayesModel extends MLReadable[NaiveBayesModel] {
+ import NaiveBayes._
+
+ private def predictRawFunction(
+ modelType: String,
+ pi: Vector,
+ theta: Matrix,
+ sigma: Matrix): Vector => Vector = {
+ modelType match {
+ case Multinomial =>
+ features: Vector => multinomialCalculation(features, pi, theta)
+ case Complement =>
+ features: Vector => complementCalculation(features, theta)
+ case Bernoulli =>
+ val localPiMinusThetaSum = bernoulliPiMinusThetaSum(pi, theta)
+ val localThetaMinusNegTheta = bernoulliThetaMinusNegTheta(theta)
+ features: Vector =>
+ bernoulliCalculation(features, localPiMinusThetaSum, localThetaMinusNegTheta)
+ case Gaussian =>
+ val localLogVarSum = gaussianLogVarSum(sigma)
+ features: Vector =>
+ gaussianCalculation(features, pi, theta, sigma, localLogVarSum)
+ }
+ }
+
+ private def multinomialCalculation(
+ features: Vector,
+ pi: Vector,
+ theta: Matrix): Vector = {
requireNonnegativeValues(features)
val prob = pi.toDense.copy
BLAS.gemv(1.0, theta, features, 1.0, prob)
prob
}
- private def complementCalculation(features: Vector) = {
+ private def complementCalculation(features: Vector, theta: Matrix): Vector = {
requireNonnegativeValues(features)
val probArray = theta.multiply(features).toArray
// the following lines equal to:
@@ -518,20 +635,48 @@ class NaiveBayesModel private[ml] (
Vectors.dense(probArray)
}
- private def bernoulliCalculation(features: Vector) = {
+ private def bernoulliThetaMinusNegTheta(theta: Matrix): Matrix = {
+ theta.map(value => value - math.log1p(-math.exp(value)))
+ }
+
+ private def bernoulliPiMinusThetaSum(pi: Vector, theta: Matrix): DenseVector = {
+ val negTheta = theta.map(value => math.log1p(-math.exp(value)))
+ val ones = new DenseVector(Array.fill(theta.numCols)(1.0))
+ val piMinusThetaSum = pi.toDense.copy
+ BLAS.gemv(1.0, negTheta, ones, 1.0, piMinusThetaSum)
+ piMinusThetaSum
+ }
+
+ private def bernoulliCalculation(
+ features: Vector,
+ piMinusThetaSum: DenseVector,
+ thetaMinusNegTheta: Matrix): Vector = {
requireZeroOneBernoulliValues(features)
val prob = piMinusThetaSum.copy
BLAS.gemv(1.0, thetaMinusNegTheta, features, 1.0, prob)
prob
}
- private def gaussianCalculation(features: Vector) = {
- val prob = Array.ofDim[Double](numClasses)
+ private def gaussianLogVarSum(sigma: Matrix): Array[Double] = {
+ Array.tabulate(sigma.numRows) { i =>
+ Iterator.range(0, sigma.numCols).map { j =>
+ math.log(sigma(i, j))
+ }.sum
+ }
+ }
+
+ private def gaussianCalculation(
+ features: Vector,
+ pi: Vector,
+ theta: Matrix,
+ sigma: Matrix,
+ logVarSum: Array[Double]): Vector = {
+ val prob = Array.ofDim[Double](pi.size)
var i = 0
- while (i < numClasses) {
+ while (i < pi.size) {
var s = 0.0
var j = 0
- while (j < numFeatures) {
+ while (j < theta.numCols) {
val d = features(j) - theta(i, j)
s += d * d / sigma(i, j)
j += 1
@@ -542,64 +687,17 @@ class NaiveBayesModel private[ml] (
Vectors.dense(prob)
}
- @transient private lazy val predictRawFunc = {
- $(modelType) match {
- case Multinomial =>
- features: Vector => multinomialCalculation(features)
- case Complement =>
- features: Vector => complementCalculation(features)
- case Bernoulli =>
- features: Vector => bernoulliCalculation(features)
- case Gaussian =>
- features: Vector => gaussianCalculation(features)
- }
- }
-
- @Since("3.0.0")
- override def predictRaw(features: Vector): Vector = predictRawFunc(features)
-
- override protected def raw2probabilityInPlace(rawPrediction: Vector): Vector = {
+ private def raw2probabilityInPlace(rawPrediction: Vector): Vector = {
rawPrediction match {
case dv: DenseVector =>
Utils.softmax(dv.values)
dv
- case sv: SparseVector =>
+ case _: SparseVector =>
throw new RuntimeException("Unexpected error in NaiveBayesModel:" +
" raw2probabilityInPlace encountered SparseVector")
}
}
- private[spark] override def estimatedSize: Long = {
- var size = this.estimateMatadataSize
- if (this.pi != null) {
- size += this.pi.getSizeInBytes
- }
- if (this.theta != null) {
- size += this.theta.getSizeInBytes
- }
- if (this.sigma != null) {
- size += this.sigma.getSizeInBytes
- }
- size
- }
-
- @Since("1.5.0")
- override def copy(extra: ParamMap): NaiveBayesModel = {
- copyValues(new NaiveBayesModel(uid, pi, theta, sigma).setParent(this.parent), extra)
- }
-
- @Since("1.5.0")
- override def toString: String = {
- s"NaiveBayesModel: uid=$uid, modelType=${$(modelType)}, numClasses=$numClasses, " +
- s"numFeatures=$numFeatures"
- }
-
- @Since("1.6.0")
- override def write: MLWriter = new NaiveBayesModel.NaiveBayesModelWriter(this)
-}
-
-@Since("1.6.0")
-object NaiveBayesModel extends MLReadable[NaiveBayesModel] {
private[ml] case class Data(pi: Vector, theta: Matrix, sigma: Matrix)
private[ml] def serializeData(data: Data, dos: DataOutputStream): Unit = {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org