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 bc9ee80d0de0 [SPARK-58856][ML] Avoid per-leaf map for tree predictions
bc9ee80d0de0 is described below
commit bc9ee80d0de08f4cca1ac8518e4b6bee64cc544d
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Fri Aug 21 12:58:58 2026 +0800
[SPARK-58856][ML] Avoid per-leaf map for tree predictions
### What changes were proposed in this pull request?
Replace the transient `Map[LeafNode, Int]` used by tree leaf prediction
with a normal integer
field on each `LeafNode`. Tree-model construction assigns the same
left-to-right DFS indices
after training and loading, so prediction only traverses the tree and reads
the reached leaf's
stored ID. The existing `NodeData` format deliberately omits this field,
while broadcast
serialization retains it. The previous `leafIndices` map was not counted in
model size
estimation; after this PR, leaf-index storage is counted as the new
`LeafNode` field.
### Why are the changes needed?
Large trees can have many leaves. When leaf prediction is used, the map
creates a hash-table
entry and boxed value for every leaf in each executor's model copy. A
single integer per leaf
keeps the lookup constant-time while substantially reducing the memory
overhead and avoids
building an executor-side lookup cache. Since the leaf-index storage is now
part of the tree
nodes, model size estimation also accounts for it more accurately.
### Does this PR introduce _any_ user-facing change?
No. Leaf IDs and their traversal order are unchanged.
### How was this patch tested?
Added persistence assertions that verify classifier and regressor leaf IDs
are unchanged after
loading. Tests have not been run yet; the relevant ML suites will be run
before this draft is
marked ready for review.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #57389 from zhengruifeng/optimize-leaf-prediction-memory-dev3.
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, 46 insertions(+), 17 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 11d43d1b31fb..78b747f65f2f 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.fromOld(oldModel.topNode, categoricalFeatures)
+ val rootNode = Node.withLeafIndices(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 4a5c3b579035..dc2377166e17 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.fromOld(oldModel.topNode, categoricalFeatures)
+ val rootNode = Node.withLeafIndices(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 0a7e37aa0c9a..72da973b5c0d 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,6 +135,25 @@ 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)
+ }
+ }
}
/**
@@ -145,7 +164,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)"
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 b3ca7f04c3de..3e4c42afa7f2 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,26 +269,31 @@ private[spark] object RandomForest extends Logging with
Serializable {
topNodes.map { rootNode =>
new DecisionTreeClassificationModel(
uid,
- rootNode.toNode(strategy.pruneTree),
+ Node.withLeafIndices(rootNode.toNode(strategy.pruneTree)),
numFeatures,
strategy.getNumClasses())
}
} else {
topNodes.map { rootNode =>
- new DecisionTreeRegressionModel(uid,
rootNode.toNode(strategy.pruneTree), numFeatures)
+ new DecisionTreeRegressionModel(
+ uid,
+ Node.withLeafIndices(rootNode.toNode(strategy.pruneTree)),
+ numFeatures)
}
}
case None =>
if (strategy.algo == OldAlgo.Classification) {
topNodes.map { rootNode =>
new DecisionTreeClassificationModel(
- rootNode.toNode(strategy.pruneTree),
+ Node.withLeafIndices(rootNode.toNode(strategy.pruneTree)),
numFeatures,
strategy.getNumClasses())
}
} else {
topNodes.map(rootNode =>
- new
DecisionTreeRegressionModel(rootNode.toNode(strategy.pruneTree), numFeatures))
+ new DecisionTreeRegressionModel(
+ Node.withLeafIndices(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 1cc64e58d980..ca923a1f751f 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,16 +101,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 = {
@@ -508,7 +506,7 @@ private[ml] object DecisionTreeModelReadWrite {
finalNodes(n.id) = node
}
// Return the root node
- finalNodes.head
+ Node.withLeafIndices(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 d298e625880f..e248d0f51b6e 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,6 +427,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()
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..7134afee2388 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()
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..9ad47963aeee 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,8 +266,9 @@ 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)
- new InternalNode(0.0, Double.NaN, Double.NaN, node1, leaf2,
- new CategoricalSplit(1, Array(0.0, 2.0), 3), null)
+ Node.withLeafIndices(
+ new InternalNode(0.0, Double.NaN, Double.NaN, node1, leaf2,
+ new CategoricalSplit(1, Array(0.0, 2.0), 3), null))
}
/**
@@ -294,8 +295,9 @@ 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)
- new InternalNode(0.0, Double.NaN, Double.NaN, leaf0, node1,
- new ContinuousSplit(2, 1.0), null)
+ Node.withLeafIndices(
+ 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]
