[SPARK-59247][ML] Optimize tree ensemble classification model transform closures

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._
@@ -304,14 +305,76 @@ class GBTClassificationModel private[ml](
 
     val outputData = super.transform(dataset)
     if ($(leafCol).nonEmpty) {
-      val leafUDF = udf { features: Vector => predictLeaf(features) }
+      val localRootNodes = _trees.map(_.rootNode)
+      val leafUDF = udf { features: Vector =>
+        TreeEnsembleModel.predictLeaf(features, localRootNodes)
+      }
       outputData.withColumn($(leafCol), leafUDF(col($(featuresCol))),
         outputSchema($(leafCol)).metadata)
     } else {
       outputData
     }
   }
 
+  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 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 +466,36 @@ 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"
 

mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala

@@ -305,14 +305,72 @@ class RandomForestClassificationModel private[ml] (
 
     val outputData = super.transform(dataset)
     if ($(leafCol).nonEmpty) {
-      val leafUDF = udf { features: Vector => predictLeaf(features) }
+      val localRootNodes = _trees.map(_.rootNode)
+      val leafUDF = udf { features: Vector =>
+        TreeEnsembleModel.predictLeaf(features, localRootNodes)
+      }
       outputData.withColumn($(leafCol), leafUDF(col($(featuresCol))),
         outputSchema($(leafCol)).metadata)
     } else {
       outputData
     }
   }
 
+  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 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 +469,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

mllib/src/main/scala/org/apache/spark/ml/tree/treeModels.scala

@@ -169,6 +169,16 @@ private[spark] trait TreeEnsembleModel[M <: DecisionTreeModel] {
 
 private[ml] object TreeEnsembleModel {
 
+  private[ml] def predictLeaf(features: Vector, rootNodes: Array[Node]): Vector = {
+    val indices = Array.ofDim[Double](rootNodes.length)
+    var i = 0
+    while (i < rootNodes.length) {
+      indices(i) = DecisionTreeModel.predictLeaf(features, rootNodes(i))
+      i += 1
+    }
+    Vectors.dense(indices)
+  }
+
   /**
    * Given a tree ensemble model, compute the importance of each feature.
    * This generalizes the idea of "Gini" importance to other losses,