[GitHub] [spark] srowen commented on a change in pull request #25485: [SPARK-28434][TESTS][ML] Fix values in dummy tree in DecisionTreeSuite

[GitBox <gi...@apache.org>]

srowen commented on a change in pull request #25485: [SPARK-28434][TESTS][ML] Fix values in dummy tree in DecisionTreeSuite
URL: https://github.com/apache/spark/pull/25485#discussion_r314963618
 
 

 ##########
 File path: mllib/src/test/scala/org/apache/spark/mllib/tree/DecisionTreeSuite.scala
 ##########
 @@ -561,34 +561,29 @@ object DecisionTreeSuite extends SparkFunSuite {
    * Create an internal node with the given node ID and feature type.
    * Note: This does NOT set the child nodes.
    */
-  private def createInternalNode(id: Int, featureType: FeatureType): Node = {
+  private def createInternalNode(
+      id: Int, featureType: FeatureType, left: Node, right: Node): Node = {
     val node = Node(nodeIndex = id, new Predict(0.0, 1.0), impurity = 0.5, isLeaf = false)
-    featureType match {
-      case Continuous =>
-        node.split = Some(new Split(feature = 0, threshold = 0.5, Continuous,
-          categories = List.empty[Double]))
-      case Categorical =>
-        node.split = Some(new Split(feature = 1, threshold = 0.0, Categorical,
-          categories = List(0.0, 1.0)))
-    }
-    // TODO: The information gain stats should be consistent with info in children: SPARK-7131
-    node.stats = Some(new InformationGainStats(gain = 0.1, impurity = 0.2,
-      leftImpurity = 0.3, rightImpurity = 0.4, new Predict(1.0, 0.4), new Predict(0.0, 0.6)))
+    node.split = Some(featureType match {
+      case Continuous => Split(feature = 0, threshold = 0.5, featureType, List.empty[Double])
+      case Categorical => Split(feature = 1, threshold = 0.0, featureType, List(0.0, 1.0))
+    })
+    node.stats = Some(new InformationGainStats(gain = 0.1, impurity = node.impurity,
 
 Review comment:
   The changes are really here; above is mostly code cleanup

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services

---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org