Github user jkbradley commented on a diff in the pull request:
https://github.com/apache/spark/pull/20686#discussion_r173593049
--- Diff:
mllib/src/test/scala/org/apache/spark/ml/feature/StringIndexerSuite.scala ---
@@ -109,16 +111,14 @@ class StringIndexerSuite
.setInputCol("label")
.setOutputCol("labelIndex")
.fit(df)
- val transformed = indexer.transform(df)
- val attr = Attribute.fromStructField(transformed.schema("labelIndex"))
- .asInstanceOf[NominalAttribute]
- assert(attr.values.get === Array("100", "300", "200"))
- val output = transformed.select("id", "labelIndex").rdd.map { r =>
- (r.getInt(0), r.getDouble(1))
- }.collect().toSet
- // 100 -> 0, 200 -> 2, 300 -> 1
- val expected = Set((0, 0.0), (1, 2.0), (2, 1.0), (3, 0.0), (4, 0.0),
(5, 1.0))
- assert(output === expected)
+ testTransformerByGlobalCheckFunc[(Int, String)](df, indexer, "id",
"labelIndex") { rows =>
+ val attr = Attribute.fromStructField(rows.head.schema("labelIndex"))
+ .asInstanceOf[NominalAttribute]
+ assert(attr.values.get === Array("100", "300", "200"))
+ // 100 -> 0, 200 -> 2, 300 -> 1
+ val expected = Seq((0, 0.0), (1, 2.0), (2, 1.0), (3, 0.0), (4, 0.0),
(5, 1.0)).toDF()
--- End diff --
And elsewhere below
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
