update

Signed-off-by: minmingzhu <minming.zhu@intel.com>
oap-project · Apr 23, 2023 · 9182e16 · 9182e16
1 parent e74ca63
commit 9182e16
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 5 deletions.
diff --git a/examples/random-forest/regression/pom.xml b/examples/random-forest/regression/pom.xml
@@ -7,7 +7,7 @@
   <version>1.3.1</version>
   <packaging>jar</packaging>
 
-  <name>KMeansExample</name>
+  <name>RandomForestRegressorExample</name>
   <url>https://github.com/oap-project/oap-mllib.git</url>
 
   <properties>

diff --git a/mllib-dal/src/main/scala/org/apache/spark/ml/regression/spark321/RandomForestRegressor.scala b/mllib-dal/src/main/scala/org/apache/spark/ml/regression/spark321/RandomForestRegressor.scala
@@ -209,8 +209,8 @@ class RandomForestRegressor @Since("1.4.0") (@Since("1.4.0") override val uid: S
     => {
       val rootNode = TreeUtils.buildTreeDFS(nodelist, metadata)
       new DecisionTreeRegressionModel(uid,
-        rootNode.toNode(),
-        numFeatures)
+                                      rootNode.toNode(),
+                                      numFeatures)
     }
     }.toArray
   }

diff --git a/mllib-dal/src/main/scala/org/apache/spark/ml/tree/TreeUtils.scala b/mllib-dal/src/main/scala/org/apache/spark/ml/tree/TreeUtils.scala
@@ -21,7 +21,7 @@ package org.apache.spark.ml.tree
 import com.intel.oap.mllib.classification.{LearningNode => LearningNodeDAL}
 
 import org.apache.spark.ml.tree.impl.DecisionTreeMetadata
-import org.apache.spark.mllib.tree.impurity.GiniCalculator
+import org.apache.spark.mllib.tree.impurity.{GiniCalculator, VarianceCalculator}
 import org.apache.spark.mllib.tree.model.ImpurityStats
 
 object TreeUtils {
@@ -58,7 +58,12 @@ object TreeUtils {
       val ln: LearningNodeDAL = nodes.get(i)
       i += 1
 
-      val impurityCalculator = new GiniCalculator(ln.probability, ln.sampleCount)
+      val impurityCalculator = if (metadata.impurity == "gini") {
+        new GiniCalculator(ln.probability, ln.sampleCount)
+      } else {
+        new VarianceCalculator(ln.probability, ln.sampleCount)
+      }
+
       val impurityStats = new ImpurityStats(0, ln.impurity, impurityCalculator, null, null)
       val node = LearningNode.apply(0, ln.isLeaf, impurityStats)
       node.split = if (!ln.isLeaf) {