fix test suite

Author: yinxusen

mllib/src/test/scala/org/apache/spark/ml/feature/PolynomialMapperSuite.scala

@@ -22,28 +22,16 @@ import org.scalatest.FunSuite
 import org.apache.spark.mllib.linalg.{DenseVector, SparseVector, Vector, Vectors}
 import org.apache.spark.mllib.util.MLlibTestSparkContext
 import org.apache.spark.mllib.util.TestingUtils._
-import org.apache.spark.sql.{DataFrame, Row, SQLContext}
+import org.apache.spark.sql.{Row, SQLContext}
+import org.scalatest.exceptions.TestFailedException
 
 class PolynomialMapperSuite extends FunSuite with MLlibTestSparkContext {
 
-  def collectResult(result: DataFrame): Array[Vector] = {
-    result.select("poly_features").collect().map {
-      case Row(features: Vector) => features
-    }
-  }
-
-  def assertTypeOfVector(lhs: Array[Vector], rhs: Array[Vector]): Unit = {
-    assert((lhs, rhs).zipped.forall {
-      case (v1: DenseVector, v2: DenseVector) => true
-      case (v1: SparseVector, v2: SparseVector) => true
-      case _ => false
-    }, "The vector type should be preserved after normalization.")
-  }
+  @transient var sqlContext: SQLContext = _
 
-  def assertValues(lhs: Array[Vector], rhs: Array[Vector]): Unit = {
-    assert((lhs, rhs).zipped.forall { (vector1, vector2) =>
-      vector1 ~== vector2 absTol 1E-1
-    }, "The vector value is not correct after normalization.")
+  override def beforeAll(): Unit = {
+    super.beforeAll()
+    sqlContext = new SQLContext(sc)
   }
 
   test("Polynomial expansion with default parameter") {
@@ -55,28 +43,27 @@ class PolynomialMapperSuite extends FunSuite with MLlibTestSparkContext {
       Vectors.sparse(3, Seq())
     )
 
-    val sqlContext = new SQLContext(sc)
-    val dataFrame = sqlContext
-      .createDataFrame(sc.parallelize(data, 2).map(Tuple1.apply)).toDF("features")
-
-    val polynomialMapper = new PolynomialMapper()
-      .setInputCol("features")
-      .setOutputCol("poly_features")
-
     val twoDegreeExpansion: Array[Vector] = Array(
       Vectors.sparse(9, Array(0, 1, 3, 4, 6), Array(-2.0, 2.3, 4.0, -4.6, 5.29)),
       Vectors.dense(-2.0, 2.3, 4.0, -4.6, 5.29),
       Vectors.dense(Array.fill[Double](9)(0.0)),
       Vectors.dense(0.6, -1.1, -3.0, 0.36, -0.66, -1.8, 1.21, 3.3, 9.0),
       Vectors.sparse(9, Array.empty[Int], Array.empty[Double]))
 
-    val result = collectResult(polynomialMapper.transform(dataFrame))
-
-    println(result.mkString("\n"))
+    val df = sqlContext.createDataFrame(data.zip(twoDegreeExpansion)).toDF("features", "expected")
 
-    assertTypeOfVector(data, result)
-
-    assertValues(result, twoDegreeExpansion)
+    val polynomialMapper = new PolynomialMapper()
+      .setInputCol("features")
+      .setOutputCol("polyFeatures")
+
+    polynomialMapper.transform(df).select("polyFeatures", "expected").collect().foreach {
+      case Row(expanded: DenseVector, expected: DenseVector) =>
+        assert(expanded ~== expected absTol 1e-1)
+      case Row(expanded: SparseVector, expected: SparseVector) =>
+        assert(expanded ~== expected absTol 1e-1)
+      case _ =>
+        throw new TestFailedException("Unmatched data types after polynomial expansion", 0)
+    }
   }
 
   test("Polynomial expansion with setter") {
@@ -88,15 +75,6 @@ class PolynomialMapperSuite extends FunSuite with MLlibTestSparkContext {
       Vectors.sparse(3, Seq())
     )
 
-    val sqlContext = new SQLContext(sc)
-    val dataFrame = sqlContext
-      .createDataFrame(sc.parallelize(data, 2).map(Tuple1.apply)).toDF("features")
-
-    val polynomialMapper = new PolynomialMapper()
-      .setInputCol("features")
-      .setOutputCol("poly_features")
-      .setDegree(3)
-
     val threeDegreeExpansion: Array[Vector] = Array(
       Vectors.sparse(19, Array(0, 1, 3, 4, 6, 9, 10, 12, 15),
         Array(-2.0, 2.3, 4.0, -4.6, 5.29, -8.0, 9.2, -10.58, 12.167)),
@@ -106,11 +84,22 @@ class PolynomialMapperSuite extends FunSuite with MLlibTestSparkContext {
         1.98, 5.4, -1.33, -3.63, -9.9, -27.0),
       Vectors.sparse(19, Array.empty[Int], Array.empty[Double]))
 
-    val result = collectResult(polynomialMapper.transform(dataFrame))
 
-    assertTypeOfVector(data, result)
+    val df = sqlContext.createDataFrame(data.zip(threeDegreeExpansion)).toDF("features", "expected")
 
-    assertValues(result, threeDegreeExpansion)
+    val polynomialMapper = new PolynomialMapper()
+      .setInputCol("features")
+      .setOutputCol("polyFeatures")
+      .setDegree(3)
+
+    polynomialMapper.transform(df).select("polyFeatures", "expected").collect().foreach {
+      case Row(expanded: DenseVector, expected: DenseVector) =>
+        assert(expanded ~== expected absTol 1e-1)
+      case Row(expanded: SparseVector, expected: SparseVector) =>
+        assert(expanded ~== expected absTol 1e-1)
+      case _ =>
+        throw new TestFailedException("Unmatched data types after polynomial expansion", 0)
+    }
   }
 }