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/08/20 23:54:26 UTC
(spark) branch master updated: [SPARK-58896][ML] Compute decision tree model stats in one traversal
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 dbf21278436d [SPARK-58896][ML] Compute decision tree model stats in one traversal
dbf21278436d is described below
commit dbf21278436dfd764f5e5b84f6ebbb9d6439b985
Author: Ruifeng Zheng <ru...@apache.org>
AuthorDate: Fri Aug 21 07:54:13 2026 +0800
[SPARK-58896][ML] Compute decision tree model stats in one traversal
### What changes were proposed in this pull request?
This PR computes decision tree model stats in a single traversal of the root node.
It adds `Node.computeStats`, which returns a `NodeStats` value containing the subtree depth,
number of descendants, and number of leaves. Decision tree classification and regression models
compute these stats once and `DecisionTreeModel` reuses them for `numNodes`, `depth`, and leaf
metadata.
This also renames the package-private `numLeave` helper to `numLeaves`.
### Why are the changes needed?
Before this change, decision tree model metadata could traverse the tree separately for depth,
node count, and leaf count. Computing these values together avoids redundant tree walks.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Manually tested with:
```
git diff --check
```
Also checked changed Scala files for long lines and non-ASCII characters.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex
Closes #58147 from zhengruifeng/tree-model-stats-single-pass-dev4.
Authored-by: Ruifeng Zheng <ru...@apache.org>
Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
---
.../ml/classification/DecisionTreeClassifier.scala | 2 ++
.../ml/regression/DecisionTreeRegressor.scala | 2 ++
.../main/scala/org/apache/spark/ml/tree/Node.scala | 28 ++++++++++++++++++++++
.../org/apache/spark/ml/tree/treeModels.scala | 18 +++++---------
.../DecisionTreeClassifierSuite.scala | 2 +-
project/MimaExcludes.scala | 9 ++++++-
6 files changed, 47 insertions(+), 14 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 3506c1a2502d..11d43d1b31fb 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
@@ -201,6 +201,8 @@ class DecisionTreeClassificationModel private[ml] (
// For ml connect only
private[ml] def this() = this("", Node.dummyNode, -1, -1)
+ override private[ml] val treeStats: NodeStats = rootNode.computeStats
+
private[spark] override def estimatedSize: Long = estimateMatadataSize + getEstimatedSize()
override def predict(features: Vector): Double = {
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 135224d9a8f9..4a5c3b579035 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
@@ -190,6 +190,8 @@ class DecisionTreeRegressionModel private[ml] (
// For ml connect only
private[ml] def this() = this("", Node.dummyNode, -1)
+ override private[ml] val treeStats: NodeStats = rootNode.computeStats
+
private[spark] override def estimatedSize: Long = estimateMatadataSize + getEstimatedSize()
override def predict(features: Vector): Double = {
diff --git a/mllib/src/main/scala/org/apache/spark/ml/tree/Node.scala b/mllib/src/main/scala/org/apache/spark/ml/tree/Node.scala
index a52b79b377fa..0a7e37aa0c9a 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/tree/Node.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/tree/Node.scala
@@ -77,8 +77,36 @@ sealed abstract class Node extends Serializable {
* @return Max feature index used in a split, or -1 if there are no splits (single leaf node).
*/
private[ml] def maxSplitFeatureIndex(): Int
+
+ private[ml] def computeStats: NodeStats = {
+ var numDescendants = 0
+ var subtreeDepth = 0
+ var numLeaves = 0
+
+ def visit(node: Node, depth: Int): Unit = {
+ if (depth > 0) {
+ numDescendants += 1
+ }
+ subtreeDepth = math.max(subtreeDepth, depth)
+ node match {
+ case _: LeafNode =>
+ numLeaves += 1
+ case internal: InternalNode =>
+ visit(internal.leftChild, depth + 1)
+ visit(internal.rightChild, depth + 1)
+ }
+ }
+
+ visit(this, 0)
+ NodeStats(subtreeDepth, numDescendants, numLeaves)
+ }
}
+private[ml] case class NodeStats(
+ subtreeDepth: Int,
+ numDescendants: Int,
+ numLeaves: Int)
+
private[ml] object Node {
/**
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 fdc4110573ba..1cc64e58d980 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
@@ -49,17 +49,13 @@ private[spark] trait DecisionTreeModel {
def rootNode: Node
/** Number of nodes in tree, including leaf nodes. */
- def numNodes: Int = {
- 1 + rootNode.numDescendants
- }
+ def numNodes: Int = treeStats.numDescendants + 1
/**
* Depth of the tree.
* E.g.: Depth 0 means 1 leaf node. Depth 1 means 1 internal node and 2 leaf nodes.
*/
- lazy val depth: Int = {
- rootNode.subtreeDepth
- }
+ def depth: Int = treeStats.subtreeDepth
/** Summary of the model */
override def toString: String = {
@@ -95,13 +91,11 @@ private[spark] trait DecisionTreeModel {
}
}
- private[ml] lazy val numLeave: Int =
- leafIterator(rootNode).size
+ private[ml] def treeStats: NodeStats
- private[ml] def leafAttr = {
- NominalAttribute.defaultAttr
- .withNumValues(numLeave)
- }
+ private[ml] def numLeaves: Int = treeStats.numLeaves
+
+ private[ml] def leafAttr = NominalAttribute.defaultAttr.withNumValues(numLeaves)
private[ml] def getLeafField(leafCol: String) = {
leafAttr.withName(leafCol).toStructField()
diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/DecisionTreeClassifierSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/DecisionTreeClassifierSuite.scala
index 765fccf6c620..d298e625880f 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/classification/DecisionTreeClassifierSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/classification/DecisionTreeClassifierSuite.scala
@@ -270,7 +270,7 @@ class DecisionTreeClassifierSuite extends MLTest with DefaultReadWriteTest {
val transformed = newTree.transform(newData)
checkNominalOnDF(transformed, "prediction", newTree.numClasses)
- checkNominalOnDF(transformed, "predictedLeafId", newTree.numLeave)
+ checkNominalOnDF(transformed, "predictedLeafId", newTree.numLeaves)
checkVectorSizeOnDF(transformed, "rawPrediction", newTree.numClasses)
checkVectorSizeOnDF(transformed, "probability", newTree.numClasses)
diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala
index 3701646cf7cd..c27ef26959ed 100644
--- a/project/MimaExcludes.scala
+++ b/project/MimaExcludes.scala
@@ -34,7 +34,14 @@ import com.typesafe.tools.mima.core.*
object MimaExcludes {
// Exclude rules for 5.0.x from 4.4.0 (add 5.0-specific filters below as needed).
- lazy val v50excludes: Seq[Problem => Boolean] = v44excludes
+ lazy val v50excludes: Seq[Problem => Boolean] = v44excludes ++ Seq(
+ // [SPARK-58896] Decision tree leaf counts moved behind NodeStats. The old package-private
+ // numLeave accessors are removed from concrete models.
+ ProblemFilters.exclude[DirectMissingMethodProblem](
+ "org.apache.spark.ml.classification.DecisionTreeClassificationModel.numLeave"),
+ ProblemFilters.exclude[DirectMissingMethodProblem](
+ "org.apache.spark.ml.regression.DecisionTreeRegressionModel.numLeave")
+ )
// Exclude rules for 4.4.x from 4.3.0 (add 4.4-specific filters below as needed).
lazy val v44excludes: Seq[Problem => Boolean] = v43excludes
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org