chore: generate Float/Double NaN (#2695)

Author: hsiang-c

spark/src/main/scala/org/apache/comet/testing/FuzzDataGenerator.scala

@@ -44,6 +44,9 @@ object FuzzDataGenerator {
   val defaultBaseDate: Long =
     new SimpleDateFormat("YYYY-MM-DD hh:mm:ss").parse("3333-05-25 12:34:56").getTime
 
+  val floatNaNLiteral = "FLOAT('NaN')"
+  val doubleNaNLiteral = "DOUBLE('NaN')"
+
   def generateSchema(options: SchemaGenOptions): StructType = {
     val primitiveTypes = options.primitiveTypes
     val dataTypes = ListBuffer[DataType]()
@@ -168,6 +171,7 @@ object FuzzDataGenerator {
             case 4 => Float.MaxValue
             case 5 => 0.0f
             case 6 if options.generateNegativeZero => -0.0f
+            case 7 if options.generateNaN => Float.NaN
             case _ => r.nextFloat()
           }
         })
@@ -181,6 +185,7 @@ object FuzzDataGenerator {
             case 4 => Double.MaxValue
             case 5 => 0.0
             case 6 if options.generateNegativeZero => -0.0
+            case 7 if options.generateNaN => Double.NaN
             case _ => r.nextDouble()
           }
         })
@@ -257,6 +262,7 @@ case class SchemaGenOptions(
 case class DataGenOptions(
     allowNull: Boolean = true,
     generateNegativeZero: Boolean = true,
+    generateNaN: Boolean = true,
     baseDate: Long = FuzzDataGenerator.defaultBaseDate,
     customStrings: Seq[String] = Seq.empty,
     maxStringLength: Int = 8)

spark/src/test/scala/org/apache/comet/CometFuzzTestSuite.scala

@@ -29,6 +29,7 @@ import org.apache.spark.sql.types._
 
 import org.apache.comet.DataTypeSupport.isComplexType
 import org.apache.comet.testing.{DataGenOptions, ParquetGenerator, SchemaGenOptions}
+import org.apache.comet.testing.FuzzDataGenerator.{doubleNaNLiteral, floatNaNLiteral}
 
 class CometFuzzTestSuite extends CometFuzzTestBase {
 
@@ -71,8 +72,20 @@ class CometFuzzTestSuite extends CometFuzzTestBase {
         // Construct the string for the default value based on the column type.
         val defaultValueString = defaultValueType match {
           // These explicit type definitions for TINYINT, SMALLINT, FLOAT, DOUBLE, and DATE are only needed for 3.4.
-          case "TINYINT" | "SMALLINT" | "FLOAT" | "DOUBLE" =>
+          case "TINYINT" | "SMALLINT" =>
             s"$defaultValueType(${defaultValueRow.get(0)})"
+          case "FLOAT" =>
+            if (Float.NaN.equals(defaultValueRow.get(0))) {
+              floatNaNLiteral
+            } else {
+              s"$defaultValueType(${defaultValueRow.get(0)})"
+            }
+          case "DOUBLE" =>
+            if (Double.NaN.equals(defaultValueRow.get(0))) {
+              doubleNaNLiteral
+            } else {
+              s"$defaultValueType(${defaultValueRow.get(0)})"
+            }
           case "DATE" => s"$defaultValueType('${defaultValueRow.get(0)}')"
           case "STRING" => s"'${defaultValueRow.get(0)}'"
           case "TIMESTAMP" | "TIMESTAMP_NTZ" => s"TIMESTAMP '${defaultValueRow.get(0)}'"
@@ -101,7 +114,7 @@ class CometFuzzTestSuite extends CometFuzzTestBase {
                 .asInstanceOf[Array[Byte]]
                 .sameElements(spark.sql(sql).collect()(0).get(0).asInstanceOf[Array[Byte]]))
           } else {
-            assert(defaultValueRow.get(0) == spark.sql(sql).collect()(0).get(0))
+            assert(defaultValueRow.get(0).equals(spark.sql(sql).collect()(0).get(0)))
           }
         }
       }