From 20f434af813693c53e58ea3ed66e0c5fbbfcaa0c Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Fri, 4 Sep 2026 08:26:39 +0000 Subject: [PATCH] [SPARK-59247][ML] Optimize tree ensemble classification model transform closures --- .../ml/classification/GBTClassifier.scala | 95 ++++++++++++++++++- .../RandomForestClassifier.scala | 88 ++++++++++++++++- .../org/apache/spark/ml/tree/treeModels.scala | 10 ++ 3 files changed, 191 insertions(+), 2 deletions(-) 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..f6fcab5724983 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._ @@ -304,7 +305,10 @@ 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 { @@ -312,6 +316,65 @@ 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 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" 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..cf990923bc7fd 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 @@ -305,7 +305,10 @@ 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 { @@ -313,6 +316,61 @@ 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 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 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 25d035e819460..99ed9ef8dc460 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 @@ -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,