[SPARK-22700] Bucketizer.transform incorrectly drops row containing NaN

• Type: Bug
• Status: Resolved
• Priority: Major
• Resolution: Fixed
• Affects Version/s: 2.2.0, 2.3.0
• Fix Version/s: 2.2.2, 2.3.0
• Component/s: ML
• Assignee: Ruifeng Zheng
• Reporter: Ruifeng Zheng
• Created: 05/Dec/17 11:44
• Updated: 22/Feb/18 01:27
• Resolved: 13/Dec/17 07:10

Description

import org.apache.spark.ml.feature._

val df = spark.createDataFrame(Seq((2.3, 3.0), (Double.NaN, 3.0), (6.7, Double.NaN))).toDF("a", "b")

val splits = Array(Double.NegativeInfinity, 3.0, Double.PositiveInfinity)

val bucketizer: Bucketizer = new Bucketizer().setInputCol("a").setOutputCol("aa").setSplits(splits)

bucketizer.setHandleInvalid("skip")

scala> df.show
+---+---+
|  a|  b|
+---+---+
|2.3|3.0|
|NaN|3.0|
|6.7|NaN|
+---+---+

scala> bucketizer.transform(df).show
+---+---+---+
|  a|  b| aa|
+---+---+---+
|2.3|3.0|0.0|
+---+---+---+

When handleInvalid is set skip, the last item in input is incorrectly droped, though colum 'b' is not an input column

Issue Links

• links to: [Github] Pull Request #19894 (zhengruifeng)
• links to: [Github] Pull Request #20539 (zhengruifeng)