Specify an ordering for NaN values.

Author: JoshRosen

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala

@@ -182,6 +182,8 @@ class CodeGenContext {
   def genComp(dataType: DataType, c1: String, c2: String): String = dataType match {
     // java boolean doesn't support > or < operator
     case BooleanType => s"($c1 == $c2 ? 0 : ($c1 ? 1 : -1))"
+    case DoubleType => s"org.apache.spark.sql.catalyst.util.TypeUtils.compareDoubles($c1, $c2)"
+    case FloatType => s"org.apache.spark.sql.catalyst.util.TypeUtils.compareFloats($c1, $c2)"
     // use c1 - c2 may overflow
     case dt: DataType if isPrimitiveType(dt) => s"($c1 > $c2 ? 1 : $c1 < $c2 ? -1 : 0)"
     case BinaryType => s"org.apache.spark.sql.catalyst.util.TypeUtils.compareBinary($c1, $c2)"

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/TypeUtils.scala

@@ -70,4 +70,24 @@ object TypeUtils {
     }
     x.length - y.length
   }
+
+  def compareDoubles(x: Double, y: Double): Int = {
+    val xIsNan: Boolean = java.lang.Double.isNaN(x)
+    val yIsNan: Boolean = java.lang.Double.isNaN(y)
+    if ((xIsNan && yIsNan) || (x == y)) 0
+    else if (xIsNan) -1
+    else if (yIsNan) 1
+    else if (x > y) -1
+    else 1
+  }
+
+  def compareFloats(x: Float, y: Float): Int = {
+    val xIsNan: Boolean = java.lang.Float.isNaN(x)
+    val yIsNan: Boolean = java.lang.Float.isNaN(y)
+    if ((xIsNan && yIsNan) || (x == y)) 0
+    else if (xIsNan) -1
+    else if (yIsNan) 1
+    else if (x > y) -1
+    else 1
+  }
 }

sql/catalyst/src/main/scala/org/apache/spark/sql/types/DoubleType.scala

@@ -23,6 +23,7 @@ import scala.reflect.runtime.universe.typeTag
 
 import org.apache.spark.annotation.DeveloperApi
 import org.apache.spark.sql.catalyst.ScalaReflectionLock
+import org.apache.spark.sql.catalyst.util.TypeUtils
 
 /**
  * :: DeveloperApi ::
@@ -39,7 +40,9 @@ class DoubleType private() extends FractionalType {
   @transient private[sql] lazy val tag = ScalaReflectionLock.synchronized { typeTag[InternalType] }
   private[sql] val numeric = implicitly[Numeric[Double]]
   private[sql] val fractional = implicitly[Fractional[Double]]
-  private[sql] val ordering = implicitly[Ordering[InternalType]]
+  private[sql] val ordering = new Ordering[Double] {
+    override def compare(x: Double, y: Double): Int = TypeUtils.compareDoubles(x, y)
+  }
   private[sql] val asIntegral = DoubleAsIfIntegral
 
   /**

sql/catalyst/src/main/scala/org/apache/spark/sql/types/FloatType.scala

@@ -23,6 +23,7 @@ import scala.reflect.runtime.universe.typeTag
 
 import org.apache.spark.annotation.DeveloperApi
 import org.apache.spark.sql.catalyst.ScalaReflectionLock
+import org.apache.spark.sql.catalyst.util.TypeUtils
 
 /**
  * :: DeveloperApi ::
@@ -39,7 +40,9 @@ class FloatType private() extends FractionalType {
   @transient private[sql] lazy val tag = ScalaReflectionLock.synchronized { typeTag[InternalType] }
   private[sql] val numeric = implicitly[Numeric[Float]]
   private[sql] val fractional = implicitly[Fractional[Float]]
-  private[sql] val ordering = implicitly[Ordering[InternalType]]
+  private[sql] val ordering = new Ordering[Float] {
+    override def compare(x: Float, y: Float): Int = TypeUtils.compareFloats(x, y)
+  }
   private[sql] val asIntegral = FloatAsIfIntegral
 
   /**

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/TypeUtilsSuite.scala

@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.util
+
+import org.apache.spark.SparkFunSuite
+
+class TypeUtilsSuite extends SparkFunSuite {
+
+  import TypeUtils._
+
+  test("compareDoubles") {
+    assert(compareDoubles(0, 0) === 0)
+    assert(compareDoubles(1, 0) === -1)
+    assert(compareDoubles(0, 1) === 1)
+    assert(compareDoubles(Double.MinValue, Double.MaxValue) === 1)
+    assert(compareDoubles(Double.NaN, Double.NaN) === 0)
+    assert(compareDoubles(Double.NaN, Double.PositiveInfinity) === -1)
+    assert(compareDoubles(Double.NaN, Double.NegativeInfinity) === -1)
+  }
+
+  test("compareFloats") {
+    assert(compareFloats(0, 0) === 0)
+    assert(compareFloats(1, 0) === -1)
+    assert(compareFloats(0, 1) === 1)
+    assert(compareFloats(Float.MinValue, Float.MaxValue) === 1)
+    assert(compareFloats(Float.NaN, Float.NaN) === 0)
+    assert(compareFloats(Float.NaN, Float.PositiveInfinity) === -1)
+    assert(compareFloats(Float.NaN, Float.NegativeInfinity) === -1)
+  }
+}

sql/core/src/test/scala/org/apache/spark/sql/DataFrameSuite.scala

@@ -739,13 +739,13 @@ class DataFrameSuite extends QueryTest {
     df.col("t.``")
   }
 
-  test("SPARK-XXXX: sort by float column containing NaN") {
+  test("SPARK-8797: sort by float column containing NaN") {
     val inputData = Seq.fill(10)(Tuple1(Float.NaN)) ++ (1 to 1000).map(x => Tuple1(x.toFloat))
     val df = Random.shuffle(inputData).toDF("a")
     df.orderBy("a").collect()
   }
 
-  test("SPARK-XXXX: sort by double column containing NaN") {
+  test("SPARK-8797: sort by double column containing NaN") {
     val inputData = Seq.fill(10)(Tuple1(Double.NaN)) ++ (1 to 1000).map(x => Tuple1(x.toDouble))
     val df = Random.shuffle(inputData).toDF("a")
     df.orderBy("a").collect()