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 e2cc1d64ab36 Revert "[SPARK-58856][ML] Avoid per-leaf map for tree
predictions"
e2cc1d64ab36 is described below
commit e2cc1d64ab363aa9fa5e4b0bdaffe20d93ff98df
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Thu Aug 27 15:13:02 2026 +0800
Revert "[SPARK-58856][ML] Avoid per-leaf map for tree predictions"
### What changes were proposed in this pull request?
This PR reverts #57389 (`bc9ee80d0de08f4cca1ac8518e4b6bee64cc544d`) on
`master`.
It restores the transient lazy leaf-to-index map used by `predictLeaf` and
removes the
post-construction leaf-index field, initialization passes, and associated
persistence assertions.
The resulting eight files exactly match their versions immediately before
#57389.
### Why are the changes needed?
#57389 replaces the per-leaf map by constructing an unindexed node graph
and then rebuilding the
entire tree to attach indices. Although trees in an ensemble are processed
sequentially, this adds
O(number of nodes in one tree) temporary driver memory during training,
loading, and old-model
conversion. A sufficiently large individual tree can therefore cause an
out-of-memory failure.
Reverting removes that regression while an implementation that assigns
immutable indices during
initial node construction is considered separately in #58330.
### Does this PR introduce _any_ user-facing change?
No compared with a released Spark version. This restores the behavior that
existed before the
unreleased change in #57389. Leaf IDs and their traversal order are
unchanged.
### How was this patch tested?
The revert was verified to apply cleanly to current `master`, and all
affected files were compared
with their pre-#57389 versions. Targeted ML tests and Scala lint have not
been run yet; the PR was
opened first as requested.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (GPT-5)
Closes #58332 from zhengruifeng/ml-revert-leaf-index-dev-7.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../ml/classification/DecisionTreeClassifier.scala | 2 +-
.../ml/regression/DecisionTreeRegressor.scala | 2 +-
.../main/scala/org/apache/spark/ml/tree/Node.scala | 22 +---------------------
.../apache/spark/ml/tree/impl/RandomForest.scala | 13 ++++---------
.../org/apache/spark/ml/tree/treeModels.scala | 10 ++++++----
.../DecisionTreeClassifierSuite.scala | 2 --
.../ml/regression/DecisionTreeRegressorSuite.scala | 2 --
.../org/apache/spark/ml/tree/impl/TreeTests.scala | 10 ++++------
8 files changed, 17 insertions(+), 46 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 78b747f65f2f..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
@@ -343,7 +343,7 @@ object DecisionTreeClassificationModel extends
MLReadable[DecisionTreeClassifica
require(oldModel.algo == OldAlgo.Classification,
s"Cannot convert non-classification DecisionTreeModel (old API) to" +
s" DecisionTreeClassificationModel (new API). Algo is:
${oldModel.algo}")
- val rootNode = Node.withLeafIndices(Node.fromOld(oldModel.topNode,
categoricalFeatures))
+ val rootNode = Node.fromOld(oldModel.topNode, categoricalFeatures)
val uid = if (parent != null) parent.uid else Identifiable.randomUID("dtc")
// Can't infer number of features from old model, so default to -1
new DecisionTreeClassificationModel(uid, rootNode, numFeatures, -1)
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 dc2377166e17..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
@@ -345,7 +345,7 @@ object DecisionTreeRegressionModel extends
MLReadable[DecisionTreeRegressionMode
require(oldModel.algo == OldAlgo.Regression,
s"Cannot convert non-regression DecisionTreeModel (old API) to" +
s" DecisionTreeRegressionModel (new API). Algo is: ${oldModel.algo}")
- val rootNode = Node.withLeafIndices(Node.fromOld(oldModel.topNode,
categoricalFeatures))
+ val rootNode = Node.fromOld(oldModel.topNode, categoricalFeatures)
val uid = if (parent != null) parent.uid else Identifiable.randomUID("dtr")
new DecisionTreeRegressionModel(uid, rootNode, numFeatures)
}
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 72da973b5c0d..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
@@ -135,25 +135,6 @@ private[ml] object Node {
val dummyNode: Node = {
new LeafNode(0.0, 0.0, ImpurityCalculator.getCalculator("gini",
Array.empty, 0))
}
-
- /** Return a copy of the tree with left-to-right DFS indices assigned to
leaves. */
- private[ml] def withLeafIndices(rootNode: Node): Node = {
- withLeafIndices(rootNode, 0)._1
- }
-
- private def withLeafIndices(node: Node, nextLeafIndex: Int): (Node, Int) = {
- node match {
- case n: LeafNode =>
- (new LeafNode(n.prediction, n.impurity, n.impurityStats,
nextLeafIndex),
- nextLeafIndex + 1)
- case n: InternalNode =>
- val (leftChild, nextIndex) = withLeafIndices(n.leftChild,
nextLeafIndex)
- val (rightChild, endIndex) = withLeafIndices(n.rightChild, nextIndex)
- val newNode = new InternalNode(n.prediction, n.impurity, n.gain,
leftChild, rightChild,
- n.split, n.impurityStats)
- (newNode, endIndex)
- }
- }
}
/**
@@ -164,8 +145,7 @@ private[ml] object Node {
class LeafNode private[ml] (
override val prediction: Double,
override val impurity: Double,
- override private[ml] val impurityStats: ImpurityCalculator,
- private[tree] val leafIndex: Int = -1) extends Node {
+ override private[ml] val impurityStats: ImpurityCalculator) extends Node {
override def toString: String =
s"LeafNode(prediction = $prediction, impurity = $impurity)"
diff --git
a/mllib/src/main/scala/org/apache/spark/ml/tree/impl/RandomForest.scala
b/mllib/src/main/scala/org/apache/spark/ml/tree/impl/RandomForest.scala
index 3e4c42afa7f2..b3ca7f04c3de 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/tree/impl/RandomForest.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/tree/impl/RandomForest.scala
@@ -269,31 +269,26 @@ private[spark] object RandomForest extends Logging with
Serializable {
topNodes.map { rootNode =>
new DecisionTreeClassificationModel(
uid,
- Node.withLeafIndices(rootNode.toNode(strategy.pruneTree)),
+ rootNode.toNode(strategy.pruneTree),
numFeatures,
strategy.getNumClasses())
}
} else {
topNodes.map { rootNode =>
- new DecisionTreeRegressionModel(
- uid,
- Node.withLeafIndices(rootNode.toNode(strategy.pruneTree)),
- numFeatures)
+ new DecisionTreeRegressionModel(uid,
rootNode.toNode(strategy.pruneTree), numFeatures)
}
}
case None =>
if (strategy.algo == OldAlgo.Classification) {
topNodes.map { rootNode =>
new DecisionTreeClassificationModel(
- Node.withLeafIndices(rootNode.toNode(strategy.pruneTree)),
+ rootNode.toNode(strategy.pruneTree),
numFeatures,
strategy.getNumClasses())
}
} else {
topNodes.map(rootNode =>
- new DecisionTreeRegressionModel(
- Node.withLeafIndices(rootNode.toNode(strategy.pruneTree)),
- numFeatures))
+ new
DecisionTreeRegressionModel(rootNode.toNode(strategy.pruneTree), numFeatures))
}
}
}
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 ca923a1f751f..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
@@ -101,14 +101,16 @@ private[spark] trait DecisionTreeModel {
leafAttr.withName(leafCol).toStructField()
}
+ @transient private lazy val leafIndices: Map[LeafNode, Int] = {
+ leafIterator(rootNode).zipWithIndex.toMap
+ }
+
/**
* @return The index of the leaf corresponding to the feature vector.
* 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
+ leafIndices(rootNode.predictImpl(features)).toDouble
}
def getEstimatedSize(): Long = {
@@ -506,7 +508,7 @@ private[ml] object DecisionTreeModelReadWrite {
finalNodes(n.id) = node
}
// Return the root node
- Node.withLeafIndices(finalNodes.head)
+ finalNodes.head
}
}
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 e248d0f51b6e..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
@@ -427,8 +427,6 @@ class DecisionTreeClassifierSuite extends MLTest with
DefaultReadWriteTest {
TreeTests.checkEqual(model, model2)
assert(model.numFeatures === model2.numFeatures)
assert(model.numClasses === model2.numClasses)
- val features = Vectors.dense(0.0, 0.0)
- assert(model.predictLeaf(features) === model2.predictLeaf(features))
}
val dt = new DecisionTreeClassifier()
diff --git
a/mllib/src/test/scala/org/apache/spark/ml/regression/DecisionTreeRegressorSuite.scala
b/mllib/src/test/scala/org/apache/spark/ml/regression/DecisionTreeRegressorSuite.scala
index 7134afee2388..679ef83ec762 100644
---
a/mllib/src/test/scala/org/apache/spark/ml/regression/DecisionTreeRegressorSuite.scala
+++
b/mllib/src/test/scala/org/apache/spark/ml/regression/DecisionTreeRegressorSuite.scala
@@ -222,8 +222,6 @@ class DecisionTreeRegressorSuite extends MLTest with
DefaultReadWriteTest {
model2: DecisionTreeRegressionModel): Unit = {
TreeTests.checkEqual(model, model2)
assert(model.numFeatures === model2.numFeatures)
- val features = Vectors.dense(0.0, 0.0)
- assert(model.predictLeaf(features) === model2.predictLeaf(features))
}
val dt = new DecisionTreeRegressor()
diff --git a/mllib/src/test/scala/org/apache/spark/ml/tree/impl/TreeTests.scala
b/mllib/src/test/scala/org/apache/spark/ml/tree/impl/TreeTests.scala
index 9ad47963aeee..afca7e993b46 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/tree/impl/TreeTests.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/tree/impl/TreeTests.scala
@@ -266,9 +266,8 @@ private[ml] object TreeTests extends SparkFunSuite {
val leaf2 = new LeafNode(0.0, Double.NaN, null)
val node1 = new InternalNode(0.0, Double.NaN, Double.NaN, leaf0, leaf1,
new ContinuousSplit(0, 0.0), null)
- Node.withLeafIndices(
- new InternalNode(0.0, Double.NaN, Double.NaN, node1, leaf2,
- new CategoricalSplit(1, Array(0.0, 2.0), 3), null))
+ new InternalNode(0.0, Double.NaN, Double.NaN, node1, leaf2,
+ new CategoricalSplit(1, Array(0.0, 2.0), 3), null)
}
/**
@@ -295,9 +294,8 @@ private[ml] object TreeTests extends SparkFunSuite {
val leaf2 = new LeafNode(0.0, Double.NaN, null)
val node1 = new InternalNode(0.0, Double.NaN, Double.NaN, leaf1, leaf2,
new CategoricalSplit(1, Array(0.0, 1.0), 3), null)
- Node.withLeafIndices(
- new InternalNode(0.0, Double.NaN, Double.NaN, leaf0, node1,
- new ContinuousSplit(2, 1.0), null))
+ new InternalNode(0.0, Double.NaN, Double.NaN, leaf0, node1,
+ new ContinuousSplit(2, 1.0), null)
}
/**
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
