Fix test to supply cols as per new invert method

prabeesh · Jul 29, 2015 · a42d773 · a42d773
1 parent 16cc3c3
commit a42d773
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/mllib/src/test/scala/org/apache/spark/ml/feature/StringIndexerSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/feature/StringIndexerSuite.scala
@@ -48,7 +48,8 @@ class StringIndexerSuite extends SparkFunSuite with MLlibTestSparkContext {
     val expected = Set((0, 0.0), (1, 2.0), (2, 1.0), (3, 0.0), (4, 0.0), (5, 1.0))
     assert(output === expected)
     // convert reverse our transform
-    val reversed = indexer.invert().transform(transformed.select("id", "labelIndex"))
+    val reversed = indexer.invert("labelIndex", "label")
+      .transform(transformed.select("id", "labelIndex"))
       .select("id", "label")
     assert(df.collect().map(r => (r.getInt(0), r.getString(1))).toSet ===
       reversed.collect().map(r => (r.getInt(0), r.getString(1))).toSet)