Github user srowen commented on a diff in the pull request: https://github.com/apache/spark/pull/13440#discussion_r218208123 --- Diff: mllib/src/main/scala/org/apache/spark/ml/tree/impl/RandomForest.scala --- @@ -670,14 +670,32 @@ private[spark] object RandomForest extends Logging { val leftImpurity = leftImpurityCalculator.calculate() // Note: This equals 0 if count = 0 val rightImpurity = rightImpurityCalculator.calculate() - val leftWeight = leftCount / totalCount.toDouble - val rightWeight = rightCount / totalCount.toDouble + val gain = metadata.impurity match { + case imp if (imp.isTestStatistic) => + // For split quality measures based on a test-statistic, run the test on the + // left and right sub-populations to get a p-value for the null hypothesis + val pval = imp.calculate(leftImpurityCalculator, rightImpurityCalculator) + // Transform the test statistic p-val into a larger-is-better gain value + Impurity.pValToGain(pval) + + case _ => + // Default purity-gain logic: + // measure the weighted decrease in impurity from parent to the left and right + val leftWeight = leftCount / totalCount.toDouble + val rightWeight = rightCount / totalCount.toDouble + + impurity - leftWeight * leftImpurity - rightWeight * rightImpurity + } - val gain = impurity - leftWeight * leftImpurity - rightWeight * rightImpurity + // If the impurity being used is a test statistic p-val, apply a standard transform into + // a larger-is-better gain value for the minimum-gain threshold + val minGain = + if (metadata.impurity.isTestStatistic) Impurity.pValToGain(metadata.minInfoGain) + else metadata.minInfoGain --- End diff -- Kind of a design question here... right now the caller has to switch logic based on what's inside metadata. Can methods like metadata.minInfoGain just implement different logic when the impurity is a test statistic, and so on? push this down towards the impurity implementation? I wonder if `isTestStatistic` can go away with the right API, but I am not familiar with the details of what that requires.
--- --------------------------------------------------------------------- To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org