update

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

Hide file tree

Showing 2 changed files with 16 additions and 15 deletions.
diff --git a/mllib-dal/src/main/native/DecisionForestRegressorOneAPIImpl.cpp b/mllib-dal/src/main/native/DecisionForestRegressorOneAPIImpl.cpp
@@ -61,10 +61,9 @@ LearningNode convertsplitToLearningNode(
     splitNode.impurity = info.get_impurity();
     splitNode.sampleCount = info.get_sample_count();
     std::unique_ptr<double[]> arr(new double[classCount]);
-    for (std::int64_t index_class = 0; index_class < classCount;
-         ++index_class) {
-        arr[index_class] = 0.0;
-    }
+    arr[0] = info.get_sample_count();
+    arr[1] = 0.0;
+    arr[2] = 0.0;
     splitNode.probability = std::move(arr);
     return splitNode;
 }
@@ -78,10 +77,9 @@ convertleafToLearningNode(const df::leaf_node_info<df::task::regression> &info,
     leafNode.impurity = info.get_impurity();
     leafNode.sampleCount = info.get_sample_count();
     std::unique_ptr<double[]> arr(new double[classCount]);
-    for (std::int64_t index_class = 0; index_class < classCount;
-         ++index_class) {
-        arr[index_class] = 0.0;
-    }
+    arr[0] = info.get_sample_count();
+    arr[1] = info.get_sample_count() * info.get_label();
+    arr[2] = info.get_sample_count() * info.get_label() * info.get_label();
     leafNode.probability = std::move(arr);
     return leafNode;
 }
@@ -294,8 +292,10 @@ static jobject doRFRegressorOneAPICompute(
                   << result_infer.get_responses() << std::endl;
 
         // convert c++ map to java hashmap
-        collect_model(env, result_train.get_model(), 0, treeForest);
-        trees = convertRFRJavaMap(env, treeForest, 0);
+        jint statsSize = 3; // spark create VarianceCalculator needs array of
+                            // sufficient statistics
+        collect_model(env, result_train.get_model(), statsSize, treeForest);
+        trees = convertRFRJavaMap(env, treeForest, statsSize);
 
         // Get the class of the input object
         jclass clazz = env->GetObjectClass(resultObj);

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, VarianceCalculator}
+import org.apache.spark.mllib.tree.impurity.{Gini, GiniCalculator, ImpurityCalculator, Variance, VarianceCalculator}
 import org.apache.spark.mllib.tree.model.ImpurityStats
 
 object TreeUtils {
@@ -58,10 +58,11 @@ object TreeUtils {
       val ln: LearningNodeDAL = nodes.get(i)
       i += 1
 
-      val impurityCalculator = if (metadata.impurity == "gini") {
-        new GiniCalculator(ln.probability, ln.sampleCount)
-      } else {
-        new VarianceCalculator(ln.probability, ln.sampleCount)
+      val impurityCalculator: ImpurityCalculator = metadata.impurity match {
+        case Gini => new GiniCalculator(ln.probability, ln.sampleCount)
+        case Variance => new VarianceCalculator(ln.probability, ln.sampleCount)
+        case _ => throw new IllegalArgumentException(s"Bad impurity parameter: " +
+          s"${metadata.impurity}")
       }
 
       val impurityStats = new ImpurityStats(0, ln.impurity, impurityCalculator, null, null)