add java-friendly versions of fit and tranform

Author: mengxr

mllib/src/main/scala/org/apache/spark/ml/Estimator.scala

@@ -18,9 +18,11 @@
 package org.apache.spark.ml
 
 import scala.annotation.varargs
+import scala.collection.JavaConverters._
 
 import org.apache.spark.ml.param.{ParamMap, ParamPair, Params}
 import org.apache.spark.sql.SchemaRDD
+import org.apache.spark.sql.api.java.JavaSchemaRDD
 
 /**
  * Abstract class for estimators that fits models to data.
@@ -63,6 +65,21 @@ abstract class Estimator[M <: Model] extends PipelineStage with Params {
     paramMaps.map(fit(dataset, _))
   }
 
+  // Java-friendly versions of fit.
+
+  @varargs
+  def fit(dataset: JavaSchemaRDD, paramPairs: ParamPair[_]*): M = {
+    fit(dataset.schemaRDD, paramPairs: _*)
+  }
+
+  def fit(dataset: JavaSchemaRDD, paramMap: ParamMap): M = {
+    fit(dataset.schemaRDD, paramMap)
+  }
+
+  def fit(dataset: JavaSchemaRDD, paramMaps: Array[ParamMap]): java.util.List[M] = {
+    fit(dataset.schemaRDD, paramMaps).asJava
+  }
+
   /**
    * Parameters for the output model.
    */

mllib/src/main/scala/org/apache/spark/ml/Transformer.scala

@@ -21,6 +21,7 @@ import scala.annotation.varargs
 
 import org.apache.spark.ml.param.{ParamMap, ParamPair, Params}
 import org.apache.spark.sql.SchemaRDD
+import org.apache.spark.sql.api.java.JavaSchemaRDD
 
 /**
  * Abstract class for transformers that transform one dataset into another.
@@ -47,4 +48,15 @@ abstract class Transformer extends PipelineStage with Params {
    * @return transformed dataset
    */
   def transform(dataset: SchemaRDD, paramMap: ParamMap): SchemaRDD
+
+  // Java-friendly versions of transform.
+
+  @varargs
+  def transform(dataset: JavaSchemaRDD, paramPairs: ParamPair[_]*): JavaSchemaRDD = {
+    transform(dataset.schemaRDD, paramPairs: _*).toJavaSchemaRDD
+  }
+
+  def transform(dataset: JavaSchemaRDD, paramMap: ParamMap): JavaSchemaRDD = {
+    transform(dataset.schemaRDD, paramMap).toJavaSchemaRDD
+  }
 }

mllib/src/test/java/org/apache/spark/ml/example/JavaLogisticRegressionSuite.java

@@ -19,6 +19,10 @@
 
 import java.io.Serializable;
 
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
 import org.apache.spark.api.java.JavaRDD;
 import org.apache.spark.api.java.JavaSparkContext;
 import org.apache.spark.ml.Pipeline;
@@ -32,10 +36,6 @@
 import org.apache.spark.sql.api.java.JavaSchemaRDD;
 import org.apache.spark.sql.api.java.Row;
 
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Test;
-
 public class JavaLogisticRegressionSuite implements Serializable {
 
   private transient JavaSparkContext jsc;
@@ -61,8 +61,8 @@ public void tearDown() {
   @Test
   public void logisticRegression() {
     LogisticRegression lr = new LogisticRegression();
-    LogisticRegressionModel model = lr.fit(dataset.schemaRDD());
-    model.transform(dataset.schemaRDD()).registerTempTable("prediction");
+    LogisticRegressionModel model = lr.fit(dataset);
+    model.transform(dataset).registerTempTable("prediction");
     JavaSchemaRDD predictions = jsql.sql("SELECT label, score, prediction FROM prediction");
     for (Row r: predictions.collect()) {
       System.out.println(r);
@@ -74,8 +74,8 @@ public void logisticRegressionWithSetters() {
     LogisticRegression lr = new LogisticRegression()
       .setMaxIter(10)
       .setRegParam(1.0);
-    LogisticRegressionModel model = lr.fit(dataset.schemaRDD());
-    model.transform(dataset.schemaRDD(), model.threshold().w(0.8)) // overwrite threshold
+    LogisticRegressionModel model = lr.fit(dataset);
+    model.transform(dataset, model.threshold().w(0.8)) // overwrite threshold
       .registerTempTable("prediction");
     JavaSchemaRDD predictions = jsql.sql("SELECT label, score, prediction FROM prediction");
     for (Row r: predictions.collect()) {
@@ -95,7 +95,7 @@ public void chainModelParams() {
   @Test
   public void logisticRegressionFitWithVarargs() {
     LogisticRegression lr = new LogisticRegression();
-    lr.fit(dataset.schemaRDD(), lr.maxIter().w(10), lr.regParam().w(1.0));
+    lr.fit(dataset, lr.maxIter().w(10), lr.regParam().w(1.0));
   }
 
   @Test
@@ -111,7 +111,7 @@ public void logisticRegressionWithCrossValidation() {
       .setEstimatorParamMaps(lrParamMaps)
       .setEvaluator(eval)
       .setNumFolds(3);
-    CrossValidatorModel bestModel = cv.fit(dataset.baseSchemaRDD());
+    CrossValidatorModel bestModel = cv.fit(dataset);
   }
 
   @Test
@@ -123,8 +123,8 @@ public void logisticRegressionWithPipeline() {
       .setFeaturesCol("scaledFeatures");
     Pipeline pipeline = new Pipeline()
       .setStages(new PipelineStage[] {scaler, lr});
-    PipelineModel model = pipeline.fit(dataset.baseSchemaRDD());
-    model.transform(dataset.baseSchemaRDD()).registerTempTable("prediction");
+    PipelineModel model = pipeline.fit(dataset);
+    model.transform(dataset).registerTempTable("prediction");
     JavaSchemaRDD predictions = jsql.sql("SELECT label, score, prediction FROM prediction");
     for (Row r: predictions.collect()) {
       System.out.println(r);