@@ -26,6 +26,8 @@ import org.apache.spark.ml.util.{DefaultReadWriteTest, MLTestingUtils}
2626import org .apache .spark .ml .util .TestingUtils ._
2727import org .apache .spark .mllib .util .MLlibTestSparkContext
2828import org .apache .spark .sql .{DataFrame , Row }
29+ import org .apache .spark .sql .functions ._
30+ import org .apache .spark .sql .types ._
2931
3032class BucketizerSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {
3133
@@ -162,6 +164,29 @@ class BucketizerSuite extends SparkFunSuite with MLlibTestSparkContext with Defa
162164 .setSplits(Array (0.1 , 0.8 , 0.9 ))
163165 testDefaultReadWrite(t)
164166 }
167+
168+ test(" Bucket numeric features" ) {
169+ val splits = Array (- 3.0 , 0.0 , 3.0 )
170+ val data = Array (- 2.0 , - 1.0 , 0.0 , 1.0 , 2.0 )
171+ val expectedBuckets = Array (0.0 , 0.0 , 1.0 , 1.0 , 1.0 )
172+ val dataFrame : DataFrame = data.zip(expectedBuckets).toSeq.toDF(" feature" , " expected" )
173+
174+ val bucketizer : Bucketizer = new Bucketizer ()
175+ .setInputCol(" feature" )
176+ .setOutputCol(" result" )
177+ .setSplits(splits)
178+
179+ val types = Seq (ShortType , IntegerType , LongType , FloatType , DoubleType ,
180+ ByteType , DecimalType (10 , 0 ))
181+ for (mType <- types) {
182+ val df = dataFrame.withColumn(" feature" , col(" feature" ).cast(mType))
183+ bucketizer.transform(df).select(" result" , " expected" ).collect().foreach {
184+ case Row (x : Double , y : Double ) =>
185+ assert(x === y, " The result is not correct after bucketing in type " +
186+ mType.toString + " . " + s " Expected $y but found $x. " )
187+ }
188+ }
189+ }
165190}
166191
167192private object BucketizerSuite extends SparkFunSuite {
0 commit comments