Github user WeichenXu123 commented on a diff in the pull request:
https://github.com/apache/spark/pull/20686#discussion_r171519138
--- Diff:
mllib/src/test/scala/org/apache/spark/ml/feature/OneHotEncoderEstimatorSuite.scala
---
@@ -103,11 +96,12 @@ class OneHotEncoderEstimatorSuite
.setInputCols(Array("size"))
.setOutputCols(Array("encoded"))
val model = encoder.fit(df)
- val output = model.transform(df)
- val group = AttributeGroup.fromStructField(output.schema("encoded"))
- assert(group.size === 2)
- assert(group.getAttr(0) ===
BinaryAttribute.defaultAttr.withName("small").withIndex(0))
- assert(group.getAttr(1) ===
BinaryAttribute.defaultAttr.withName("medium").withIndex(1))
+ testTransformerByGlobalCheckFunc[(Double)](df, model, "encoded") {
rows =>
+ val group =
AttributeGroup.fromStructField(rows.head.schema("encoded"))
+ assert(group.size === 2)
+ assert(group.getAttr(0) ===
BinaryAttribute.defaultAttr.withName("small").withIndex(0))
+ assert(group.getAttr(1) ===
BinaryAttribute.defaultAttr.withName("medium").withIndex(1))
+ }
--- End diff --
I think for streaming , we don't need to test functions about attributes,
so this part just keep old testing code.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
