test varargs and chain model params

Author: mengxr

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

@@ -17,24 +17,28 @@
 
 package org.apache.spark.ml
 
-import org.apache.spark.ml.param.{ParamMap, Params, ParamPair}
-import org.apache.spark.sql.SchemaRDD
-
 import scala.annotation.varargs
 
+import org.apache.spark.ml.param.{ParamMap, ParamPair, Params}
+import org.apache.spark.sql.SchemaRDD
+
 /**
  * Abstract class for estimators that fits models to data.
  */
-abstract class Estimator[M <: Model] extends Identifiable with Params with PipelineStage {
+abstract class Estimator[M <: Model] extends PipelineStage with Params {
 
   /**
-   * Fits a single model to the input data with default parameters.
+   * Fits a single model to the input data with optional parameters.
    *
    * @param dataset input dataset
+   * @param paramPairs optional list of param pairs, overwrite embedded params
    * @return fitted model
    */
-  def fit(dataset: SchemaRDD): M = {
-    fit(dataset, ParamMap.empty)
+  @varargs
+  def fit(dataset: SchemaRDD, paramPairs: ParamPair[_]*): M = {
+    val map = new ParamMap()
+    paramPairs.foreach(map.put(_))
+    fit(dataset, map)
   }
 
   /**
@@ -46,25 +50,6 @@ abstract class Estimator[M <: Model] extends Identifiable with Params with Pipel
    */
   def fit(dataset: SchemaRDD, paramMap: ParamMap): M
 
-  /**
-   * Fits a single model to the input data with provided parameters.
-   *
-   * @param dataset input dataset
-   * @param firstParamPair first parameter
-   * @param otherParamPairs other parameters
-   * @return fitted model
-   */
-  @varargs
-  def fit[T](
-      dataset: SchemaRDD,
-      firstParamPair: ParamPair[_],
-      otherParamPairs: ParamPair[_]*): M = {
-    val map = new ParamMap()
-    map.put(firstParamPair)
-    otherParamPairs.foreach(map.put(_))
-    fit(dataset, map)
-  }
-
   /**
    * Fits multiple models to the input data with multiple sets of parameters.
    * The default implementation uses a for loop on each parameter map.
@@ -74,7 +59,7 @@ abstract class Estimator[M <: Model] extends Identifiable with Params with Pipel
    * @param paramMaps an array of parameter maps
    * @return fitted models, matching the input parameter maps
    */
-  def fit(dataset: SchemaRDD, paramMaps: Array[ParamMap]): Seq[M] = {
+  def fit(dataset: SchemaRDD, paramMaps: Array[ParamMap]): Seq[M] = { // how to return an array?
     paramMaps.map(fit(dataset, _))
   }
 

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

@@ -17,20 +17,24 @@
 
 package org.apache.spark.ml
 
-import org.apache.spark.ml.param.{ParamMap, Params, ParamPair}
+import scala.annotation.varargs
+
+import org.apache.spark.ml.param.{ParamMap, ParamPair, Params}
 import org.apache.spark.sql.SchemaRDD
 
 /**
  * Abstract class for transformers that transform one dataset into another.
  */
-abstract class Transformer extends Identifiable with Params with PipelineStage {
+abstract class Transformer extends PipelineStage with Params {
 
   /**
-   * Transforms the dataset with the default parameters.
+   * Transforms the dataset with optional parameters
    * @param dataset input dataset
+   * @param paramPairs optional list of param pairs, overwrite embedded params
    * @return transformed dataset
    */
-  def transform(dataset: SchemaRDD): SchemaRDD = {
+  @varargs
+  def transform(dataset: SchemaRDD, paramPairs: ParamPair[_]*): SchemaRDD = {
     transform(dataset, ParamMap.empty)
   }
 
@@ -41,31 +45,4 @@ abstract class Transformer extends Identifiable with Params with PipelineStage {
    * @return transformed dataset
    */
   def transform(dataset: SchemaRDD, paramMap: ParamMap): SchemaRDD
-
-  /**
-   * Transforms the dataset with provided parameter pairs.
-   * @param dataset input dataset
-   * @param firstParamPair first parameter pair
-   * @param otherParamPairs second parameter pair
-   * @return transformed dataset
-   */
-  def transform(
-      dataset: SchemaRDD,
-      firstParamPair: ParamPair[_],
-      otherParamPairs: ParamPair[_]*): SchemaRDD = {
-    val map = new ParamMap()
-    map.put(firstParamPair)
-    otherParamPairs.foreach(map.put(_))
-    transform(dataset, map)
-  }
-
-  /**
-   * Transforms the dataset with multiple sets of parameters.
-   * @param dataset input dataset
-   * @param paramMaps an array of parameter maps
-   * @return transformed dataset
-   */
-  def transform(dataset: SchemaRDD, paramMaps: Array[ParamMap]): Array[SchemaRDD] = {
-    paramMaps.map(transform(dataset, _))
-  }
 }

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

@@ -64,15 +64,20 @@ public void logisticRegression() {
       .setMaxIter(10)
       .setRegParam(1.0);
     lr.model().setThreshold(0.8);
-    // In Java we can access baseSchemaRDD, while in Scala we cannot.
-    LogisticRegressionModel model = lr.fit(dataset.baseSchemaRDD());
-    model.transform(dataset.baseSchemaRDD()).registerTempTable("prediction");
+    LogisticRegressionModel model = lr.fit(dataset.schemaRDD());
+    model.transform(dataset.schemaRDD()).registerTempTable("prediction");
     JavaSchemaRDD predictions = jsql.sql("SELECT label, score, prediction FROM prediction");
     for (Row r: predictions.collect()) {
       System.out.println(r);
     }
   }
 
+  @Test
+  public void logisticRegressionFitWithVarargs() {
+    LogisticRegression lr = new LogisticRegression();
+    lr.fit(dataset.schemaRDD(), lr.maxIter().w(10), lr.regParam().w(1.0));
+  }
+
   @Test
   public void logisticRegressionWithCrossValidation() {
     LogisticRegression lr = new LogisticRegression();

mllib/src/test/scala/org/apache/spark/ml/example/LogisticRegressionSuite.scala

@@ -35,11 +35,16 @@ class LogisticRegressionSuite extends FunSuite {
       .setMaxIter(10)
       .setRegParam(1.0)
     val model = lr.fit(dataset)
-    model.transform(dataset, model.threshold -> 0.8) // overwrite threshold
+    model.transform(dataset, lr.model.threshold -> 0.8) // overwrite threshold
       .select('label, 'score, 'prediction).collect()
       .foreach(println)
   }
 
+  test("logistic regression fit with varargs") {
+    val lr = new LogisticRegression
+    lr.fit(dataset, lr.maxIter -> 10, lr.regParam -> 1.0)
+  }
+
   test("logistic regression with cross validation") {
     val lr = new LogisticRegression
     val lrParamMaps = new ParamGridBuilder()