Github user jkbradley commented on a diff in the pull request:
https://github.com/apache/spark/pull/15428#discussion_r84563124
--- Diff:
mllib/src/test/scala/org/apache/spark/ml/feature/QuantileDiscretizerSuite.scala
---
@@ -82,14 +82,23 @@ class QuantileDiscretizerSuite
.setInputCol("input")
.setOutputCol("result")
.setNumBuckets(numBuckets)
-
- // Reserve extra one bucket for NaN
- val expectedNumBuckets = discretizer.fit(df).getSplits.length - 1
- val result = discretizer.fit(df).transform(df)
- val observedNumBuckets = result.select("result").distinct.count
- assert(observedNumBuckets == expectedNumBuckets,
- s"Observed number of buckets are not correct." +
- s" Expected $expectedNumBuckets but found $observedNumBuckets")
+ val splits = Array(Double.NegativeInfinity, 1.0,
Double.PositiveInfinity)
+
+ for (handleNaN <- Array("keep", "skip")) {
+ discretizer.setHandleNaN(handleNaN)
+ val expectedNumBuckets = {
--- End diff --
This is kind of complex logic just to get the expected number of splits.
I'd recommend just putting the expected bucket *values* in the original
DataFrame:
```
val df = sc.parallelize(Seq(
(1.0, /*expected value for option "keep"*/, /*expected value for option
"skip"*/),
...
)).toDF("input", "expectedKeep", "expectedSkip")
```
Then you can compare with the actual values. That'll be an easier test to
understand IMO.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastruct...@apache.org or file a JIRA ticket
with INFRA.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
