Skip to content

[SPARK-30082][SQL] Depend on Scala type coercion when building replace query #26738

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
johnhany97 wants to merge 13 commits into from
Closed

[SPARK-30082][SQL] Depend on Scala type coercion when building replace query #26738

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
8de494e
Do not cast NaN to an Integer, Long, Short or Byte
johnhany97 Dec 2, 2019
fad3304
Revert "Do not cast NaN to an Integer, Long, Short or Byte"
johnhany97 Dec 2, 2019
4a9573c
Take it back a step
johnhany97 Dec 2, 2019
12f219b
Merge branch 'master' of github.com:apache/spark into SPARK-30082
johnhany97 Dec 2, 2019
532f447
Add tests
johnhany97 Dec 2, 2019
1c5f1eb
Avoid repeating df used in tests
johnhany97 Dec 2, 2019
ed6f08d
Improve coding style
johnhany97 Dec 2, 2019
279a9fd
Opt for a more simple fix
johnhany97 Dec 3, 2019
6b5d26d
Adjust unit tests to reflect new behaviour
johnhany97 Dec 3, 2019
10c91d6
Adjust whitespace
johnhany97 Dec 3, 2019
1744b28
Do not change param name
johnhany97 Dec 3, 2019
1295633
Clean unit tests df generator
johnhany97 Dec 3, 2019
b3709a1
Clean tests themselves as well
johnhany97 Dec 3, 2019
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
14 changes: 13 additions & 1 deletion sql/core/src/main/scala/org/apache/spark/sql/DataFrameNaFunctions.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -456,11 +456,23 @@ final class DataFrameNaFunctions private[sql](df: DataFrame) {
val keyExpr = df.col(col.name).expr
def buildExpr(v: Any) = Cast(Literal(v), keyExpr.dataType)
val branches = replacementMap.flatMap { case (source, target) =>
Seq(buildExpr(source), buildExpr(target))
if (isNaN(source) || isNaN(target)) {
col.dataType match {
case IntegerType | LongType | ShortType | ByteType => Seq.empty
case _ => Seq(buildExpr(source), buildExpr(target))
}
} else {
Seq(buildExpr(source), buildExpr(target))
}
}.toSeq
new Column(CaseKeyWhen(keyExpr, branches :+ keyExpr)).as(col.name)
}

// Check if NaN
private[this] def isNaN(value: Any): Boolean =
(value.isInstanceOf[Double] && value.asInstanceOf[Double].isNaN) ||
(value.isInstanceOf[Float] && value.asInstanceOf[Float].isNaN)

private def convertToDouble(v: Any): Double = v match {
case v: Float => v.toDouble
case v: Double => v
Expand Down
54 changes: 54 additions & 0 deletions sql/core/src/test/scala/org/apache/spark/sql/DataFrameNaFunctionsSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,16 @@ class DataFrameNaFunctionsSuite extends QueryTest with SharedSparkSession {
).toDF("name", "age", "height")
}

def createNaNDF(): DataFrame = {
Seq[(java.lang.Integer, java.lang.Long, java.lang.Short,
java.lang.Byte, java.lang.Float, java.lang.Double)](
(1, new java.lang.Long(1), new java.lang.Short("1"),
new java.lang.Byte("1"), new java.lang.Float(1.0), 1.0),
(0, new java.lang.Long(0), new java.lang.Short("0"),
new java.lang.Byte("0"), java.lang.Float.NaN, java.lang.Double.NaN)
).toDF("int", "long", "short", "byte", "float", "double")
}

test("drop") {
val input = createDF()
val rows = input.collect()
Expand Down Expand Up @@ -404,4 +414,48 @@ class DataFrameNaFunctionsSuite extends QueryTest with SharedSparkSession {
df.na.drop("any"),
Row("5", "6", "6") :: Nil)
}

test("replace nan with float") {
checkAnswer(
createNaNDF().na.replace("*", Map(
Float.NaN -> 10f
)),
Row(1, new java.lang.Long(1), new java.lang.Short("1"),
new java.lang.Byte("1"), new java.lang.Float(1.0), 1.0) ::
Row(0, new java.lang.Long(0), new java.lang.Short("0"),
new java.lang.Byte("0"), new java.lang.Float(10), new java.lang.Double(10)) :: Nil)
}

test("replace nan with double") {
checkAnswer(
createNaNDF().na.replace("*", Map(
Double.NaN -> 10.toDouble
)),
Row(1, new java.lang.Long(1), new java.lang.Short("1"),
new java.lang.Byte("1"), new java.lang.Float(1.0), 1.0) ::
Row(0, new java.lang.Long(0), new java.lang.Short("0"),
new java.lang.Byte("0"), new java.lang.Float(10), new java.lang.Double(10)) :: Nil)
}

test("replace float with nan") {
checkAnswer(
createNaNDF().na.replace("*", Map(
1.0f -> Float.NaN
)),
Row(1, new java.lang.Long(1), new java.lang.Short("1"),
new java.lang.Byte("1"), java.lang.Float.NaN, java.lang.Double.NaN) ::
Row(0, new java.lang.Long(0), new java.lang.Short("0"),
new java.lang.Byte("0"), java.lang.Float.NaN, java.lang.Double.NaN) :: Nil)
}

test("replace double with nan") {
checkAnswer(
createNaNDF().na.replace("*", Map(
1.toDouble -> Double.NaN
)),
Row(1, new java.lang.Long(1), new java.lang.Short("1"),
new java.lang.Byte("1"), java.lang.Float.NaN, java.lang.Double.NaN) ::
Row(0, new java.lang.Long(0), new java.lang.Short("0"),
new java.lang.Byte("0"), java.lang.Float.NaN, java.lang.Double.NaN) :: Nil)
}
}