Refactor Parquet readers [databricks]

sql-plugin/src/main/java/com/nvidia/spark/rapids/iceberg/parquet/GpuParquetReader.java

@@ -25,6 +25,7 @@
 
 import scala.collection.Seq;
 
+import com.nvidia.spark.rapids.DateTimeRebaseCorrected$;
 import com.nvidia.spark.rapids.GpuMetric;
 import com.nvidia.spark.rapids.GpuParquetUtils;
 import com.nvidia.spark.rapids.ParquetPartitionReader;
@@ -139,8 +140,8 @@ public org.apache.iceberg.io.CloseableIterator<ColumnarBatch> iterator() {
           new Path(input.location()), clippedBlocks, fileReadSchema, caseSensitive,
           partReaderSparkSchema, debugDumpPrefix, debugDumpAlways,
           maxBatchSizeRows, maxBatchSizeBytes, targetBatchSizeBytes, useChunkedReader, metrics,
-          true, // isCorrectedInt96RebaseMode
-          true, // isCorrectedRebaseMode
+          DateTimeRebaseCorrected$.MODULE$, // dateRebaseMode
+          DateTimeRebaseCorrected$.MODULE$, // timestampRebaseMode
           true, // hasInt96Timestamps
           false // useFieldId
       );

sql-plugin/src/main/java/com/nvidia/spark/rapids/iceberg/spark/source/GpuMultiFileBatchReader.java

@@ -220,10 +220,12 @@ static class IcebergParquetExtraInfo extends ParquetExtraInfo {
     private final Schema expectedSchema;
     private final PartitionSpec partitionSpec;
 
-    IcebergParquetExtraInfo(boolean isCorrectedRebaseMode,
-        boolean isCorrectedInt96RebaseMode, boolean hasInt96Timestamps,
-        Map<Integer, ?> idToConstant, Schema expectedSchema, PartitionSpec partitionSpec) {
-      super(isCorrectedRebaseMode, isCorrectedInt96RebaseMode, hasInt96Timestamps);
+    IcebergParquetExtraInfo(DateTimeRebaseMode dateRebaseMode,
+                            DateTimeRebaseMode timestampRebaseMode,
+                            boolean hasInt96Timestamps,
+                            Map<Integer, ?> idToConstant, Schema expectedSchema,
+                            PartitionSpec partitionSpec) {
+      super(dateRebaseMode, timestampRebaseMode, hasInt96Timestamps);
       this.idToConstant = idToConstant;
       this.expectedSchema = expectedSchema;
       this.partitionSpec = partitionSpec;
@@ -309,8 +311,8 @@ protected FilteredParquetFileInfo filterParquetBlocks(FileScanTask fst) {
         ParquetFileInfoWithBlockMeta parquetBlockMeta = ParquetFileInfoWithBlockMeta.apply(
             new Path(new URI(fst.file().path().toString())), clippedBlocks,
             InternalRow.empty(), fileReadSchema, partReaderSparkSchema,
-            true, // isCorrectedInt96RebaseMode
-            true, // isCorrectedRebaseMode
+            DateTimeRebaseCorrected$.MODULE$, // dateRebaseMode
+            DateTimeRebaseCorrected$.MODULE$, // timestampRebaseMode
             true //  hasInt96Timestamps
         );
         return new FilteredParquetFileInfo(parquetBlockMeta, updatedConstants, updatedSchema);
@@ -397,8 +399,8 @@ protected FilePartitionReaderBase createRapidsReader(PartitionedFile[] pFiles,
                 ParquetSchemaWrapper.apply(filteredInfo.parquetBlockMeta.schema()),
                 filteredInfo.parquetBlockMeta.readSchema(),
                 new IcebergParquetExtraInfo(
-                    filteredInfo.parquetBlockMeta.isCorrectedRebaseMode(),
-                    filteredInfo.parquetBlockMeta.isCorrectedInt96RebaseMode(),
+                    filteredInfo.parquetBlockMeta.dateRebaseMode(),
+                    filteredInfo.parquetBlockMeta.timestampRebaseMode(),
                     filteredInfo.parquetBlockMeta.hasInt96Timestamps(),
                     filteredInfo.idToConstant(),
                     filteredInfo.expectedSchema(),

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuParquetFileFormat.scala

@@ -19,7 +19,6 @@ package com.nvidia.spark.rapids
 import java.time.ZoneId
 
 import ai.rapids.cudf._
-import com.nvidia.spark.RebaseHelper
 import com.nvidia.spark.rapids.Arm.withResource
 import com.nvidia.spark.rapids.RapidsPluginImplicits.AutoCloseableProducingArray
 import com.nvidia.spark.rapids.jni.DateTimeRebase
@@ -121,28 +120,31 @@ object GpuParquetFileFormat {
       }
     }
 
-    SparkShimImpl.int96ParquetRebaseWrite(sqlConf) match {
-      case "EXCEPTION" =>
-      case "CORRECTED" =>
-      case "LEGACY" =>
+    DateTimeRebaseMode.fromName(SparkShimImpl.int96ParquetRebaseWrite(sqlConf)) match {
+      case DateTimeRebaseException | DateTimeRebaseCorrected => // Good
+      case DateTimeRebaseLegacy =>
         if (schemaHasTimestamps) {
           meta.willNotWorkOnGpu("LEGACY rebase mode for int96 timestamps is not supported")
         }
-      case other =>
-        meta.willNotWorkOnGpu(s"$other is not a supported rebase mode for int96")
+      // This should never be reached out, since invalid mode is handled in
+      // `DateTimeRebaseMode.fromName`.
+      case other => meta.willNotWorkOnGpu(
+        DateTimeRebaseUtils.invalidRebaseModeMessage(other.getClass.getName))
     }
 
-    SparkShimImpl.parquetRebaseWrite(sqlConf) match {
-      case "EXCEPTION" | "CORRECTED" => // Good
-      case "LEGACY" =>
+    DateTimeRebaseMode.fromName(SparkShimImpl.parquetRebaseWrite(sqlConf)) match {
+      case DateTimeRebaseException | DateTimeRebaseCorrected => // Good
+      case DateTimeRebaseLegacy =>
         if (!TypeChecks.areTimestampsSupported()) {
           meta.willNotWorkOnGpu("Only UTC timezone is supported in LEGACY rebase mode. " +
             s"Current timezone settings: (JVM : ${ZoneId.systemDefault()}, " +
             s"session: ${SQLConf.get.sessionLocalTimeZone}). " +
             " Set both of the timezones to UTC to enable LEGACY rebase support.")
         }
-      case other =>
-        meta.willNotWorkOnGpu(s"$other is not a supported rebase mode")
+      // This should never be reached out, since invalid mode is handled in
+      // `DateTimeRebaseMode.fromName`.
+      case other => meta.willNotWorkOnGpu(
+        DateTimeRebaseUtils.invalidRebaseModeMessage(other.getClass.getName))
     }
 
     if (meta.canThisBeReplaced) {
@@ -193,9 +195,11 @@ class GpuParquetFileFormat extends ColumnarFileFormat with Logging {
     val conf = ContextUtil.getConfiguration(job)
 
     val outputTimestampType = sqlConf.parquetOutputTimestampType
-    val dateTimeRebaseMode = sparkSession.sqlContext.getConf(SparkShimImpl.parquetRebaseWriteKey)
+    val dateTimeRebaseMode = DateTimeRebaseMode.fromName(
+      sparkSession.sqlContext.getConf(SparkShimImpl.parquetRebaseWriteKey))
     val timestampRebaseMode = if (outputTimestampType.equals(ParquetOutputTimestampType.INT96)) {
-      sparkSession.sqlContext.getConf(SparkShimImpl.int96ParquetRebaseWriteKey)
+      DateTimeRebaseMode.fromName(
+        sparkSession.sqlContext.getConf(SparkShimImpl.int96ParquetRebaseWriteKey))
     } else {
       dateTimeRebaseMode
     }
@@ -302,19 +306,20 @@ class GpuParquetWriter(
     dataSchema: StructType,
     compressionType: CompressionType,
     outputTimestampType: String,
-    dateRebaseMode: String,
-    timestampRebaseMode: String,
+    dateRebaseMode: DateTimeRebaseMode,
+    timestampRebaseMode: DateTimeRebaseMode,
     context: TaskAttemptContext,
     parquetFieldIdEnabled: Boolean)
   extends ColumnarOutputWriter(context, dataSchema, "Parquet", true) {
   override def throwIfRebaseNeededInExceptionMode(batch: ColumnarBatch): Unit = {
     val cols = GpuColumnVector.extractBases(batch)
     cols.foreach { col =>
-      if (dateRebaseMode.equals("EXCEPTION") && RebaseHelper.isDateRebaseNeededInWrite(col)) {
+      if (dateRebaseMode == DateTimeRebaseException &&
+        DateTimeRebaseUtils.isDateRebaseNeededInWrite(col)) {
         throw DataSourceUtils.newRebaseExceptionInWrite("Parquet")
       }
-      else if (timestampRebaseMode.equals("EXCEPTION") &&
-               RebaseHelper.isTimeRebaseNeededInWrite(col)) {
+      else if (timestampRebaseMode == DateTimeRebaseException &&
+               DateTimeRebaseUtils.isTimeRebaseNeededInWrite(col)) {
         throw DataSourceUtils.newRebaseExceptionInWrite("Parquet")
       }
     }
@@ -334,14 +339,14 @@ class GpuParquetWriter(
     ColumnCastUtil.deepTransform(cv, Some(dt)) {
       case (cv, _) if cv.getType.isTimestampType =>
         if(cv.getType == DType.TIMESTAMP_DAYS) {
-          if (dateRebaseMode.equals("LEGACY")) {
+          if (dateRebaseMode == DateTimeRebaseLegacy) {
             DateTimeRebase.rebaseGregorianToJulian(cv)
           } else {
             cv.copyToColumnVector()
           }
         } else { /* timestamp */
           val typeMillis = ParquetOutputTimestampType.TIMESTAMP_MILLIS.toString
-          if (timestampRebaseMode.equals("LEGACY")) {
+          if (timestampRebaseMode == DateTimeRebaseLegacy) {
             val rebasedTimestampAsMicros = if(cv.getType == DType.TIMESTAMP_MICROSECONDS) {
               DateTimeRebase.rebaseGregorianToJulian(cv)
             } else {