[SPARK-59400][ML] Reuse TreeEnsembleModel prediction helper

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

@@ -369,7 +369,7 @@ class GBTClassificationModel private[ml](
       }).apply(features)
     } else {
       udf((features: Vector) => {
-        val margin = TreeEnsembleModel.predict(features, localRootNodes, localTreeWeights)
+        val margin = TreeEnsembleModel.predictRaw(features, localRootNodes, localTreeWeights)
         if (margin > 0.0) 1.0 else 0.0
       }).apply(features)
     }
@@ -380,13 +380,13 @@ class GBTClassificationModel private[ml](
     if (isDefined(thresholds)) {
       super.predict(features)
     } else {
-      if (margin(features) > 0.0) 1.0 else 0.0
+      if (TreeEnsembleModel.predictRaw(features, _trees, _treeWeights) > 0.0) 1.0 else 0.0
     }
   }
 
   @Since("3.0.0")
   override def predictRaw(features: Vector): Vector = {
-    val prediction: Double = margin(features)
+    val prediction = TreeEnsembleModel.predictRaw(features, _trees, _treeWeights)
     Vectors.dense(Array(-prediction, prediction))
   }
 
@@ -428,17 +428,6 @@ class GBTClassificationModel private[ml](
   lazy val featureImportances: Vector =
     TreeEnsembleModel.featureImportances(trees, numFeatures, perTreeNormalization = false)
 
-  /** Raw prediction for the positive class. */
-  private def margin(features: Vector): Double = {
-    var prediction = 0.0
-    var i = 0
-    while (i < _trees.length) {
-      prediction += _trees(i).rootNode.predictImpl(features).prediction * _treeWeights(i)
-      i += 1
-    }
-    prediction
-  }
-
   /** (private[ml]) Convert to a model in the old API */
   private[ml] def toOld: OldGBTModel = {
     new OldGBTModel(OldAlgo.Classification, _trees.map(_.toOld), _treeWeights)
@@ -470,7 +459,7 @@ object GBTClassificationModel extends MLReadable[GBTClassificationModel] {
       features: Vector,
       rootNodes: Array[Node],
       treeWeights: Array[Double]): Vector = {
-    val prediction = TreeEnsembleModel.predict(features, rootNodes, treeWeights)
+    val prediction = TreeEnsembleModel.predictRaw(features, rootNodes, treeWeights)
     Vectors.dense(-prediction, prediction)
   }
 

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

@@ -372,24 +372,8 @@ class RandomForestClassificationModel private[ml] (
   }
 
   @Since("3.0.0")
-  override def predictRaw(features: Vector): 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)
-    _trees.foreach { tree =>
-      val classCounts = tree.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)
-  }
+  override def predictRaw(features: Vector): Vector =
+    RandomForestClassificationModel.predictRaw(features, _trees, numClasses)
 
   override protected def raw2probabilityInPlace(rawPrediction: Vector): Vector = {
     rawPrediction match {
@@ -469,13 +453,29 @@ class RandomForestClassificationModel private[ml] (
 @Since("2.0.0")
 object RandomForestClassificationModel extends MLReadable[RandomForestClassificationModel] {
 
+  private def predictRaw(
+      features: Vector,
+      trees: Array[DecisionTreeClassificationModel],
+      numClasses: Int): Vector = {
+    val votes = Array.ofDim[Double](numClasses)
+    trees.foreach { tree =>
+      val classCounts = tree.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 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

mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala

@@ -280,7 +280,7 @@ class GBTRegressionModel private[ml](
       if ($(predictionCol).nonEmpty) {
         val predUDF = udf { features: Vector =>
           val (rootNodes, treeWeights) = bcTreeData.value
-          TreeEnsembleModel.predict(features, rootNodes, treeWeights)
+          TreeEnsembleModel.predictRaw(features, rootNodes, treeWeights)
         }
         predColNames :+= $(predictionCol)
         predCols :+= predUDF(col($(featuresCol)))
@@ -304,17 +304,8 @@ class GBTRegressionModel private[ml](
     }
   }
 
-  override def predict(features: Vector): Double = {
-    // TODO: When we add a generic Boosting class, handle transform there?  SPARK-7129
-    // Classifies by thresholding sum of weighted tree predictions
-    var prediction = 0.0
-    var i = 0
-    while (i < _trees.length) {
-      prediction += _trees(i).rootNode.predictImpl(features).prediction * _treeWeights(i)
-      i += 1
-    }
-    prediction
-  }
+  override def predict(features: Vector): Double =
+    TreeEnsembleModel.predictRaw(features, _trees, _treeWeights)
 
   @Since("1.4.0")
   override def copy(extra: ParamMap): GBTRegressionModel = {

mllib/src/main/scala/org/apache/spark/ml/regression/RandomForestRegressor.scala

@@ -262,12 +262,8 @@ class RandomForestRegressionModel private[ml] (
     }
   }
 
-  override def predict(features: Vector): Double = {
-    // TODO: When we add a generic Bagging class, handle transform there.  SPARK-7128
-    // Predict average of tree predictions.
-    // Ignore the weights since all are 1.0 for now.
-    _trees.map(_.rootNode.predictImpl(features).prediction).sum / getNumTrees
-  }
+  override def predict(features: Vector): Double =
+    TreeEnsembleModel.predictRaw(features, _trees) / getNumTrees
 
   @Since("1.4.0")
   override def copy(extra: ParamMap): RandomForestRegressionModel = {

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

@@ -169,7 +169,32 @@ private[spark] trait TreeEnsembleModel[M <: DecisionTreeModel] {
 
 private[ml] object TreeEnsembleModel {
 
-  private[ml] def predict(
+  private[ml] def predictRaw[M <: DecisionTreeModel](
+      features: Vector,
+      trees: Array[M],
+      treeWeights: Array[Double]): Double = {
+    var prediction = 0.0
+    var i = 0
+    while (i < trees.length) {
+      prediction += trees(i).rootNode.predictImpl(features).prediction * treeWeights(i)
+      i += 1
+    }
+    prediction
+  }
+
+  private[ml] def predictRaw[M <: DecisionTreeModel](
+      features: Vector,
+      trees: Array[M]): Double = {
+    var prediction = 0.0
+    var i = 0
+    while (i < trees.length) {
+      prediction += trees(i).rootNode.predictImpl(features).prediction
+      i += 1
+    }
+    prediction
+  }
+
+  private[ml] def predictRaw(
       features: Vector,
       rootNodes: Array[Node],
       treeWeights: Array[Double]): Double = {