Skip to content
Snippets Groups Projects
Commit 68abc1b4 authored by Yuhao Yang's avatar Yuhao Yang Committed by Sean Owen
Browse files

[SPARK-14814][MLLIB] API: Java compatibility, docs

## What changes were proposed in this pull request?
jira: https://issues.apache.org/jira/browse/SPARK-14814
fix a java compatibility function in mllib DecisionTreeModel. As synced in jira, other compatibility issues don't need fixes.

## How was this patch tested?

existing ut

Author: Yuhao Yang <hhbyyh@gmail.com>

Closes #12971 from hhbyyh/javacompatibility.
parent 635ef407
No related branches found
No related tags found
No related merge requests found
......@@ -75,8 +75,8 @@ class DecisionTreeModel @Since("1.0.0") (
* @return JavaRDD of predictions for each of the given data points
*/
@Since("1.2.0")
def predict(features: JavaRDD[Vector]): JavaRDD[Double] = {
predict(features.rdd)
def predict(features: JavaRDD[Vector]): JavaRDD[java.lang.Double] = {
predict(features.rdd).toJavaRDD().asInstanceOf[JavaRDD[java.lang.Double]]
}
/**
......
......@@ -28,6 +28,8 @@ import org.junit.Test;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.mllib.linalg.Vector;
import org.apache.spark.mllib.regression.LabeledPoint;
import org.apache.spark.mllib.tree.configuration.Algo;
import org.apache.spark.mllib.tree.configuration.Strategy;
......@@ -95,6 +97,14 @@ public class JavaDecisionTreeSuite implements Serializable {
DecisionTreeModel model = DecisionTree$.MODULE$.train(rdd.rdd(), strategy);
// java compatibility test
JavaRDD<Double> predictions = model.predict(rdd.map(new Function<LabeledPoint, Vector>() {
@Override
public Vector call(LabeledPoint v1) {
return v1.features();
}
}));
int numCorrect = validatePrediction(arr, model);
Assert.assertTrue(numCorrect == rdd.count());
}
......
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