Github user mengxr commented on a diff in the pull request:
https://github.com/apache/spark/pull/475#discussion_r12311733
--- Diff:
mllib/src/main/scala/org/apache/spark/mllib/tree/DecisionTree.scala ---
@@ -269,7 +296,50 @@ object DecisionTree extends Serializable with Logging {
level: Int,
filters: Array[List[Filter]],
splits: Array[Array[Split]],
- bins: Array[Array[Bin]]): Array[(Split, InformationGainStats)] = {
+ bins: Array[Array[Bin]],
+ maxLevelForSingleGroup: Int): Array[(Split, InformationGainStats)] =
{
+ // split into groups to avoid memory overflow during aggregation
+ if (level > maxLevelForSingleGroup) {
+ val numGroups = math.pow(2, (level - maxLevelForSingleGroup)).toInt
+ logDebug("numGroups = " + numGroups)
+ var groupIndex = 0
+ var bestSplits = new Array[(Split, InformationGainStats)](0)
+ while (groupIndex < numGroups) {
+ val bestSplitsForGroup = findBestSplitsPerGroup(input,
parentImpurities, strategy, level,
+ filters, splits, bins, numGroups, groupIndex)
+ bestSplits = Array.concat(bestSplits, bestSplitsForGroup)
+ groupIndex += 1
+ }
+ bestSplits
+ } else {
+ findBestSplitsPerGroup(input, parentImpurities, strategy, level,
filters, splits, bins)
+ }
+ }
+
+ /**
+ * Returns an array of optimal splits for a group of nodes at a given
level
+ *
+ * @param input RDD of
[[org.apache.spark.mllib.regression.LabeledPoint]] used as training data
+ * for DecisionTree
+ * @param parentImpurities Impurities for all parent nodes for the
current level
+ * @param strategy
[[org.apache.spark.mllib.tree.configuration.Strategy]] instance containing
+ * parameters for construction the DecisionTree
+ * @param level Level of the tree
+ * @param filters Filters for all nodes at a given level
+ * @param splits possible splits for all features
+ * @param bins possible bins for all features
+ * @return array of splits with best splits for all nodes at a given
level.
--- End diff --
Add docs for `numGroups` and `groupIndex`.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastruct...@apache.org or file a JIRA ticket
with INFRA.
---
