remove noSplit and set Predict private to tree

apache · Sep 10, 2014 · 0278a11 · 0278a11
1 parent d593ec7
commit 0278a11
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 31 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/tree/DecisionTree.scala b/mllib/src/main/scala/org/apache/spark/mllib/tree/DecisionTree.scala
@@ -835,11 +835,7 @@ object DecisionTree extends Serializable with Logging {
               calculateGainForSplit(leftChildStats, rightChildStats, nodeImpurity, level, metadata)
             (splitIdx, gainStats)
           }.maxBy(_._2.gain)
-        if (bestFeatureGainStats == InformationGainStats.invalidInformationGainStats) {
-          (Split.noSplit, InformationGainStats.invalidInformationGainStats)
-        } else {
-          (splits(featureIndex)(bestFeatureSplitIndex), bestFeatureGainStats)
-        }
+        (splits(featureIndex)(bestFeatureSplitIndex), bestFeatureGainStats)
       } else if (metadata.isUnordered(featureIndex)) {
         // Unordered categorical feature
         val (leftChildOffset, rightChildOffset) =
@@ -853,11 +849,7 @@ object DecisionTree extends Serializable with Logging {
               calculateGainForSplit(leftChildStats, rightChildStats, nodeImpurity, level, metadata)
             (splitIndex, gainStats)
           }.maxBy(_._2.gain)
-        if (bestFeatureGainStats == InformationGainStats.invalidInformationGainStats) {
-          (Split.noSplit, InformationGainStats.invalidInformationGainStats)
-        } else {
-          (splits(featureIndex)(bestFeatureSplitIndex), bestFeatureGainStats)
-        }
+        (splits(featureIndex)(bestFeatureSplitIndex), bestFeatureGainStats)
       } else {
         // Ordered categorical feature
         val nodeFeatureOffset = binAggregates.getNodeFeatureOffset(nodeIndex, featureIndex)
@@ -929,15 +921,11 @@ object DecisionTree extends Serializable with Logging {
               calculateGainForSplit(leftChildStats, rightChildStats, nodeImpurity, level, metadata)
             (splitIndex, gainStats)
           }.maxBy(_._2.gain)
-        if (bestFeatureGainStats == InformationGainStats.invalidInformationGainStats) {
-          (Split.noSplit, InformationGainStats.invalidInformationGainStats)
-        } else {
-          val categoriesForSplit =
-            categoriesSortedByCentroid.map(_._1.toDouble).slice(0, bestFeatureSplitIndex + 1)
-          val bestFeatureSplit =
-            new Split(featureIndex, Double.MinValue, Categorical, categoriesForSplit)
-          (bestFeatureSplit, bestFeatureGainStats)
-        }
+        val categoriesForSplit =
+          categoriesSortedByCentroid.map(_._1.toDouble).slice(0, bestFeatureSplitIndex + 1)
+        val bestFeatureSplit =
+          new Split(featureIndex, Double.MinValue, Categorical, categoriesForSplit)
+        (bestFeatureSplit, bestFeatureGainStats)
       }
     }.maxBy(_._2.gain)
 

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/tree/model/Predict.scala b/mllib/src/main/scala/org/apache/spark/mllib/tree/model/Predict.scala
@@ -26,7 +26,7 @@ import org.apache.spark.annotation.DeveloperApi
  * @param prob probability of the label (classification only)
  */
 @DeveloperApi
-class Predict(
+private[tree] class Predict(
     val predict: Double,
     val prob: Double = 0.0) extends Serializable{
 

diff --git a/mllib/src/main/scala/org/apache/spark/mllib/tree/model/Split.scala b/mllib/src/main/scala/org/apache/spark/mllib/tree/model/Split.scala
@@ -68,11 +68,3 @@ private[tree] class DummyHighSplit(feature: Int, featureType: FeatureType)
 private[tree] class DummyCategoricalSplit(feature: Int, featureType: FeatureType)
   extends Split(feature, Double.MaxValue, featureType, List())
 
-private[tree] object Split {
-  /**
-   * A [[org.apache.spark.mllib.tree.model.Split]] object to denote that
-   * we can't find a valid split that satisfies minimum info gain
-   * or minimum number of instances per node.
-   */
-  val noSplit = new Split(-1, Double.MinValue, FeatureType.Continuous, List())
-}
diff --git a/mllib/src/test/scala/org/apache/spark/mllib/tree/DecisionTreeSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/tree/DecisionTreeSuite.scala
@@ -28,7 +28,7 @@ import org.apache.spark.mllib.tree.configuration.FeatureType._
 import org.apache.spark.mllib.tree.configuration.Strategy
 import org.apache.spark.mllib.tree.impl.{DecisionTreeMetadata, TreePoint}
 import org.apache.spark.mllib.tree.impurity.{Entropy, Gini, Variance}
-import org.apache.spark.mllib.tree.model.{Split, DecisionTreeModel, Node}
+import org.apache.spark.mllib.tree.model.{InformationGainStats, DecisionTreeModel, Node}
 import org.apache.spark.mllib.util.LocalSparkContext
 
 class DecisionTreeSuite extends FunSuite with LocalSparkContext {
@@ -689,11 +689,21 @@ class DecisionTreeSuite extends FunSuite with LocalSparkContext {
     val model = DecisionTree.train(input, strategy)
     assert(model.topNode.isLeaf)
     assert(model.topNode.predict == 0.0)
-    assert(model.topNode.split.get == Split.noSplit)
     val predicts = input.map(p => model.predict(p.features)).collect()
     predicts.foreach { predict =>
       assert(predict == 0.0)
     }
+
+    // test for findBestSplits when no valid split can be found
+    val metadata = DecisionTreeMetadata.buildMetadata(input, strategy)
+    val (splits, bins) = DecisionTree.findSplitsBins(input, metadata)
+    val treeInput = TreePoint.convertToTreeRDD(input, bins, metadata)
+    val bestSplits = DecisionTree.findBestSplits(treeInput, new Array(8), metadata, 0,
+      new Array[Node](0), splits, bins, 10)
+
+    assert(bestSplits.length === 1)
+    val bestInfoStats = bestSplits(0)._2
+    assert(bestInfoStats == InformationGainStats.invalidInformationGainStats)
   }
 
   test("split must satisfy min info gain requirements") {
@@ -709,11 +719,21 @@ class DecisionTreeSuite extends FunSuite with LocalSparkContext {
     val model = DecisionTree.train(input, strategy)
     assert(model.topNode.isLeaf)
     assert(model.topNode.predict == 0.0)
-    assert(model.topNode.split.get == Split.noSplit)
     val predicts = input.map(p => model.predict(p.features)).collect()
     predicts.foreach { predict =>
       assert(predict == 0.0)
     }
+
+    // test for findBestSplits when no valid split can be found
+    val metadata = DecisionTreeMetadata.buildMetadata(input, strategy)
+    val (splits, bins) = DecisionTree.findSplitsBins(input, metadata)
+    val treeInput = TreePoint.convertToTreeRDD(input, bins, metadata)
+    val bestSplits = DecisionTree.findBestSplits(treeInput, new Array(8), metadata, 0,
+      new Array[Node](0), splits, bins, 10)
+
+    assert(bestSplits.length === 1)
+    val bestInfoStats = bestSplits(0)._2
+    assert(bestInfoStats == InformationGainStats.invalidInformationGainStats)
   }
 }