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
actuaryzhang 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 3 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,9 @@ 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 +165,28 @@ class BucketizerSuite extends SparkFunSuite with MLlibTestSparkContext with Defa
.setSplits(Array(0.1, 0.8, 0.9))
testDefaultReadWrite(t)
}

test("Bucket non-double 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)
Copy link
Contributor

@MLnick MLnick May 4, 2017

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Other tests for supporting numeric types have included DecimalType - often DecimalType(10, 0), as well as ByteType. See the various Estimator tests which use MLTestingUtils.genClassifDFWithNumericLabelCol and MLTestingUtils.genRegressionDFWithNumericLabelCol

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 feature value is not correct after bucketing in type " +
Copy link
Contributor

@yanboliang yanboliang May 4, 2017

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

feature value -> result

mType.toString + ". " + s"Expected $y but found $x.")
}
}
}
}

private object BucketizerSuite extends SparkFunSuite {
Expand Down