fix unit tests

Author: mengxr

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

@@ -18,26 +18,17 @@
 package org.apache.spark.ml.classification;
 
 import java.io.Serializable;
+import java.util.List;
 
 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;
-import org.apache.spark.ml.PipelineModel;
-import org.apache.spark.ml.PipelineStage;
-import org.apache.spark.ml.evaluation.BinaryClassificationEvaluator;
-import org.apache.spark.ml.feature.StandardScaler;
-import org.apache.spark.ml.param.ParamMap;
-import org.apache.spark.ml.tuning.CrossValidator;
-import org.apache.spark.ml.tuning.CrossValidatorModel;
-import org.apache.spark.ml.tuning.ParamGridBuilder;
 import org.apache.spark.mllib.regression.LabeledPoint;
 import org.apache.spark.sql.api.java.JavaSQLContext;
 import org.apache.spark.sql.api.java.JavaSchemaRDD;
-import org.apache.spark.sql.api.java.Row;
 import static org.apache.spark.mllib.classification.LogisticRegressionSuite
   .generateLogisticInputAsList;
 
@@ -51,9 +42,8 @@ public class JavaLogisticRegressionSuite implements Serializable {
   public void setUp() {
     jsc = new JavaSparkContext("local", "JavaLogisticRegressionSuite");
     jsql = new JavaSQLContext(jsc);
-    JavaRDD<LabeledPoint> points =
-      jsc.parallelize(generateLogisticInputAsList(1.0, 1.0, 100, 42), 2);
-    dataset = jsql.applySchema(points, LabeledPoint.class);
+    List<LabeledPoint> points = generateLogisticInputAsList(1.0, 1.0, 100, 42);
+    dataset = jsql.applySchema(jsc.parallelize(points, 2), LabeledPoint.class);
   }
 
   @After

mllib/src/test/java/org/apache/spark/ml/tuning/JavaCrossValidatorSuite.java

@@ -17,12 +17,14 @@
 
 package org.apache.spark.ml.tuning;
 
+import java.io.Serializable;
+import java.util.List;
+
 import org.junit.After;
 import org.junit.Assert;
 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.classification.LogisticRegression;
 import org.apache.spark.ml.evaluation.BinaryClassificationEvaluator;
@@ -33,7 +35,7 @@
 import static org.apache.spark.mllib.classification.LogisticRegressionSuite
   .generateLogisticInputAsList;
 
-public class JavaCrossValidatorSuite {
+public class JavaCrossValidatorSuite implements Serializable {
 
   private transient JavaSparkContext jsc;
   private transient JavaSQLContext jsql;
@@ -43,9 +45,8 @@ public class JavaCrossValidatorSuite {
   public void setUp() {
     jsc = new JavaSparkContext("local", "JavaLogisticRegressionSuite");
     jsql = new JavaSQLContext(jsc);
-    JavaRDD<LabeledPoint> points =
-      jsc.parallelize(generateLogisticInputAsList(1.0, 1.0, 100, 42), 2);
-    dataset = jsql.applySchema(points, LabeledPoint.class);
+    List<LabeledPoint> points = generateLogisticInputAsList(1.0, 1.0, 100, 42);
+    dataset = jsql.applySchema(jsc.parallelize(points, 2), LabeledPoint.class);
   }
 
   @After

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

@@ -23,9 +23,15 @@ import org.apache.spark.sql.SchemaRDD
 import org.apache.spark.mllib.classification.LogisticRegressionSuite.generateLogisticInput
 import org.apache.spark.sql.test.TestSQLContext._
 
-class LogisticRegressionSuite extends FunSuite with BeforeAndAfterAll {
+class LogisticRegressionSuite extends FunSuite with BeforeAndAfterAll with Serializable {
 
-  var dataset: SchemaRDD = sparkContext.parallelize(generateLogisticInput(1.0, 1.0, 1000, 42), 2)
+  var dataset: SchemaRDD = _
+
+  override def beforeAll(): Unit = {
+    super.beforeAll()
+    val points = generateLogisticInput(1.0, 1.0, 100, 42)
+    dataset = sparkContext.parallelize(points, 2)
+  }
 
   test("logistic regression") {
     val lr = new LogisticRegression

mllib/src/test/scala/org/apache/spark/ml/tuning/CrossValidatorSuite.scala

@@ -17,17 +17,22 @@
 
 package org.apache.spark.ml.tuning
 
-import org.scalatest.FunSuite
+import org.scalatest.{BeforeAndAfterAll, FunSuite}
 
 import org.apache.spark.ml.classification.LogisticRegression
 import org.apache.spark.ml.evaluation.BinaryClassificationEvaluator
 import org.apache.spark.mllib.classification.LogisticRegressionSuite.generateLogisticInput
 import org.apache.spark.sql.SchemaRDD
 import org.apache.spark.sql.test.TestSQLContext._
 
-class CrossValidatorSuite extends FunSuite {
+class CrossValidatorSuite extends FunSuite with BeforeAndAfterAll with Serializable {
 
-  val dataset: SchemaRDD = sparkContext.makeRDD(generateLogisticInput(1.0, 1.0, 100, 42), 2)
+  var dataset: SchemaRDD = _
+
+  override def beforeAll(): Unit = {
+    val points = generateLogisticInput(1.0, 1.0, 100, 42)
+    dataset = sparkContext.parallelize(points, 2)
+  }
 
   test("cross validation with logistic regression") {
     val lr = new LogisticRegression