Skip to content

[SPARK-20574][ML] Allow Bucketizer to handle non-Double numeric column #17840

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 4 commits into from
Closed

[SPARK-20574][ML] Allow Bucketizer to handle non-Double numeric column #17840

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
7d93d5a
allow bucketizer to work for non-double column
May 3, 2017
a86fbde
update test for non-Double types
May 3, 2017
b4a5b6f
update test
May 3, 2017
d149350
include all numeric types
May 4, 2017
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions mllib/src/main/scala/org/apache/spark/ml/feature/Bucketizer.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,7 @@ final class Bucketizer @Since("1.4.0") (@Since("1.4.0") override val uid: String
Bucketizer.binarySearchForBuckets($(splits), feature, keepInvalid)
}

val newCol = bucketizer(filteredDataset($(inputCol)))
val newCol = bucketizer(filteredDataset($(inputCol)).cast(DoubleType))
val newField = prepOutputField(filteredDataset.schema)
filteredDataset.withColumn($(outputCol), newCol, newField.metadata)
}
Expand All @@ -130,7 +130,7 @@ final class Bucketizer @Since("1.4.0") (@Since("1.4.0") override val uid: String

@Since("1.4.0")
override def transformSchema(schema: StructType): StructType = {
SchemaUtils.checkColumnType(schema, $(inputCol), DoubleType)
SchemaUtils.checkNumericType(schema, $(inputCol))
SchemaUtils.appendColumn(schema, prepOutputField(schema))
}

Expand Down
25 changes: 25 additions & 0 deletions mllib/src/test/scala/org/apache/spark/ml/feature/BucketizerSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@ import org.apache.spark.ml.util.{DefaultReadWriteTest, MLTestingUtils}
import org.apache.spark.ml.util.TestingUtils._
import org.apache.spark.mllib.util.MLlibTestSparkContext
import org.apache.spark.sql.{DataFrame, Row}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.types._

class BucketizerSuite extends SparkFunSuite with MLlibTestSparkContext with DefaultReadWriteTest {

Expand Down Expand Up @@ -162,6 +164,29 @@ class BucketizerSuite extends SparkFunSuite with MLlibTestSparkContext with Defa
.setSplits(Array(0.1, 0.8, 0.9))
testDefaultReadWrite(t)
}

test("Bucket numeric features") {
val splits = Array(-3.0, 0.0, 3.0)
val data = Array(-2.0, -1.0, 0.0, 1.0, 2.0)
val expectedBuckets = Array(0.0, 0.0, 1.0, 1.0, 1.0)
val dataFrame: DataFrame = data.zip(expectedBuckets).toSeq.toDF("feature", "expected")

val bucketizer: Bucketizer = new Bucketizer()
.setInputCol("feature")
.setOutputCol("result")
.setSplits(splits)

val types = Seq(ShortType, IntegerType, LongType, FloatType, DoubleType,
ByteType, DecimalType(10, 0))
for (mType <- types) {
val df = dataFrame.withColumn("feature", col("feature").cast(mType))
bucketizer.transform(df).select("result", "expected").collect().foreach {
case Row(x: Double, y: Double) =>
assert(x === y, "The result is not correct after bucketing in type " +
mType.toString + ". " + s"Expected $y but found $x.")
}
}
}
}

private object BucketizerSuite extends SparkFunSuite {
Expand Down