chore: Parquet fuzz testing (#1623)

Author: andygrove

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

@@ -34,11 +34,14 @@ import org.apache.spark.sql.types.{ArrayType, DataType, DataTypes, DecimalType,
 object ParquetGenerator {
 
   /**
-   * Arbitrary date to use as base for generating temporal columns. Random integers will be added
-   * to or subtracted from this value.
+   * Date to use as base for generating temporal columns. Random integers will be added to or
+   * subtracted from this value.
+   *
+   * Date was chosen to trigger generating a timestamp that's larger than a 64-bit nanosecond
+   * timestamp can represent so that we test support for INT96 timestamps.
    */
-  private val baseDate =
-    new SimpleDateFormat("YYYY-MM-DD hh:mm:ss").parse("2024-05-25 12:34:56").getTime
+  val defaultBaseDate: Long =
+    new SimpleDateFormat("YYYY-MM-DD hh:mm:ss").parse("3333-05-25 12:34:56").getTime
 
   private val primitiveTypes = Seq(
     DataTypes.BooleanType,
@@ -217,13 +220,13 @@ object ParquetGenerator {
               null
           }
       case DataTypes.DateType =>
-        Range(0, numRows).map(_ => new java.sql.Date(baseDate + r.nextInt()))
+        Range(0, numRows).map(_ => new java.sql.Date(options.baseDate + r.nextInt()))
       case DataTypes.TimestampType =>
-        Range(0, numRows).map(_ => new Timestamp(baseDate + r.nextInt()))
+        Range(0, numRows).map(_ => new Timestamp(options.baseDate + r.nextInt()))
       case DataTypes.TimestampNTZType =>
         Range(0, numRows).map(_ =>
           LocalDateTime.ofInstant(
-            Instant.ofEpochMilli(baseDate + r.nextInt()),
+            Instant.ofEpochMilli(options.baseDate + r.nextInt()),
             ZoneId.systemDefault()))
       case _ => throw new IllegalStateException(s"Cannot generate data for $dataType yet")
     }
@@ -234,6 +237,7 @@ object ParquetGenerator {
 case class DataGenOptions(
     allowNull: Boolean = true,
     generateNegativeZero: Boolean = true,
+    baseDate: Long = ParquetGenerator.defaultBaseDate,
     generateArray: Boolean = false,
     generateStruct: Boolean = false,
     generateMap: Boolean = false)

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

@@ -20,6 +20,7 @@
 package org.apache.comet
 
 import java.io.File
+import java.text.SimpleDateFormat
 
 import scala.util.Random
 
@@ -32,6 +33,8 @@ import org.apache.spark.sql.comet.{CometNativeScanExec, CometScanExec}
 import org.apache.spark.sql.execution.SparkPlan
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
 import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.internal.SQLConf.ParquetOutputTimestampType
+import org.apache.spark.sql.types.{ArrayType, DataType, DataTypes, StructType}
 
 import org.apache.comet.testing.{DataGenOptions, ParquetGenerator}
 
@@ -57,7 +60,13 @@ class CometFuzzTestSuite extends CometTestBase with AdaptiveSparkPlanHelper {
       CometConf.COMET_ENABLED.key -> "false",
       SQLConf.SESSION_LOCAL_TIMEZONE.key -> defaultTimezone) {
       val options =
-        DataGenOptions(generateArray = true, generateStruct = true, generateNegativeZero = false)
+        DataGenOptions(
+          generateArray = true,
+          generateStruct = true,
+          generateNegativeZero = false,
+          // override base date due to known issues with experimental scans
+          baseDate =
+            new SimpleDateFormat("YYYY-MM-DD hh:mm:ss").parse("2024-05-25 12:34:56").getTime)
       ParquetGenerator.makeParquetFile(random, spark, filename, 1000, options)
     }
   }
@@ -166,6 +175,75 @@ class CometFuzzTestSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     }
   }
 
+  test("Parquet temporal types written as INT96") {
+
+    // there are known issues with INT96 support in the new experimental scans
+    // https://github.com/apache/datafusion-comet/issues/1441
+    assume(!CometConf.isExperimentalNativeScan)
+
+    testParquetTemporalTypes(ParquetOutputTimestampType.INT96)
+  }
+
+  test("Parquet temporal types written as TIMESTAMP_MICROS") {
+    testParquetTemporalTypes(ParquetOutputTimestampType.TIMESTAMP_MICROS)
+  }
+
+  test("Parquet temporal types written as TIMESTAMP_MILLIS") {
+    testParquetTemporalTypes(ParquetOutputTimestampType.TIMESTAMP_MILLIS)
+  }
+
+  private def testParquetTemporalTypes(
+      outputTimestampType: ParquetOutputTimestampType.Value): Unit = {
+
+    val options =
+      DataGenOptions(generateArray = true, generateStruct = true, generateNegativeZero = false)
+
+    withTempPath { filename =>
+      val random = new Random(42)
+      withSQLConf(
+        CometConf.COMET_ENABLED.key -> "false",
+        SQLConf.PARQUET_OUTPUT_TIMESTAMP_TYPE.key -> outputTimestampType.toString,
+        SQLConf.SESSION_LOCAL_TIMEZONE.key -> defaultTimezone) {
+        ParquetGenerator.makeParquetFile(random, spark, filename.toString, 100, options)
+      }
+
+      Seq(defaultTimezone, "UTC", "America/Denver").foreach { tz =>
+        Seq(true, false).foreach { inferTimestampNtzEnabled =>
+          Seq(true, false).foreach { int96TimestampConversion =>
+            Seq(true, false).foreach { int96AsTimestamp =>
+              withSQLConf(
+                CometConf.COMET_ENABLED.key -> "true",
+                SQLConf.SESSION_LOCAL_TIMEZONE.key -> tz,
+                SQLConf.PARQUET_INT96_AS_TIMESTAMP.key -> int96AsTimestamp.toString,
+                SQLConf.PARQUET_INT96_TIMESTAMP_CONVERSION.key -> int96TimestampConversion.toString,
+                SQLConf.PARQUET_INFER_TIMESTAMP_NTZ_ENABLED.key -> inferTimestampNtzEnabled.toString) {
+
+                val df = spark.read.parquet(filename.toString)
+                df.createOrReplaceTempView("t1")
+
+                def hasTemporalType(t: DataType): Boolean = t match {
+                  case DataTypes.DateType | DataTypes.TimestampType |
+                      DataTypes.TimestampNTZType =>
+                    true
+                  case t: StructType => t.exists(f => hasTemporalType(f.dataType))
+                  case t: ArrayType => hasTemporalType(t.elementType)
+                  case _ => false
+                }
+
+                val columns =
+                  df.schema.fields.filter(f => hasTemporalType(f.dataType)).map(_.name)
+
+                for (col <- columns) {
+                  checkSparkAnswer(s"SELECT $col FROM t1 ORDER BY $col")
+                }
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+
   override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit
       pos: Position): Unit = {
     Seq("native", "jvm").foreach { shuffleMode =>