Fixed DecisionTreeSuite bug for training second level. Needed to update treePointToNodeIndex with groupShift.

Author: jkbradley

mllib/src/main/scala/org/apache/spark/mllib/tree/DecisionTree.scala

@@ -635,8 +635,8 @@ object DecisionTree extends Serializable with Logging {
         0
       } else {
         val globalNodeIndex = predictNodeIndex(nodes(0), treePoint.binnedFeatures)
-        // Get index for this level.
-        globalNodeIndex - levelOffset
+        // Get index for this (level, group).
+        globalNodeIndex - levelOffset - groupShift
       }
     }
 
@@ -660,12 +660,6 @@ object DecisionTree extends Serializable with Logging {
         numClasses * numBins * featureIndex +
         numClasses * treePoint.binnedFeatures(featureIndex) +
         treePoint.label.toInt
-      if (aggIndex >= agg.size) {
-        println(s"nodeIndex = $nodeIndex, featureIndex = $featureIndex," +
-          s" treePoint.binnedFeatures(featureIndex) = ${treePoint.binnedFeatures(featureIndex)}," +
-          s" label = ${treePoint.label.toInt}, agg.size = ${agg.size}")
-        println(s"numClasses = $numClasses, numFeatures = $numFeatures, numBins = $numBins")
-      }
       agg(aggIndex) += 1
     }
 
@@ -808,7 +802,10 @@ object DecisionTree extends Serializable with Logging {
      */
     def binSeqOp(agg: Array[Double], treePoint: TreePoint): Array[Double] = {
       val nodeIndex = treePointToNodeIndex(treePoint)
-      if (nodeIndex >= 0) { // Otherwise, example does not reach this level.
+      // If the example does not reach this level, then nodeIndex < 0.
+      // If the example reaches this level but is handled in a different group,
+      //  then either nodeIndex < 0 (previous group) or nodeIndex >= numNodes (later group).
+      if (nodeIndex >= 0 && nodeIndex < numNodes) {
         strategy.algo match {
           case Classification =>
             if (isMulticlassClassificationWithCategoricalFeatures) {
@@ -825,7 +822,6 @@ object DecisionTree extends Serializable with Logging {
     // Calculate bin aggregate length for classification or regression.
     val binAggregateLength = numNodes * getElementsPerNode(numFeatures, numBins, numClasses,
         isMulticlassClassificationWithCategoricalFeatures, strategy.algo)
-    println("binAggregateLength = " + binAggregateLength)
     logDebug("binAggregateLength = " + binAggregateLength)
 
     /**

mllib/src/test/scala/org/apache/spark/mllib/tree/DecisionTreeSuite.scala

@@ -585,9 +585,7 @@ class DecisionTreeSuite extends FunSuite with LocalSparkContext {
     assert(bestSplits(0)._2.predict === 1)
   }
 
-  // TODO: Decide about testing 2nd level
   test("second level node building with/without groups") {
-    println("START second level node building with/without groups")
     val arr = DecisionTreeSuite.generateOrderedLabeledPoints()
     assert(arr.length === 1000)
     val rdd = sc.parallelize(arr)
@@ -605,10 +603,9 @@ class DecisionTreeSuite extends FunSuite with LocalSparkContext {
     val modelOneNode = DecisionTree.train(rdd, strategyOneNode)
     val nodes: Array[Node] = new Array[Node](7)
     nodes(0) = modelOneNode.topNode
+    nodes(0).leftNode = None
+    nodes(0).rightNode = None
 
-/*    val leftFilter = Filter(new Split(0, 400, FeatureType.Continuous, List()), -1)
-    val rightFilter = Filter(new Split(0, 400, FeatureType.Continuous, List()) ,1)
-    val filters = Array[List[Filter]](List(), List(leftFilter), List(rightFilter))*/
     val parentImpurities = Array(0.5, 0.5, 0.5)
 
     // Single group second level tree construction.