From f1635b5adb8b3458a492ed8a11d90d65f8bbd6d6 Mon Sep 17 00:00:00 2001
From: Ruifeng Zheng <ruifengz@apache.org>
Date: Tue, 1 Sep 2026 15:45:07 +0000
Subject: [PATCH 1/4] [SPARK-59154][ML] Optimize tree classification model
 transform closures

---
 .../DecisionTreeClassifier.scala              | 61 +++++++++++-
 .../ml/classification/GBTClassifier.scala     | 96 +++++++++++++++++++
 .../RandomForestClassifier.scala              | 94 ++++++++++++++++++
 3 files changed, 250 insertions(+), 1 deletion(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala
index 11d43d1b31fb1..381d1b96d631a 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala
@@ -209,6 +209,55 @@ class DecisionTreeClassificationModel private[ml] (
     rootNode.predictImpl(features).prediction
   }
 
+  override protected def predictRawColumn(features: Column): Column = {
+    val localRootNode = rootNode
+    udf((features: Vector) =>
+      DecisionTreeClassificationModel.predictRaw(features, localRootNode)
+    ).apply(features)
+  }
+
+  override protected def raw2probabilityColumn(rawPrediction: Column): Column = {
+    udf((rawPrediction: Vector) =>
+      DecisionTreeClassificationModel.raw2probability(rawPrediction)
+    ).apply(rawPrediction)
+  }
+
+  override protected def predictProbabilityColumn(features: Column): Column = {
+    val localRootNode = rootNode
+    udf((features: Vector) => {
+      val rawPrediction = DecisionTreeClassificationModel.predictRaw(features, localRootNode)
+      DecisionTreeClassificationModel.raw2probability(rawPrediction)
+    }).apply(features)
+  }
+
+  override protected def raw2predictionColumn(rawPrediction: Column): Column = {
+    if (isDefined(thresholds)) {
+      val localThresholds = getThresholds.clone()
+      udf((rawPrediction: Vector) => {
+        val probability = DecisionTreeClassificationModel.raw2probability(rawPrediction)
+        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+      }).apply(rawPrediction)
+    } else {
+      udf((rawPrediction: Vector) => rawPrediction.argmax.toDouble).apply(rawPrediction)
+    }
+  }
+
+  override protected def probability2predictionColumn(probability: Column): Column = {
+    if (isDefined(thresholds)) {
+      val localThresholds = getThresholds.clone()
+      udf((probability: Vector) =>
+        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+      ).apply(probability)
+    } else {
+      udf((probability: Vector) => probability.argmax.toDouble).apply(probability)
+    }
+  }
+
+  override protected def predictionColumn(features: Column): Column = {
+    val localRootNode = rootNode
+    udf((features: Vector) => localRootNode.predictImpl(features).prediction).apply(features)
+  }
+
   @Since("3.0.0")
   override def transformSchema(schema: StructType): StructType = {
     var outputSchema = super.transformSchema(schema)
@@ -233,7 +282,7 @@ class DecisionTreeClassificationModel private[ml] (
 
   @Since("3.0.0")
   override def predictRaw(features: Vector): Vector = {
-    Vectors.dense(rootNode.predictImpl(features).impurityStats.stats.clone())
+    DecisionTreeClassificationModel.predictRaw(features, rootNode)
   }
 
   override protected def raw2probabilityInPlace(rawPrediction: Vector): Vector = {
@@ -291,6 +340,16 @@ class DecisionTreeClassificationModel private[ml] (
 @Since("2.0.0")
 object DecisionTreeClassificationModel extends MLReadable[DecisionTreeClassificationModel] {
 
+  private def predictRaw(features: Vector, rootNode: Node): Vector = {
+    Vectors.dense(rootNode.predictImpl(features).impurityStats.stats.clone())
+  }
+
+  private def raw2probability(rawPrediction: Vector): Vector = {
+    val probability = rawPrediction.copy.toDense
+    ProbabilisticClassificationModel.normalizeToProbabilitiesInPlace(probability)
+    probability
+  }
+
   @Since("2.0.0")
   override def read: MLReader[DecisionTreeClassificationModel] =
     new DecisionTreeClassificationModelReader
diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala
index 19e49f299cbe5..83e2c0eef35a4 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala
@@ -32,6 +32,7 @@ import org.apache.spark.ml.util.DatasetUtils.extractInstances
 import org.apache.spark.ml.util.DefaultParamsReader.Metadata
 import org.apache.spark.ml.util.Instrumentation.instrumented
 import org.apache.spark.mllib.tree.configuration.{Algo => OldAlgo}
+import org.apache.spark.mllib.tree.loss.{ClassificationLoss => OldClassificationLoss}
 import org.apache.spark.mllib.tree.model.{GradientBoostedTreesModel => OldGBTModel}
 import org.apache.spark.sql._
 import org.apache.spark.sql.functions._
@@ -312,6 +313,76 @@ class GBTClassificationModel private[ml](
     }
   }
 
+  override protected def predictRawColumn(features: Column): Column = {
+    val localRootNodes = _trees.map(_.rootNode)
+    val localTreeWeights = _treeWeights.clone()
+    udf((features: Vector) =>
+      GBTClassificationModel.predictRaw(features, localRootNodes, localTreeWeights)
+    ).apply(features)
+  }
+
+  override protected def raw2probabilityColumn(rawPrediction: Column): Column = {
+    val localLoss = loss
+    udf((rawPrediction: Vector) =>
+      GBTClassificationModel.raw2probability(rawPrediction, localLoss)
+    ).apply(rawPrediction)
+  }
+
+  override protected def predictProbabilityColumn(features: Column): Column = {
+    val localRootNodes = _trees.map(_.rootNode)
+    val localTreeWeights = _treeWeights.clone()
+    val localLoss = loss
+    udf((features: Vector) => {
+      val rawPrediction =
+        GBTClassificationModel.predictRaw(features, localRootNodes, localTreeWeights)
+      GBTClassificationModel.raw2probability(rawPrediction, localLoss)
+    }).apply(features)
+  }
+
+  override protected def raw2predictionColumn(rawPrediction: Column): Column = {
+    if (isDefined(thresholds)) {
+      val localThresholds = getThresholds.clone()
+      val localLoss = loss
+      udf((rawPrediction: Vector) => {
+        val probability = GBTClassificationModel.raw2probability(rawPrediction, localLoss)
+        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+      }).apply(rawPrediction)
+    } else {
+      udf((rawPrediction: Vector) => rawPrediction.argmax.toDouble).apply(rawPrediction)
+    }
+  }
+
+  override protected def probability2predictionColumn(probability: Column): Column = {
+    if (isDefined(thresholds)) {
+      val localThresholds = getThresholds.clone()
+      udf((probability: Vector) =>
+        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+      ).apply(probability)
+    } else {
+      udf((probability: Vector) => probability.argmax.toDouble).apply(probability)
+    }
+  }
+
+  override protected def predictionColumn(features: Column): Column = {
+    val localRootNodes = _trees.map(_.rootNode)
+    val localTreeWeights = _treeWeights.clone()
+    if (isDefined(thresholds)) {
+      val localThresholds = getThresholds.clone()
+      val localLoss = loss
+      udf((features: Vector) => {
+        val rawPrediction =
+          GBTClassificationModel.predictRaw(features, localRootNodes, localTreeWeights)
+        val probability = GBTClassificationModel.raw2probability(rawPrediction, localLoss)
+        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+      }).apply(features)
+    } else {
+      udf((features: Vector) => {
+        val margin = GBTClassificationModel.margin(features, localRootNodes, localTreeWeights)
+        if (margin > 0.0) 1.0 else 0.0
+      }).apply(features)
+    }
+  }
+
   override def predict(features: Vector): Double = {
     // If thresholds defined, use predictRaw to get probabilities, otherwise use optimization
     if (isDefined(thresholds)) {
@@ -403,6 +474,31 @@ class GBTClassificationModel private[ml](
 @Since("2.0.0")
 object GBTClassificationModel extends MLReadable[GBTClassificationModel] {
 
+  private def margin(
+      features: Vector,
+      rootNodes: Array[Node],
+      treeWeights: Array[Double]): Double = {
+    val treePredictions = rootNodes.map(_.predictImpl(features).prediction)
+    BLAS.nativeBLAS.ddot(rootNodes.length, treePredictions, 1, treeWeights, 1)
+  }
+
+  private def predictRaw(
+      features: Vector,
+      rootNodes: Array[Node],
+      treeWeights: Array[Double]): Vector = {
+    val prediction = margin(features, rootNodes, treeWeights)
+    Vectors.dense(-prediction, prediction)
+  }
+
+  private def raw2probability(
+      rawPrediction: Vector,
+      loss: OldClassificationLoss): Vector = {
+    val probability = rawPrediction.copy.toDense
+    probability.values(0) = loss.computeProbability(probability.values(0))
+    probability.values(1) = 1.0 - probability.values(0)
+    probability
+  }
+
   private val numFeaturesKey: String = "numFeatures"
   private val numTreesKey: String = "numTrees"
 
diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala
index 41bcf3729e5a1..adc490437b2ed 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala
@@ -313,6 +313,72 @@ class RandomForestClassificationModel private[ml] (
     }
   }
 
+  override protected def predictRawColumn(features: Column): Column = {
+    val localRootNodes = _trees.map(_.rootNode)
+    val localNumClasses = numClasses
+    udf((features: Vector) =>
+      RandomForestClassificationModel.predictRaw(features, localRootNodes, localNumClasses)
+    ).apply(features)
+  }
+
+  override protected def raw2probabilityColumn(rawPrediction: Column): Column = {
+    udf((rawPrediction: Vector) =>
+      RandomForestClassificationModel.raw2probability(rawPrediction)
+    ).apply(rawPrediction)
+  }
+
+  override protected def predictProbabilityColumn(features: Column): Column = {
+    val localRootNodes = _trees.map(_.rootNode)
+    val localNumClasses = numClasses
+    udf((features: Vector) => {
+      val rawPrediction =
+        RandomForestClassificationModel.predictRaw(features, localRootNodes, localNumClasses)
+      RandomForestClassificationModel.raw2probability(rawPrediction)
+    }).apply(features)
+  }
+
+  override protected def raw2predictionColumn(rawPrediction: Column): Column = {
+    if (isDefined(thresholds)) {
+      val localThresholds = getThresholds.clone()
+      udf((rawPrediction: Vector) => {
+        val probability = RandomForestClassificationModel.raw2probability(rawPrediction)
+        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+      }).apply(rawPrediction)
+    } else {
+      udf((rawPrediction: Vector) => rawPrediction.argmax.toDouble).apply(rawPrediction)
+    }
+  }
+
+  override protected def probability2predictionColumn(probability: Column): Column = {
+    if (isDefined(thresholds)) {
+      val localThresholds = getThresholds.clone()
+      udf((probability: Vector) =>
+        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+      ).apply(probability)
+    } else {
+      udf((probability: Vector) => probability.argmax.toDouble).apply(probability)
+    }
+  }
+
+  override protected def predictionColumn(features: Column): Column = {
+    val localRootNodes = _trees.map(_.rootNode)
+    val localNumClasses = numClasses
+    if (isDefined(thresholds)) {
+      val localThresholds = getThresholds.clone()
+      udf((features: Vector) => {
+        val rawPrediction =
+          RandomForestClassificationModel.predictRaw(features, localRootNodes, localNumClasses)
+        val probability = RandomForestClassificationModel.raw2probability(rawPrediction)
+        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+      }).apply(features)
+    } else {
+      udf((features: Vector) =>
+        RandomForestClassificationModel.predictRaw(
+          features, localRootNodes, localNumClasses).argmax.toDouble
+      ).apply(features)
+    }
+  }
+
   @Since("3.0.0")
   override def predictRaw(features: Vector): Vector = {
     // TODO: When we add a generic Bagging class, handle transform there: SPARK-7128
@@ -411,6 +477,34 @@ class RandomForestClassificationModel private[ml] (
 @Since("2.0.0")
 object RandomForestClassificationModel extends MLReadable[RandomForestClassificationModel] {
 
+  private def predictRaw(
+      features: Vector,
+      rootNodes: Array[Node],
+      numClasses: Int): Vector = {
+    // TODO: When we add a generic Bagging class, handle transform there: SPARK-7128
+    // Classifies using majority votes.
+    // Ignore the tree weights since all are 1.0 for now.
+    val votes = Array.ofDim[Double](numClasses)
+    rootNodes.foreach { rootNode =>
+      val classCounts = rootNode.predictImpl(features).impurityStats.stats
+      val total = classCounts.sum
+      if (total != 0) {
+        var i = 0
+        while (i < numClasses) {
+          votes(i) += classCounts(i) / total
+          i += 1
+        }
+      }
+    }
+    Vectors.dense(votes)
+  }
+
+  private def raw2probability(rawPrediction: Vector): Vector = {
+    val probability = rawPrediction.copy.toDense
+    ProbabilisticClassificationModel.normalizeToProbabilitiesInPlace(probability)
+    probability
+  }
+
   @Since("2.0.0")
   override def read: MLReader[RandomForestClassificationModel] =
     new RandomForestClassificationModelReader

From 62591ea241354104b20099b098c820c6e9a43f35 Mon Sep 17 00:00:00 2001
From: Ruifeng Zheng <ruifengz@apache.org>
Date: Wed, 2 Sep 2026 05:36:25 +0000
Subject: [PATCH 2/4] [SPARK-59154][ML] Share tree classification transform
 logic

---
 .../DecisionTreeClassifier.scala              | 48 ++-----------
 .../ml/classification/GBTClassifier.scala     | 67 +++++--------------
 .../ProbabilisticClassifier.scala             | 59 ++++++++++++++++
 .../RandomForestClassifier.scala              | 58 +++-------------
 4 files changed, 94 insertions(+), 138 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala
index 381d1b96d631a..8044532c8e9cd 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala
@@ -185,6 +185,7 @@ class DecisionTreeClassificationModel private[ml] (
     @Since("1.6.0")override val numFeatures: Int,
     @Since("1.5.0")override val numClasses: Int)
   extends ProbabilisticClassificationModel[Vector, DecisionTreeClassificationModel]
+  with ProbabilisticClassificationModelTransform[Vector, DecisionTreeClassificationModel]
   with DecisionTreeModel with DecisionTreeClassifierParams with MLWritable with Serializable {
 
   require(rootNode != null,
@@ -209,53 +210,18 @@ class DecisionTreeClassificationModel private[ml] (
     rootNode.predictImpl(features).prediction
   }
 
-  override protected def predictRawColumn(features: Column): Column = {
+  override protected def predictRawFunction: Vector => Vector = {
     val localRootNode = rootNode
-    udf((features: Vector) =>
-      DecisionTreeClassificationModel.predictRaw(features, localRootNode)
-    ).apply(features)
+    features => DecisionTreeClassificationModel.predictRaw(features, localRootNode)
   }
 
-  override protected def raw2probabilityColumn(rawPrediction: Column): Column = {
-    udf((rawPrediction: Vector) =>
-      DecisionTreeClassificationModel.raw2probability(rawPrediction)
-    ).apply(rawPrediction)
+  override protected def raw2probabilityFunction: Vector => Vector = {
+    rawPrediction => DecisionTreeClassificationModel.raw2probability(rawPrediction)
   }
 
-  override protected def predictProbabilityColumn(features: Column): Column = {
+  override protected def predictionFunction: Vector => Double = {
     val localRootNode = rootNode
-    udf((features: Vector) => {
-      val rawPrediction = DecisionTreeClassificationModel.predictRaw(features, localRootNode)
-      DecisionTreeClassificationModel.raw2probability(rawPrediction)
-    }).apply(features)
-  }
-
-  override protected def raw2predictionColumn(rawPrediction: Column): Column = {
-    if (isDefined(thresholds)) {
-      val localThresholds = getThresholds.clone()
-      udf((rawPrediction: Vector) => {
-        val probability = DecisionTreeClassificationModel.raw2probability(rawPrediction)
-        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
-      }).apply(rawPrediction)
-    } else {
-      udf((rawPrediction: Vector) => rawPrediction.argmax.toDouble).apply(rawPrediction)
-    }
-  }
-
-  override protected def probability2predictionColumn(probability: Column): Column = {
-    if (isDefined(thresholds)) {
-      val localThresholds = getThresholds.clone()
-      udf((probability: Vector) =>
-        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
-      ).apply(probability)
-    } else {
-      udf((probability: Vector) => probability.argmax.toDouble).apply(probability)
-    }
-  }
-
-  override protected def predictionColumn(features: Column): Column = {
-    val localRootNode = rootNode
-    udf((features: Vector) => localRootNode.predictImpl(features).prediction).apply(features)
+    features => localRootNode.predictImpl(features).prediction
   }
 
   @Since("3.0.0")
diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala
index 83e2c0eef35a4..75386ab3a8265 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala
@@ -242,6 +242,7 @@ class GBTClassificationModel private[ml](
     @Since("1.6.0") override val numFeatures: Int,
     @Since("2.2.0") override val numClasses: Int)
   extends ProbabilisticClassificationModel[Vector, GBTClassificationModel]
+  with ProbabilisticClassificationModelTransform[Vector, GBTClassificationModel]
   with GBTClassifierParams with TreeEnsembleModel[DecisionTreeRegressionModel]
   with MLWritable with Serializable {
 
@@ -313,73 +314,34 @@ class GBTClassificationModel private[ml](
     }
   }
 
-  override protected def predictRawColumn(features: Column): Column = {
+  override protected def predictRawFunction: Vector => Vector = {
     val localRootNodes = _trees.map(_.rootNode)
     val localTreeWeights = _treeWeights.clone()
-    udf((features: Vector) =>
-      GBTClassificationModel.predictRaw(features, localRootNodes, localTreeWeights)
-    ).apply(features)
+    features => GBTClassificationModel.predictRaw(features, localRootNodes, localTreeWeights)
   }
 
-  override protected def raw2probabilityColumn(rawPrediction: Column): Column = {
+  override protected def raw2probabilityFunction: Vector => Vector = {
     val localLoss = loss
-    udf((rawPrediction: Vector) =>
-      GBTClassificationModel.raw2probability(rawPrediction, localLoss)
-    ).apply(rawPrediction)
+    rawPrediction => GBTClassificationModel.raw2probability(rawPrediction, localLoss)
   }
 
-  override protected def predictProbabilityColumn(features: Column): Column = {
-    val localRootNodes = _trees.map(_.rootNode)
-    val localTreeWeights = _treeWeights.clone()
-    val localLoss = loss
-    udf((features: Vector) => {
-      val rawPrediction =
-        GBTClassificationModel.predictRaw(features, localRootNodes, localTreeWeights)
-      GBTClassificationModel.raw2probability(rawPrediction, localLoss)
-    }).apply(features)
-  }
-
-  override protected def raw2predictionColumn(rawPrediction: Column): Column = {
-    if (isDefined(thresholds)) {
-      val localThresholds = getThresholds.clone()
-      val localLoss = loss
-      udf((rawPrediction: Vector) => {
-        val probability = GBTClassificationModel.raw2probability(rawPrediction, localLoss)
-        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
-      }).apply(rawPrediction)
-    } else {
-      udf((rawPrediction: Vector) => rawPrediction.argmax.toDouble).apply(rawPrediction)
-    }
-  }
-
-  override protected def probability2predictionColumn(probability: Column): Column = {
-    if (isDefined(thresholds)) {
-      val localThresholds = getThresholds.clone()
-      udf((probability: Vector) =>
-        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
-      ).apply(probability)
-    } else {
-      udf((probability: Vector) => probability.argmax.toDouble).apply(probability)
-    }
-  }
-
-  override protected def predictionColumn(features: Column): Column = {
+  override protected def predictionFunction: Vector => Double = {
     val localRootNodes = _trees.map(_.rootNode)
     val localTreeWeights = _treeWeights.clone()
     if (isDefined(thresholds)) {
       val localThresholds = getThresholds.clone()
       val localLoss = loss
-      udf((features: Vector) => {
+      features => {
         val rawPrediction =
           GBTClassificationModel.predictRaw(features, localRootNodes, localTreeWeights)
         val probability = GBTClassificationModel.raw2probability(rawPrediction, localLoss)
         ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
-      }).apply(features)
+      }
     } else {
-      udf((features: Vector) => {
+      features => {
         val margin = GBTClassificationModel.margin(features, localRootNodes, localTreeWeights)
         if (margin > 0.0) 1.0 else 0.0
-      }).apply(features)
+      }
     }
   }
 
@@ -478,8 +440,13 @@ object GBTClassificationModel extends MLReadable[GBTClassificationModel] {
       features: Vector,
       rootNodes: Array[Node],
       treeWeights: Array[Double]): Double = {
-    val treePredictions = rootNodes.map(_.predictImpl(features).prediction)
-    BLAS.nativeBLAS.ddot(rootNodes.length, treePredictions, 1, treeWeights, 1)
+    var prediction = 0.0
+    var i = 0
+    while (i < rootNodes.length) {
+      prediction += rootNodes(i).predictImpl(features).prediction * treeWeights(i)
+      i += 1
+    }
+    prediction
   }
 
   private def predictRaw(
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 2e3e681df9391..845500093fb6e 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
@@ -351,6 +351,65 @@ abstract class ProbabilisticClassificationModel[
   }
 }
 
+private[classification] trait ProbabilisticClassificationModelTransform[
+    FeaturesType,
+    M <: ProbabilisticClassificationModel[FeaturesType, M]]
+  extends ProbabilisticClassificationModel[FeaturesType, M] {
+
+  protected def predictRawFunction: FeaturesType => Vector
+
+  protected def raw2probabilityFunction: Vector => Vector
+
+  protected def predictionFunction: FeaturesType => Double
+
+  override protected def predictRawColumn(features: Column): Column = {
+    val localPredictRaw = predictRawFunction
+    udf { value: Any => localPredictRaw(value.asInstanceOf[FeaturesType]) }.apply(features)
+  }
+
+  override protected def raw2probabilityColumn(rawPrediction: Column): Column = {
+    val localRaw2probability = raw2probabilityFunction
+    udf((rawPrediction: Vector) => localRaw2probability(rawPrediction)).apply(rawPrediction)
+  }
+
+  override protected def predictProbabilityColumn(features: Column): Column = {
+    val localPredictRaw = predictRawFunction
+    val localRaw2probability = raw2probabilityFunction
+    udf((value: Any) => {
+      localRaw2probability(localPredictRaw(value.asInstanceOf[FeaturesType]))
+    }).apply(features)
+  }
+
+  override protected def raw2predictionColumn(rawPrediction: Column): Column = {
+    if (isDefined(thresholds)) {
+      val localRaw2probability = raw2probabilityFunction
+      val localThresholds = getThresholds.clone()
+      udf((rawPrediction: Vector) => {
+        val probability = localRaw2probability(rawPrediction)
+        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+      }).apply(rawPrediction)
+    } else {
+      udf((rawPrediction: Vector) => rawPrediction.argmax.toDouble).apply(rawPrediction)
+    }
+  }
+
+  override protected def probability2predictionColumn(probability: Column): Column = {
+    if (isDefined(thresholds)) {
+      val localThresholds = getThresholds.clone()
+      udf((probability: Vector) =>
+        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+      ).apply(probability)
+    } else {
+      udf((probability: Vector) => probability.argmax.toDouble).apply(probability)
+    }
+  }
+
+  override protected def predictionColumn(features: Column): Column = {
+    val localPredict = predictionFunction
+    udf { value: Any => localPredict(value.asInstanceOf[FeaturesType]) }.apply(features)
+  }
+}
+
 private[ml] object ProbabilisticClassificationModel {
 
   def probability2prediction(probability: Vector, thresholds: Array[Double]): Double = {
diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala
index adc490437b2ed..e4a7021fd4a8d 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala
@@ -225,6 +225,7 @@ class RandomForestClassificationModel private[ml] (
     @Since("1.6.0") override val numFeatures: Int,
     @Since("1.5.0") override val numClasses: Int)
   extends ProbabilisticClassificationModel[Vector, RandomForestClassificationModel]
+  with ProbabilisticClassificationModelTransform[Vector, RandomForestClassificationModel]
   with RandomForestClassifierParams with TreeEnsembleModel[DecisionTreeClassificationModel]
   with MLWritable with Serializable
   with HasTrainingSummary[RandomForestClassificationTrainingSummary] {
@@ -313,69 +314,32 @@ class RandomForestClassificationModel private[ml] (
     }
   }
 
-  override protected def predictRawColumn(features: Column): Column = {
+  override protected def predictRawFunction: Vector => Vector = {
     val localRootNodes = _trees.map(_.rootNode)
     val localNumClasses = numClasses
-    udf((features: Vector) =>
+    features =>
       RandomForestClassificationModel.predictRaw(features, localRootNodes, localNumClasses)
-    ).apply(features)
   }
 
-  override protected def raw2probabilityColumn(rawPrediction: Column): Column = {
-    udf((rawPrediction: Vector) =>
-      RandomForestClassificationModel.raw2probability(rawPrediction)
-    ).apply(rawPrediction)
+  override protected def raw2probabilityFunction: Vector => Vector = {
+    rawPrediction => RandomForestClassificationModel.raw2probability(rawPrediction)
   }
 
-  override protected def predictProbabilityColumn(features: Column): Column = {
+  override protected def predictionFunction: Vector => Double = {
     val localRootNodes = _trees.map(_.rootNode)
     val localNumClasses = numClasses
-    udf((features: Vector) => {
-      val rawPrediction =
-        RandomForestClassificationModel.predictRaw(features, localRootNodes, localNumClasses)
-      RandomForestClassificationModel.raw2probability(rawPrediction)
-    }).apply(features)
-  }
-
-  override protected def raw2predictionColumn(rawPrediction: Column): Column = {
-    if (isDefined(thresholds)) {
-      val localThresholds = getThresholds.clone()
-      udf((rawPrediction: Vector) => {
-        val probability = RandomForestClassificationModel.raw2probability(rawPrediction)
-        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
-      }).apply(rawPrediction)
-    } else {
-      udf((rawPrediction: Vector) => rawPrediction.argmax.toDouble).apply(rawPrediction)
-    }
-  }
-
-  override protected def probability2predictionColumn(probability: Column): Column = {
     if (isDefined(thresholds)) {
       val localThresholds = getThresholds.clone()
-      udf((probability: Vector) =>
-        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
-      ).apply(probability)
-    } else {
-      udf((probability: Vector) => probability.argmax.toDouble).apply(probability)
-    }
-  }
-
-  override protected def predictionColumn(features: Column): Column = {
-    val localRootNodes = _trees.map(_.rootNode)
-    val localNumClasses = numClasses
-    if (isDefined(thresholds)) {
-      val localThresholds = getThresholds.clone()
-      udf((features: Vector) => {
+      features => {
         val rawPrediction =
           RandomForestClassificationModel.predictRaw(features, localRootNodes, localNumClasses)
         val probability = RandomForestClassificationModel.raw2probability(rawPrediction)
         ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
-      }).apply(features)
+      }
     } else {
-      udf((features: Vector) =>
-        RandomForestClassificationModel.predictRaw(
-          features, localRootNodes, localNumClasses).argmax.toDouble
-      ).apply(features)
+      features =>
+        RandomForestClassificationModel.predictRaw(features, localRootNodes, localNumClasses)
+          .argmax.toDouble
     }
   }
 

From 579155f01749f76bc95ff2c35bad6965a3a48888 Mon Sep 17 00:00:00 2001
From: Ruifeng Zheng <ruifengz@apache.org>
Date: Wed, 2 Sep 2026 08:50:23 +0000
Subject: [PATCH 3/4] [SPARK-59154][ML] Narrow optimization to decision tree
 models

---
 .../DecisionTreeClassifier.scala              | 53 +++++++++++++---
 .../ml/classification/GBTClassifier.scala     | 63 -------------------
 .../ProbabilisticClassifier.scala             | 59 -----------------
 .../RandomForestClassifier.scala              | 58 -----------------
 .../ml/regression/DecisionTreeRegressor.scala | 13 +++-
 .../org/apache/spark/ml/tree/treeModels.scala | 13 +++-
 6 files changed, 65 insertions(+), 194 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala
index 8044532c8e9cd..7b99e15c40a55 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala
@@ -185,7 +185,6 @@ class DecisionTreeClassificationModel private[ml] (
     @Since("1.6.0")override val numFeatures: Int,
     @Since("1.5.0")override val numClasses: Int)
   extends ProbabilisticClassificationModel[Vector, DecisionTreeClassificationModel]
-  with ProbabilisticClassificationModelTransform[Vector, DecisionTreeClassificationModel]
   with DecisionTreeModel with DecisionTreeClassifierParams with MLWritable with Serializable {
 
   require(rootNode != null,
@@ -210,18 +209,53 @@ class DecisionTreeClassificationModel private[ml] (
     rootNode.predictImpl(features).prediction
   }
 
-  override protected def predictRawFunction: Vector => Vector = {
+  override protected def predictRawColumn(features: Column): Column = {
     val localRootNode = rootNode
-    features => DecisionTreeClassificationModel.predictRaw(features, localRootNode)
+    udf((features: Vector) =>
+      DecisionTreeClassificationModel.predictRaw(features, localRootNode)
+    ).apply(features)
   }
 
-  override protected def raw2probabilityFunction: Vector => Vector = {
-    rawPrediction => DecisionTreeClassificationModel.raw2probability(rawPrediction)
+  override protected def raw2probabilityColumn(rawPrediction: Column): Column = {
+    udf((rawPrediction: Vector) =>
+      DecisionTreeClassificationModel.raw2probability(rawPrediction)
+    ).apply(rawPrediction)
   }
 
-  override protected def predictionFunction: Vector => Double = {
+  override protected def predictProbabilityColumn(features: Column): Column = {
     val localRootNode = rootNode
-    features => localRootNode.predictImpl(features).prediction
+    udf((features: Vector) => {
+      val rawPrediction = DecisionTreeClassificationModel.predictRaw(features, localRootNode)
+      DecisionTreeClassificationModel.raw2probability(rawPrediction)
+    }).apply(features)
+  }
+
+  override protected def raw2predictionColumn(rawPrediction: Column): Column = {
+    if (isDefined(thresholds)) {
+      val localThresholds = getThresholds.clone()
+      udf((rawPrediction: Vector) => {
+        val probability = DecisionTreeClassificationModel.raw2probability(rawPrediction)
+        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+      }).apply(rawPrediction)
+    } else {
+      udf((rawPrediction: Vector) => rawPrediction.argmax.toDouble).apply(rawPrediction)
+    }
+  }
+
+  override protected def probability2predictionColumn(probability: Column): Column = {
+    if (isDefined(thresholds)) {
+      val localThresholds = getThresholds.clone()
+      udf((probability: Vector) =>
+        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+      ).apply(probability)
+    } else {
+      udf((probability: Vector) => probability.argmax.toDouble).apply(probability)
+    }
+  }
+
+  override protected def predictionColumn(features: Column): Column = {
+    val localRootNode = rootNode
+    udf((features: Vector) => localRootNode.predictImpl(features).prediction).apply(features)
   }
 
   @Since("3.0.0")
@@ -238,7 +272,10 @@ class DecisionTreeClassificationModel private[ml] (
 
     val outputData = super.transform(dataset)
     if ($(leafCol).nonEmpty) {
-      val leafUDF = udf { features: Vector => predictLeaf(features) }
+      val localRootNode = rootNode
+      val leafUDF = udf { features: Vector =>
+        DecisionTreeModel.predictLeaf(features, localRootNode)
+      }
       outputData.withColumn($(leafCol), leafUDF(col($(featuresCol))),
         outputSchema($(leafCol)).metadata)
     } else {
diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala
index 75386ab3a8265..19e49f299cbe5 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala
@@ -32,7 +32,6 @@ import org.apache.spark.ml.util.DatasetUtils.extractInstances
 import org.apache.spark.ml.util.DefaultParamsReader.Metadata
 import org.apache.spark.ml.util.Instrumentation.instrumented
 import org.apache.spark.mllib.tree.configuration.{Algo => OldAlgo}
-import org.apache.spark.mllib.tree.loss.{ClassificationLoss => OldClassificationLoss}
 import org.apache.spark.mllib.tree.model.{GradientBoostedTreesModel => OldGBTModel}
 import org.apache.spark.sql._
 import org.apache.spark.sql.functions._
@@ -242,7 +241,6 @@ class GBTClassificationModel private[ml](
     @Since("1.6.0") override val numFeatures: Int,
     @Since("2.2.0") override val numClasses: Int)
   extends ProbabilisticClassificationModel[Vector, GBTClassificationModel]
-  with ProbabilisticClassificationModelTransform[Vector, GBTClassificationModel]
   with GBTClassifierParams with TreeEnsembleModel[DecisionTreeRegressionModel]
   with MLWritable with Serializable {
 
@@ -314,37 +312,6 @@ class GBTClassificationModel private[ml](
     }
   }
 
-  override protected def predictRawFunction: Vector => Vector = {
-    val localRootNodes = _trees.map(_.rootNode)
-    val localTreeWeights = _treeWeights.clone()
-    features => GBTClassificationModel.predictRaw(features, localRootNodes, localTreeWeights)
-  }
-
-  override protected def raw2probabilityFunction: Vector => Vector = {
-    val localLoss = loss
-    rawPrediction => GBTClassificationModel.raw2probability(rawPrediction, localLoss)
-  }
-
-  override protected def predictionFunction: Vector => Double = {
-    val localRootNodes = _trees.map(_.rootNode)
-    val localTreeWeights = _treeWeights.clone()
-    if (isDefined(thresholds)) {
-      val localThresholds = getThresholds.clone()
-      val localLoss = loss
-      features => {
-        val rawPrediction =
-          GBTClassificationModel.predictRaw(features, localRootNodes, localTreeWeights)
-        val probability = GBTClassificationModel.raw2probability(rawPrediction, localLoss)
-        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
-      }
-    } else {
-      features => {
-        val margin = GBTClassificationModel.margin(features, localRootNodes, localTreeWeights)
-        if (margin > 0.0) 1.0 else 0.0
-      }
-    }
-  }
-
   override def predict(features: Vector): Double = {
     // If thresholds defined, use predictRaw to get probabilities, otherwise use optimization
     if (isDefined(thresholds)) {
@@ -436,36 +403,6 @@ class GBTClassificationModel private[ml](
 @Since("2.0.0")
 object GBTClassificationModel extends MLReadable[GBTClassificationModel] {
 
-  private def margin(
-      features: Vector,
-      rootNodes: Array[Node],
-      treeWeights: Array[Double]): Double = {
-    var prediction = 0.0
-    var i = 0
-    while (i < rootNodes.length) {
-      prediction += rootNodes(i).predictImpl(features).prediction * treeWeights(i)
-      i += 1
-    }
-    prediction
-  }
-
-  private def predictRaw(
-      features: Vector,
-      rootNodes: Array[Node],
-      treeWeights: Array[Double]): Vector = {
-    val prediction = margin(features, rootNodes, treeWeights)
-    Vectors.dense(-prediction, prediction)
-  }
-
-  private def raw2probability(
-      rawPrediction: Vector,
-      loss: OldClassificationLoss): Vector = {
-    val probability = rawPrediction.copy.toDense
-    probability.values(0) = loss.computeProbability(probability.values(0))
-    probability.values(1) = 1.0 - probability.values(0)
-    probability
-  }
-
   private val numFeaturesKey: String = "numFeatures"
   private val numTreesKey: String = "numTrees"
 
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 845500093fb6e..2e3e681df9391 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
@@ -351,65 +351,6 @@ abstract class ProbabilisticClassificationModel[
   }
 }
 
-private[classification] trait ProbabilisticClassificationModelTransform[
-    FeaturesType,
-    M <: ProbabilisticClassificationModel[FeaturesType, M]]
-  extends ProbabilisticClassificationModel[FeaturesType, M] {
-
-  protected def predictRawFunction: FeaturesType => Vector
-
-  protected def raw2probabilityFunction: Vector => Vector
-
-  protected def predictionFunction: FeaturesType => Double
-
-  override protected def predictRawColumn(features: Column): Column = {
-    val localPredictRaw = predictRawFunction
-    udf { value: Any => localPredictRaw(value.asInstanceOf[FeaturesType]) }.apply(features)
-  }
-
-  override protected def raw2probabilityColumn(rawPrediction: Column): Column = {
-    val localRaw2probability = raw2probabilityFunction
-    udf((rawPrediction: Vector) => localRaw2probability(rawPrediction)).apply(rawPrediction)
-  }
-
-  override protected def predictProbabilityColumn(features: Column): Column = {
-    val localPredictRaw = predictRawFunction
-    val localRaw2probability = raw2probabilityFunction
-    udf((value: Any) => {
-      localRaw2probability(localPredictRaw(value.asInstanceOf[FeaturesType]))
-    }).apply(features)
-  }
-
-  override protected def raw2predictionColumn(rawPrediction: Column): Column = {
-    if (isDefined(thresholds)) {
-      val localRaw2probability = raw2probabilityFunction
-      val localThresholds = getThresholds.clone()
-      udf((rawPrediction: Vector) => {
-        val probability = localRaw2probability(rawPrediction)
-        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
-      }).apply(rawPrediction)
-    } else {
-      udf((rawPrediction: Vector) => rawPrediction.argmax.toDouble).apply(rawPrediction)
-    }
-  }
-
-  override protected def probability2predictionColumn(probability: Column): Column = {
-    if (isDefined(thresholds)) {
-      val localThresholds = getThresholds.clone()
-      udf((probability: Vector) =>
-        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
-      ).apply(probability)
-    } else {
-      udf((probability: Vector) => probability.argmax.toDouble).apply(probability)
-    }
-  }
-
-  override protected def predictionColumn(features: Column): Column = {
-    val localPredict = predictionFunction
-    udf { value: Any => localPredict(value.asInstanceOf[FeaturesType]) }.apply(features)
-  }
-}
-
 private[ml] object ProbabilisticClassificationModel {
 
   def probability2prediction(probability: Vector, thresholds: Array[Double]): Double = {
diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala
index e4a7021fd4a8d..41bcf3729e5a1 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala
@@ -225,7 +225,6 @@ class RandomForestClassificationModel private[ml] (
     @Since("1.6.0") override val numFeatures: Int,
     @Since("1.5.0") override val numClasses: Int)
   extends ProbabilisticClassificationModel[Vector, RandomForestClassificationModel]
-  with ProbabilisticClassificationModelTransform[Vector, RandomForestClassificationModel]
   with RandomForestClassifierParams with TreeEnsembleModel[DecisionTreeClassificationModel]
   with MLWritable with Serializable
   with HasTrainingSummary[RandomForestClassificationTrainingSummary] {
@@ -314,35 +313,6 @@ class RandomForestClassificationModel private[ml] (
     }
   }
 
-  override protected def predictRawFunction: Vector => Vector = {
-    val localRootNodes = _trees.map(_.rootNode)
-    val localNumClasses = numClasses
-    features =>
-      RandomForestClassificationModel.predictRaw(features, localRootNodes, localNumClasses)
-  }
-
-  override protected def raw2probabilityFunction: Vector => Vector = {
-    rawPrediction => RandomForestClassificationModel.raw2probability(rawPrediction)
-  }
-
-  override protected def predictionFunction: Vector => Double = {
-    val localRootNodes = _trees.map(_.rootNode)
-    val localNumClasses = numClasses
-    if (isDefined(thresholds)) {
-      val localThresholds = getThresholds.clone()
-      features => {
-        val rawPrediction =
-          RandomForestClassificationModel.predictRaw(features, localRootNodes, localNumClasses)
-        val probability = RandomForestClassificationModel.raw2probability(rawPrediction)
-        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
-      }
-    } else {
-      features =>
-        RandomForestClassificationModel.predictRaw(features, localRootNodes, localNumClasses)
-          .argmax.toDouble
-    }
-  }
-
   @Since("3.0.0")
   override def predictRaw(features: Vector): Vector = {
     // TODO: When we add a generic Bagging class, handle transform there: SPARK-7128
@@ -441,34 +411,6 @@ class RandomForestClassificationModel private[ml] (
 @Since("2.0.0")
 object RandomForestClassificationModel extends MLReadable[RandomForestClassificationModel] {
 
-  private def predictRaw(
-      features: Vector,
-      rootNodes: Array[Node],
-      numClasses: Int): Vector = {
-    // TODO: When we add a generic Bagging class, handle transform there: SPARK-7128
-    // Classifies using majority votes.
-    // Ignore the tree weights since all are 1.0 for now.
-    val votes = Array.ofDim[Double](numClasses)
-    rootNodes.foreach { rootNode =>
-      val classCounts = rootNode.predictImpl(features).impurityStats.stats
-      val total = classCounts.sum
-      if (total != 0) {
-        var i = 0
-        while (i < numClasses) {
-          votes(i) += classCounts(i) / total
-          i += 1
-        }
-      }
-    }
-    Vectors.dense(votes)
-  }
-
-  private def raw2probability(rawPrediction: Vector): Vector = {
-    val probability = rawPrediction.copy.toDense
-    ProbabilisticClassificationModel.normalizeToProbabilitiesInPlace(probability)
-    probability
-  }
-
   @Since("2.0.0")
   override def read: MLReader[RandomForestClassificationModel] =
     new RandomForestClassificationModelReader
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala
index 4a5c3b5790354..79339807a2b6b 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala
@@ -218,26 +218,33 @@ class DecisionTreeRegressionModel private[ml] (
   @Since("2.0.0")
   override def transform(dataset: Dataset[_]): DataFrame = {
     val outputSchema = transformSchema(dataset.schema, logging = true)
+    val localRootNode = rootNode
 
     var predictionColNames = Seq.empty[String]
     var predictionColumns = Seq.empty[Column]
 
     if ($(predictionCol).nonEmpty) {
-      val predictUDF = udf { features: Vector => predict(features) }
+      val predictUDF = udf { features: Vector =>
+        localRootNode.predictImpl(features).prediction
+      }
       predictionColNames :+= $(predictionCol)
       predictionColumns :+= predictUDF(col($(featuresCol)))
         .as($(predictionCol), outputSchema($(predictionCol)).metadata)
     }
 
     if (isDefined(varianceCol) && $(varianceCol).nonEmpty) {
-      val predictVarianceUDF = udf { features: Vector => predictVariance(features) }
+      val predictVarianceUDF = udf { features: Vector =>
+        localRootNode.predictImpl(features).impurityStats.calculate()
+      }
       predictionColNames :+= $(varianceCol)
       predictionColumns :+= predictVarianceUDF(col($(featuresCol)))
         .as($(varianceCol), outputSchema($(varianceCol)).metadata)
     }
 
     if ($(leafCol).nonEmpty) {
-      val leafUDF = udf { features: Vector => predictLeaf(features) }
+      val leafUDF = udf { features: Vector =>
+        DecisionTreeModel.predictLeaf(features, localRootNode)
+      }
       predictionColNames :+= $(leafCol)
       predictionColumns :+= leafUDF(col($(featuresCol)))
         .as($(leafCol), outputSchema($(leafCol)).metadata)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/tree/treeModels.scala b/mllib/src/main/scala/org/apache/spark/ml/tree/treeModels.scala
index 5f45cc3315621..70d5589ab7b6a 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/tree/treeModels.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/tree/treeModels.scala
@@ -94,9 +94,7 @@ private[spark] trait DecisionTreeModel {
    *         Leaves are indexed in pre-order from 0.
    */
   def predictLeaf(features: Vector): Double = {
-    val leaf = rootNode.predictImpl(features)
-    assert(leaf.leafIndex >= 0, "Leaf indices are not assigned.")
-    leaf.leafIndex.toDouble
+    DecisionTreeModel.predictLeaf(features, rootNode)
   }
 
   def getEstimatedSize(): Long = {
@@ -104,6 +102,15 @@ private[spark] trait DecisionTreeModel {
   }
 }
 
+private[spark] object DecisionTreeModel {
+
+  private[ml] def predictLeaf(features: Vector, rootNode: Node): Double = {
+    val leaf = rootNode.predictImpl(features)
+    assert(leaf.leafIndex >= 0, "Leaf indices are not assigned.")
+    leaf.leafIndex.toDouble
+  }
+}
+
 /**
  * Abstraction for models which are ensembles of decision trees
  * @tparam M  Type of tree model in this ensemble

From 706d8309d1621139aef6d89d11296b6cd9db1ce7 Mon Sep 17 00:00:00 2001
From: Ruifeng Zheng <ruifengz@apache.org>
Date: Wed, 2 Sep 2026 11:38:26 +0000
Subject: [PATCH 4/4] [SPARK-59154][ML] Share probability prediction column
 logic

---
 .../ml/classification/DecisionTreeClassifier.scala    | 11 -----------
 .../ml/classification/ProbabilisticClassifier.scala   |  9 ++++++++-
 2 files changed, 8 insertions(+), 12 deletions(-)

diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala
index 7b99e15c40a55..d83df697a04a6 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala
@@ -242,17 +242,6 @@ class DecisionTreeClassificationModel private[ml] (
     }
   }
 
-  override protected def probability2predictionColumn(probability: Column): Column = {
-    if (isDefined(thresholds)) {
-      val localThresholds = getThresholds.clone()
-      udf((probability: Vector) =>
-        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
-      ).apply(probability)
-    } else {
-      udf((probability: Vector) => probability.argmax.toDouble).apply(probability)
-    }
-  }
-
   override protected def predictionColumn(features: Column): Column = {
     val localRootNode = rootNode
     udf((features: Vector) => localRootNode.predictImpl(features).prediction).apply(features)
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 2e3e681df9391..8be033d22db84 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
@@ -194,7 +194,14 @@ abstract class ProbabilisticClassificationModel[
    * @note This method honors [[thresholds]] when they are set.
    */
   protected def probability2predictionColumn(probability: Column): Column = {
-    udf(probability2prediction _).apply(probability)
+    if (isDefined(thresholds)) {
+      val localThresholds = getThresholds.clone()
+      udf((probability: Vector) =>
+        ProbabilisticClassificationModel.probability2prediction(probability, localThresholds)
+      ).apply(probability)
+    } else {
+      udf((probability: Vector) => probability.argmax.toDouble).apply(probability)
+    }
   }
 
   /** @group setParam */