[SPARK-34591][ML] Add decision tree pruning as a parameter

[WeichenXu123]

What changes were proposed in this pull request?

This PR adds a parameter to enable/disable a featuer where LearningNodes are merged after a RF model is trained.

This PR takes over #32813

Why are the changes needed?

2 Reasons:

1. In addition to basic classification, another use case for decision trees are the probabilities associated with predictions.
   Once pruned, these predictions are lost and it makes the trees/predictions challenging to work with if not unusable.

2. It is not in line with the default behavior in sklearn. In sklearn, the trees are left unpruned by default.

Please see Jira ticket for more explanation.

Does this PR introduce any user-facing change?

New params:
adds a parameter pruneTree that is exposed to the Tree based classifiers. Will add tests here to ensure parameter is exposed correctly.

How was this patch tested?

Unit tests.

[Copilot]

Pull request overview

Adds a new pruneTree parameter to Spark ML tree-based classification estimators to control post-training pruning (merging redundant subtrees), wiring the flag through the Scala training pipeline and exposing it in PySpark.

Changes:

• Introduce pruneTree as a ML param (default true) and propagate it into the underlying old-API Strategy used by training.
• Expose pruneTree in PySpark DecisionTreeClassifier / RandomForestClassifier (constructor, setParams, setter/getter).
• Update Scala unit tests to toggle pruning via Strategy.pruneTree (instead of a testing-only prune argument).

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 8 comments.

Show a summary per file

File	Description
python/pyspark/ml/tree.py	Adds pruneTree Param + getter in shared tree classifier param mixin.
python/pyspark/ml/classification.py	Wires pruneTree through PySpark DT/RF classifier defaults, signatures, and setters.
mllib/src/test/scala/org/apache/spark/ml/tree/impl/RandomForestSuite.scala	Updates pruning-related tests to use Strategy.pruneTree.
mllib/src/main/scala/org/apache/spark/mllib/tree/configuration/Strategy.scala	Adds pruneTree to the old-API Strategy configuration.
mllib/src/main/scala/org/apache/spark/ml/tree/treeParams.scala	Adds the ML pruneTree param to TreeClassifierParams.
mllib/src/main/scala/org/apache/spark/ml/tree/impl/RandomForest.scala	Removes test-only prune arg; uses strategy.pruneTree for model materialization.
mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala	Exposes setPruneTree and propagates param to old Strategy.
mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala	Exposes setPruneTree and propagates param to old Strategy.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

dismiss