This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 82d058857b91 [SPARK-58856][ML] Assign leaf indices during tree
construction
82d058857b91 is described below
commit 82d058857b91482e4cc50badae7bcbf2717f9a1a
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Fri Aug 28 15:37:19 2026 +0800
[SPARK-58856][ML] Assign leaf indices during tree construction
### What changes were proposed in this pull request?
#58332 reverted the leaf-index optimization from #57389. This PR
reintroduces that optimization
without rebuilding the tree after construction.
`LeafNode` stores its leaf index as an immutable integer, and each
node-construction path assigns
the index directly:
- `LearningNode.toNode` threads the next index through its left-first
recursion after accounting
for pruning.
- `Node.fromOld` threads the next index through legacy-model conversion.
- `buildTreeFromNodes` uses the persisted preorder node IDs: its reverse
traversal assigns leaf
indices from right to left, starting at `leafCount - 1`.
`predictLeaf` reads the index from the reached leaf, so the transient
per-leaf map and the
post-construction `Node.withLeafIndices` copy are both unnecessary. The
persisted `NodeData` format
remains unchanged; loading reconstructs the indices.
Test fixtures provide explicit leaf indices, tree equality checks compare
them, and legacy
conversion tests exercise `predictLeaf`.
### Why are the changes needed?
The transient `Map[LeafNode, Int]` creates a hash-table entry and boxed
integer for every leaf in
each executor's model copy when leaf prediction is used. Storing one
integer on each leaf reduces
that retained memory and makes model-size estimation account for the index
storage.
The reverted implementation assigned indices by constructing a second
complete node graph. That
added O(number of nodes in one tree) temporary driver memory during
training, loading, and legacy
conversion. Assigning indices during initial construction provides the
memory benefit without that
peak-memory regression and preserves immutable model nodes.
### Does this PR introduce _any_ user-facing change?
No. Leaf IDs and their left-to-right traversal order remain unchanged.
### How was this patch tested?
Ran:
`JAVA_HOME=/usr/lib/jvm/java-17-openjdk-amd64 ./build/sbt 'mllib/testOnly
org.apache.spark.ml.tree.impl.RandomForestSuite
org.apache.spark.ml.classification.DecisionTreeClassifierSuite
org.apache.spark.ml.regression.DecisionTreeRegressorSuite'`
All 59 tests passed. `./dev/lint-scala` also passed Scalastyle and Scalafmt
checks.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex (GPT-5)
Closes #58330 from zhengruifeng/ml-leaf-index-construction-dev-7.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit 3a7590655875254771e8611d48a679c1563c5da6)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../main/scala/org/apache/spark/ml/tree/Node.scala | 48 +++++++++++++++-------
.../org/apache/spark/ml/tree/treeModels.scala | 25 ++++-------
.../DecisionTreeClassifierSuite.scala | 33 ++++++++++++++-
.../ml/regression/DecisionTreeRegressorSuite.scala | 6 ++-
.../org/apache/spark/ml/tree/impl/TreeTests.scala | 31 ++++++++++----
5 files changed, 102 insertions(+), 41 deletions(-)
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 0a7e37aa0c9a..d664d0530a2f 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
@@ -113,21 +113,33 @@ private[ml] object Node {
* Create a new Node from the old Node format, recursively creating child
nodes as needed.
*/
def fromOld(oldNode: OldNode, categoricalFeatures: Map[Int, Int]): Node = {
+ fromOld(oldNode, categoricalFeatures, 0)._1
+ }
+
+ private def fromOld(
+ oldNode: OldNode,
+ categoricalFeatures: Map[Int, Int],
+ nextLeafIndex: Int): (Node, Int) = {
if (oldNode.isLeaf) {
// TODO: Once the implementation has been moved to this API, then
include sufficient
// statistics here.
- new LeafNode(prediction = oldNode.predict.predict,
- impurity = oldNode.impurity, impurityStats = null)
+ (new LeafNode(prediction = oldNode.predict.predict,
+ impurity = oldNode.impurity, impurityStats = null, leafIndex =
nextLeafIndex),
+ nextLeafIndex + 1)
} else {
val gain = if (oldNode.stats.nonEmpty) {
oldNode.stats.get.gain
} else {
0.0
}
- new InternalNode(prediction = oldNode.predict.predict, impurity =
oldNode.impurity,
- gain = gain, leftChild = fromOld(oldNode.leftNode.get,
categoricalFeatures),
- rightChild = fromOld(oldNode.rightNode.get, categoricalFeatures),
- split = Split.fromOld(oldNode.split.get, categoricalFeatures),
impurityStats = null)
+ val (leftChild, nextIndex) =
+ fromOld(oldNode.leftNode.get, categoricalFeatures, nextLeafIndex)
+ val (rightChild, endIndex) =
+ fromOld(oldNode.rightNode.get, categoricalFeatures, nextIndex)
+ (new InternalNode(prediction = oldNode.predict.predict, impurity =
oldNode.impurity,
+ gain = gain, leftChild = leftChild, rightChild = rightChild,
+ split = Split.fromOld(oldNode.split.get, categoricalFeatures),
impurityStats = null),
+ endIndex)
}
}
@@ -145,7 +157,8 @@ private[ml] object Node {
class LeafNode private[ml] (
override val prediction: Double,
override val impurity: Double,
- override private[ml] val impurityStats: ImpurityCalculator) extends Node {
+ override private[ml] val impurityStats: ImpurityCalculator,
+ private[tree] val leafIndex: Int = -1) extends Node {
override def toString: String =
s"LeafNode(prediction = $prediction, impurity = $impurity)"
@@ -320,24 +333,31 @@ private[tree] class LearningNode(
* Convert this [[LearningNode]] to a regular [[Node]], and recurse on any
children.
*/
def toNode(prune: Boolean = true): Node = {
+ toNode(prune, 0)._1
+ }
+ private def toNode(prune: Boolean, nextLeafIndex: Int): (Node, Int) = {
if (!leftChild.isEmpty || !rightChild.isEmpty) {
assert(leftChild.nonEmpty && rightChild.nonEmpty && split.nonEmpty &&
stats != null,
"Unknown error during Decision Tree learning. Could not convert
LearningNode to Node.")
- (leftChild.get.toNode(prune), rightChild.get.toNode(prune)) match {
+ val (leftNode, nextIndex) = leftChild.get.toNode(prune, nextLeafIndex)
+ val (rightNode, endIndex) = rightChild.get.toNode(prune, nextIndex)
+ (leftNode, rightNode) match {
case (l: LeafNode, r: LeafNode) if prune && l.prediction ==
r.prediction =>
- new LeafNode(l.prediction, stats.impurity, stats.impurityCalculator)
+ (new LeafNode(l.prediction, stats.impurity,
stats.impurityCalculator, nextLeafIndex),
+ nextLeafIndex + 1)
case (l, r) =>
- new InternalNode(stats.impurityCalculator.predict, stats.impurity,
stats.gain,
- l, r, split.get, stats.impurityCalculator)
+ (new InternalNode(stats.impurityCalculator.predict, stats.impurity,
stats.gain,
+ l, r, split.get, stats.impurityCalculator), endIndex)
}
} else {
if (stats.valid) {
- new LeafNode(stats.impurityCalculator.predict, stats.impurity,
- stats.impurityCalculator)
+ (new LeafNode(stats.impurityCalculator.predict, stats.impurity,
+ stats.impurityCalculator, nextLeafIndex), nextLeafIndex + 1)
} else {
// Here we want to keep same behavior with the old
mllib.DecisionTreeModel
- new LeafNode(stats.impurityCalculator.predict, -1.0,
stats.impurityCalculator)
+ (new LeafNode(stats.impurityCalculator.predict, -1.0,
stats.impurityCalculator,
+ nextLeafIndex), nextLeafIndex + 1)
}
}
}
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 2b275ce6fd64..5f45cc331562 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
@@ -79,18 +79,6 @@ private[spark] trait DecisionTreeModel {
/** Convert to spark.mllib DecisionTreeModel (losing some information) */
private[spark] def toOld: OldDecisionTreeModel
- /**
- * @return an iterator that traverses (DFS, left to right) the leaves
- * in the subtree of this node.
- */
- private def leafIterator(node: Node): Iterator[LeafNode] = {
- node match {
- case l: LeafNode => Iterator.single(l)
- case n: InternalNode =>
- leafIterator(n.leftChild) ++ leafIterator(n.rightChild)
- }
- }
-
private[ml] def treeStats: NodeStats
private[ml] def numLeaves: Int = treeStats.numLeaves
@@ -101,16 +89,14 @@ 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 = {
- leafIndices(rootNode.predictImpl(features)).toDouble
+ val leaf = rootNode.predictImpl(features)
+ assert(leaf.leafIndex >= 0, "Leaf indices are not assigned.")
+ leaf.leafIndex.toDouble
}
def getEstimatedSize(): Long = {
@@ -494,6 +480,7 @@ private[ml] object DecisionTreeModelReadWrite {
// We fill `finalNodes` in reverse order. Since node IDs are assigned via
a pre-order
// traversal, this guarantees that child nodes will be built before parent
nodes.
val finalNodes = new Array[Node](nodes.length)
+ var leafIndex = nodes.count(_.leftChild == -1) - 1
nodes.reverseIterator.foreach { case n: NodeData =>
val impurityStats =
ImpurityCalculator.getCalculator(impurityType, n.impurityStats,
n.rawCount)
@@ -503,7 +490,9 @@ private[ml] object DecisionTreeModelReadWrite {
new InternalNode(n.prediction, n.impurity, n.gain, leftChild,
rightChild,
n.split.getSplit, impurityStats)
} else {
- new LeafNode(n.prediction, n.impurity, impurityStats)
+ val leaf = new LeafNode(n.prediction, n.impurity, impurityStats,
leafIndex)
+ leafIndex -= 1
+ leaf
}
finalNodes(n.id) = node
}
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 d298e625880f..db498cf2d6dc 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
@@ -359,6 +359,33 @@ class DecisionTreeClassifierSuite extends MLTest with
DefaultReadWriteTest {
}
}
+ test("leaf indices are correct after training, loading, and conversion from
old API") {
+ val rdd = TreeTests.getTreeReadWriteData(sc)
+ val categoricalFeatures = Map(0 -> 2, 1 -> 3)
+ val numClasses = 2
+ val data = TreeTests.setMetadata(rdd, categoricalFeatures, numClasses)
+ val dt = new DecisionTreeClassifier().setMaxDepth(2)
+ val trainedModel = dt.fit(data)
+
+ assert(trainedModel.numLeaves > 1)
+ TreeTests.checkLeafIndices(trainedModel.rootNode)
+
+ withTempDir { dir =>
+ val path = s"${dir.getAbsolutePath}/model"
+ trainedModel.write.save(path)
+ val loadedModel = DecisionTreeClassificationModel.load(path)
+ assert(loadedModel.numLeaves > 1)
+ TreeTests.checkLeafIndices(loadedModel.rootNode)
+ }
+
+ val oldStrategy = dt.getOldStrategy(categoricalFeatures, numClasses)
+ val oldTree = OldDecisionTree.train(rdd.map(OldLabeledPoint.fromML),
oldStrategy)
+ val convertedModel = DecisionTreeClassificationModel.fromOld(
+ oldTree, dt, categoricalFeatures)
+ assert(convertedModel.numLeaves > 1)
+ TreeTests.checkLeafIndices(convertedModel.rootNode)
+ }
+
test("should support all NumericType labels and not support other types") {
val dt = new DecisionTreeClassifier().setMaxDepth(1)
MLTestingUtils.checkNumericTypes[DecisionTreeClassificationModel,
DecisionTreeClassifier](
@@ -427,6 +454,8 @@ 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()
@@ -488,7 +517,8 @@ private[ml] object DecisionTreeClassifierSuite extends
SparkFunSuite {
dt: DecisionTreeClassifier,
categoricalFeatures: Map[Int, Int],
numClasses: Int): Unit = {
- val numFeatures = data.first().features.size
+ val vec = data.first().features
+ val numFeatures = vec.size
val oldStrategy = dt.getOldStrategy(categoricalFeatures, numClasses)
val oldTree = OldDecisionTree.train(data.map(OldLabeledPoint.fromML),
oldStrategy)
val newData: DataFrame = TreeTests.setMetadata(data, categoricalFeatures,
numClasses)
@@ -497,6 +527,7 @@ private[ml] object DecisionTreeClassifierSuite extends
SparkFunSuite {
val oldTreeAsNew = DecisionTreeClassificationModel.fromOld(
oldTree, newTree.parent.asInstanceOf[DecisionTreeClassifier],
categoricalFeatures)
TreeTests.checkEqual(oldTreeAsNew, newTree)
+ assert(oldTreeAsNew.predictLeaf(vec) === newTree.predictLeaf(vec))
assert(newTree.numFeatures === numFeatures)
}
}
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 679ef83ec762..349877b6b9eb 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,6 +222,8 @@ 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()
@@ -270,7 +272,8 @@ private[ml] object DecisionTreeRegressorSuite extends
SparkFunSuite {
data: RDD[LabeledPoint],
dt: DecisionTreeRegressor,
categoricalFeatures: Map[Int, Int]): Unit = {
- val numFeatures = data.first().features.size
+ val vec = data.first().features
+ val numFeatures = vec.size
val oldStrategy = dt.getOldStrategy(categoricalFeatures)
val oldTree = OldDecisionTree.train(data.map(OldLabeledPoint.fromML),
oldStrategy)
val newData: DataFrame = TreeTests.setMetadata(data, categoricalFeatures,
numClasses = 0)
@@ -279,6 +282,7 @@ private[ml] object DecisionTreeRegressorSuite extends
SparkFunSuite {
val oldTreeAsNew = DecisionTreeRegressionModel.fromOld(
oldTree, newTree.parent.asInstanceOf[DecisionTreeRegressor],
categoricalFeatures)
TreeTests.checkEqual(oldTreeAsNew, newTree)
+ assert(oldTreeAsNew.predictLeaf(vec) === newTree.predictLeaf(vec))
assert(newTree.numFeatures === numFeatures)
}
}
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 afca7e993b46..f24631c40ba2 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
@@ -141,12 +141,29 @@ private[ml] object TreeTests extends SparkFunSuite {
assert(aye.split === bee.split)
checkEqual(aye.leftChild, bee.leftChild)
checkEqual(aye.rightChild, bee.rightChild)
- case (aye: LeafNode, bee: LeafNode) => // do nothing
+ case (aye: LeafNode, bee: LeafNode) =>
+ assert(aye.leafIndex === bee.leafIndex)
case _ =>
fail("Found mismatched nodes")
}
}
+ /** Check that leaf indices are consecutive and assigned from left to right.
*/
+ def checkLeafIndices(root: Node): Unit = {
+ var expectedLeafIndex = 0
+
+ def traverse(node: Node): Unit = node match {
+ case internal: InternalNode =>
+ traverse(internal.leftChild)
+ traverse(internal.rightChild)
+ case leaf: LeafNode =>
+ assert(leaf.leafIndex === expectedLeafIndex)
+ expectedLeafIndex += 1
+ }
+
+ traverse(root)
+ }
+
/**
* Check if the two models are exactly the same.
* If the models are not equal, this throws an exception.
@@ -261,9 +278,9 @@ private[ml] object TreeTests extends SparkFunSuite {
* / \
* leaf0 leaf1
*/
- val leaf0 = new LeafNode(0.0, Double.NaN, null)
- val leaf1 = new LeafNode(1.0, Double.NaN, null)
- val leaf2 = new LeafNode(0.0, Double.NaN, null)
+ val leaf0 = new LeafNode(0.0, Double.NaN, null, 0)
+ val leaf1 = new LeafNode(1.0, Double.NaN, null, 1)
+ val leaf2 = new LeafNode(0.0, Double.NaN, null, 2)
val node1 = new InternalNode(0.0, Double.NaN, Double.NaN, leaf0, leaf1,
new ContinuousSplit(0, 0.0), null)
new InternalNode(0.0, Double.NaN, Double.NaN, node1, leaf2,
@@ -289,9 +306,9 @@ private[ml] object TreeTests extends SparkFunSuite {
* / \
* leaf1 leaf2
*/
- val leaf0 = new LeafNode(0.0, Double.NaN, null)
- val leaf1 = new LeafNode(1.0, Double.NaN, null)
- val leaf2 = new LeafNode(0.0, Double.NaN, null)
+ val leaf0 = new LeafNode(0.0, Double.NaN, null, 0)
+ val leaf1 = new LeafNode(1.0, Double.NaN, null, 1)
+ val leaf2 = new LeafNode(0.0, Double.NaN, null, 2)
val node1 = new InternalNode(0.0, Double.NaN, Double.NaN, leaf1, leaf2,
new CategoricalSplit(1, Array(0.0, 1.0), 3), null)
new InternalNode(0.0, Double.NaN, Double.NaN, leaf0, node1,
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
