Skip to content
Snippets Groups Projects
Commit 7d432af8 authored by 颜发才(Yan Facai)'s avatar 颜发才(Yan Facai) Committed by Joseph K. Bradley
Browse files

[SPARK-20043][ML] DecisionTreeModel: ImpurityCalculator builder fails for...

[SPARK-20043][ML] DecisionTreeModel: ImpurityCalculator builder fails for uppercase impurity type Gini

Fix bug: DecisionTreeModel can't recongnize Impurity "Gini" when loading

TODO:
+ [x] add unit test
+ [x] fix the bug

Author: 颜发才(Yan Facai) <facai.yan@gmail.com>

Closes #17407 from facaiy/BUG/decision_tree_loader_failer_with_Gini_impurity.
parent 92e385e0
No related branches found
No related tags found
No related merge requests found
......@@ -184,7 +184,7 @@ private[spark] object ImpurityCalculator {
* the given stats.
*/
def getCalculator(impurity: String, stats: Array[Double]): ImpurityCalculator = {
impurity match {
impurity.toLowerCase match {
case "gini" => new GiniCalculator(stats)
case "entropy" => new EntropyCalculator(stats)
case "variance" => new VarianceCalculator(stats)
......
......@@ -385,6 +385,20 @@ class DecisionTreeClassifierSuite
testEstimatorAndModelReadWrite(dt, continuousData, allParamSettings ++ Map("maxDepth" -> 0),
allParamSettings ++ Map("maxDepth" -> 0), checkModelData)
}
test("SPARK-20043: " +
"ImpurityCalculator builder fails for uppercase impurity type Gini in model read/write") {
val rdd = TreeTests.getTreeReadWriteData(sc)
val data: DataFrame =
TreeTests.setMetadata(rdd, Map.empty[Int, Int], numClasses = 2)
val dt = new DecisionTreeClassifier()
.setImpurity("Gini")
.setMaxDepth(2)
val model = dt.fit(data)
testDefaultReadWrite(model)
}
}
private[ml] object DecisionTreeClassifierSuite extends SparkFunSuite {
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment