Improve dateFormat support in GpuJsonScan and make tests consistent with GpuStructsToJson [databricks]

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

@@ -359,7 +359,7 @@ abstract class CSVPartitionReaderBase[BUFF <: LineBufferer, FACT <: LineBufferer
     }
   }
 
-  override def dateFormat: String = GpuCsvUtils.dateFormatInRead(parsedOptions)
+  override def dateFormat: Option[String] = Some(GpuCsvUtils.dateFormatInRead(parsedOptions))
   override def timestampFormat: String = GpuCsvUtils.timestampFormatInRead(parsedOptions)
 }
 

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

@@ -19,17 +19,14 @@ package com.nvidia.spark.rapids
 import java.text.SimpleDateFormat
 import java.time.DateTimeException
 import java.util.Optional
-
 import scala.collection.mutable.ArrayBuffer
-
-import ai.rapids.cudf.{BinaryOp, CaptureGroups, ColumnVector, ColumnView, DecimalUtils, DType, RegexProgram, Scalar}
+import ai.rapids.cudf.{BinaryOp, CaptureGroups, ColumnVector, ColumnView, DType, DecimalUtils, RegexProgram, Scalar}
 import ai.rapids.cudf
 import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource}
 import com.nvidia.spark.rapids.RapidsPluginImplicits._
 import com.nvidia.spark.rapids.jni.CastStrings
 import com.nvidia.spark.rapids.shims.{AnsiUtil, GpuCastShims, GpuIntervalUtils, GpuTypeShims, SparkShimImpl, YearParseUtil}
 import org.apache.commons.text.StringEscapeUtils
-
 import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.expressions.{Cast, Expression, NullIntolerant, TimeZoneAwareExpression, UnaryExpression}
 import org.apache.spark.sql.catalyst.util.DateTimeConstants.MICROS_PER_SECOND
@@ -1284,7 +1281,8 @@ object GpuCast {
   def convertDateOrNull(
       input: ColumnVector,
       regex: String,
-      cudfFormat: String): ColumnVector = {
+      cudfFormat: String,
+      failOnInvalid: Boolean = false): ColumnVector = {
 
     val prog = new RegexProgram(regex, CaptureGroups.NON_CAPTURE)
     val isValidDate = withResource(input.matchesRe(prog)) { isMatch =>
@@ -1294,6 +1292,13 @@ object GpuCast {
     }
 
     withResource(isValidDate) { _ =>
+      if (failOnInvalid) {
+        withResource(isValidDate.all()) { all =>
+          if (all.isValid && !all.getBoolean) {
+            throw new DateTimeException("One or more values is not a valid date")
+          }
+        }
+      }
       withResource(Scalar.fromNull(DType.TIMESTAMP_DAYS)) { orElse =>
         withResource(input.asTimestampDays(cudfFormat)) { asDays =>
           isValidDate.ifElse(asDays, orElse)
@@ -1376,7 +1381,7 @@ object GpuCast {
     }
   }
 
-  private def castStringToDateAnsi(input: ColumnVector, ansiMode: Boolean): ColumnVector = {
+  def castStringToDateAnsi(input: ColumnVector, ansiMode: Boolean): ColumnVector = {
     val result = castStringToDate(input)
     if (ansiMode) {
       // When ANSI mode is enabled, we need to throw an exception if any values could not be

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

@@ -16,7 +16,7 @@
 
 package com.nvidia.spark.rapids
 
-import java.time.DateTimeException
+import java.util
 import java.util.Optional
 
 import scala.collection.mutable.ListBuffer
@@ -25,8 +25,7 @@ import ai.rapids.cudf.{CaptureGroups, ColumnVector, DType, HostColumnVector, Hos
 import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource}
 import com.nvidia.spark.rapids.DateUtils.{toStrf, TimestampFormatConversionException}
 import com.nvidia.spark.rapids.jni.CastStrings
-import com.nvidia.spark.rapids.shims.GpuTypeShims
-import java.util
+import com.nvidia.spark.rapids.shims.{GpuJsonToStructsShim, GpuTypeShims}
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
 import org.apache.hadoop.io.compress.CompressionCodecFactory
@@ -36,7 +35,7 @@ import org.apache.spark.sql.connector.read.PartitionReader
 import org.apache.spark.sql.execution.QueryExecutionException
 import org.apache.spark.sql.execution.datasources.{HadoopFileLinesReader, PartitionedFile}
 import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.rapids.{ExceptionTimeParserPolicy, GpuToTimestamp, LegacyTimeParserPolicy}
+import org.apache.spark.sql.rapids.{GpuToTimestamp, LegacyTimeParserPolicy}
 import org.apache.spark.sql.types.{DataTypes, DecimalType, StructField, StructType}
 import org.apache.spark.sql.vectorized.ColumnarBatch
 
@@ -372,31 +371,11 @@ abstract class GpuTextBasedPartitionReader[BUFF <: LineBufferer, FACT <: LineBuf
     }
   }
 
-  def dateFormat: String
+  def dateFormat: Option[String]
   def timestampFormat: String
 
   def castStringToDate(input: ColumnVector, dt: DType): ColumnVector = {
-    castStringToDate(input, dt, failOnInvalid = true)
-  }
-
-  def castStringToDate(input: ColumnVector, dt: DType, failOnInvalid: Boolean): ColumnVector = {
-    val cudfFormat = DateUtils.toStrf(dateFormat, parseString = true)
-    withResource(input.strip()) { stripped =>
-      withResource(stripped.isTimestamp(cudfFormat)) { isDate =>
-        if (failOnInvalid && GpuOverrides.getTimeParserPolicy == ExceptionTimeParserPolicy) {
-          withResource(isDate.all()) { all =>
-            if (all.isValid && !all.getBoolean) {
-              throw new DateTimeException("One or more values is not a valid date")
-            }
-          }
-        }
-        withResource(stripped.asTimestamp(dt, cudfFormat)) { asDate =>
-          withResource(Scalar.fromNull(dt)) { nullScalar =>
-            isDate.ifElse(asDate, nullScalar)
-          }
-        }
-      }
-    }
+    GpuJsonToStructsShim.castJsonStringToDateFromScan(input, dt, dateFormat)
   }
 
   def castStringToTimestamp(

sql-plugin/src/main/scala/org/apache/spark/sql/catalyst/json/rapids/GpuJsonScan.scala

@@ -25,7 +25,7 @@ import ai.rapids.cudf
 import ai.rapids.cudf.{CaptureGroups, ColumnVector, DType, NvtxColor, RegexProgram, Scalar, Schema, Table}
 import com.nvidia.spark.rapids._
 import com.nvidia.spark.rapids.Arm.withResource
-import com.nvidia.spark.rapids.shims.{ColumnDefaultValuesShims, LegacyBehaviorPolicyShim, ShimFilePartitionReaderFactory}
+import com.nvidia.spark.rapids.shims.{ColumnDefaultValuesShims, GpuJsonToStructsShim, LegacyBehaviorPolicyShim, ShimFilePartitionReaderFactory}
 import org.apache.hadoop.conf.Configuration
 
 import org.apache.spark.broadcast.Broadcast
@@ -113,16 +113,15 @@ object GpuJsonScan {
 
     val hasDates = TrampolineUtil.dataTypeExistsRecursively(dt, _.isInstanceOf[DateType])
     if (hasDates) {
-      GpuJsonUtils.optionalDateFormatInRead(parsedOptions) match {
-        case None | Some("yyyy-MM-dd") =>
-          // this is fine
-        case dateFormat =>
-          meta.willNotWorkOnGpu(s"GpuJsonToStructs unsupported dateFormat $dateFormat")
-      }
+      GpuJsonToStructsShim.tagDateFormatSupport(meta,
+        GpuJsonUtils.optionalDateFormatInRead(parsedOptions))
     }
 
     val hasTimestamps = TrampolineUtil.dataTypeExistsRecursively(dt, _.isInstanceOf[TimestampType])
     if (hasTimestamps) {
+      GpuJsonToStructsShim.tagTimestampFormatSupport(meta,
+        GpuJsonUtils.optionalTimestampFormatInRead(parsedOptions))
+
       GpuJsonUtils.optionalTimestampFormatInRead(parsedOptions) match {
         case None | Some("yyyy-MM-dd'T'HH:mm:ss[.SSS][XXX]") =>
           // this is fine
@@ -163,10 +162,16 @@ object GpuJsonScan {
     tagSupportOptions(parsedOptions, meta)
 
     val types = readSchema.map(_.dataType)
-    if (types.contains(DateType)) {
+
+    val hasDates = TrampolineUtil.dataTypeExistsRecursively(readSchema, _.isInstanceOf[DateType])
+    if (hasDates) {
+
       GpuTextBasedDateUtils.tagCudfFormat(meta,
         GpuJsonUtils.dateFormatInRead(parsedOptions), parseString = true)
 
+      GpuJsonToStructsShim.tagDateFormatSupportFromScan(meta,
+        GpuJsonUtils.optionalDateFormatInRead(parsedOptions))
+
       // For date type, timezone needs to be checked also. This is because JVM timezone is used
       // to get days offset before rebasing Julian to Gregorian in Spark while not in Rapids.
       //
@@ -490,6 +495,6 @@ class JsonPartitionReader(
     }
   }
 
-  override def dateFormat: String = GpuJsonUtils.dateFormatInRead(parsedOptions)
+  override def dateFormat: Option[String] = GpuJsonUtils.optionalDateFormatInRead(parsedOptions)
   override def timestampFormat: String = GpuJsonUtils.timestampFormatInRead(parsedOptions)
 }

sql-plugin/src/main/spark311/scala/com/nvidia/spark/rapids/shims/GpuJsonToStructsShim.scala

@@ -17,35 +17,28 @@
 {"spark": "311"}
 {"spark": "312"}
 {"spark": "313"}
-{"spark": "320"}
-{"spark": "321"}
-{"spark": "321cdh"}
-{"spark": "321db"}
-{"spark": "322"}
-{"spark": "323"}
-{"spark": "324"}
-{"spark": "330"}
-{"spark": "330cdh"}
-{"spark": "330db"}
-{"spark": "331"}
-{"spark": "332"}
-{"spark": "332cdh"}
-{"spark": "332db"}
-{"spark": "333"}
 spark-rapids-shim-json-lines ***/
 package com.nvidia.spark.rapids.shims
 
-import ai.rapids.cudf.{ColumnVector, Scalar}
+import ai.rapids.cudf.{ColumnVector, DType, Scalar}
 import com.nvidia.spark.rapids.Arm.withResource
-import com.nvidia.spark.rapids.GpuCast
+import com.nvidia.spark.rapids.{GpuCast, RapidsMeta, GpuOverrides}
 
 import org.apache.spark.sql.catalyst.json.GpuJsonUtils
+import org.apache.spark.sql.rapids.ExceptionTimeParserPolicy
 
 object GpuJsonToStructsShim {
+  def tagDateFormatSupport(meta: RapidsMeta[_, _, _], dateFormat: Option[String]): Unit = {
+    dateFormat match {
+      case None | Some("yyyy-MM-dd") =>
+      case dateFormat =>
+        meta.willNotWorkOnGpu(s"GpuJsonToStructs unsupported dateFormat $dateFormat")
+    }
+  }
 
   def castJsonStringToDate(input: ColumnVector, options: Map[String, String]): ColumnVector = {
-    GpuJsonUtils.dateFormatInRead(options) match {
-      case "yyyy-MM-dd" =>
+    GpuJsonUtils.optionalDateFormatInRead(options) match {
+      case None | Some("yyyy-MM-dd") =>
         withResource(Scalar.fromString(" ")) { space =>
           withResource(input.strip(space)) { trimmed =>
             GpuCast.castStringToDate(trimmed)
@@ -57,6 +50,27 @@ object GpuJsonToStructsShim {
     }
   }
 
+  def tagDateFormatSupportFromScan(meta: RapidsMeta[_, _, _], dateFormat: Option[String]): Unit = {
+    tagDateFormatSupport(meta, dateFormat)
+  }
+
+  def castJsonStringToDateFromScan(input: ColumnVector, dt: DType,
+      dateFormat: Option[String]): ColumnVector = {
+    dateFormat match {
+      case None | Some("yyyy-MM-dd") =>
+        withResource(input.strip()) { trimmed =>
+          GpuCast.castStringToDateAnsi(trimmed, ansiMode =
+            GpuOverrides.getTimeParserPolicy == ExceptionTimeParserPolicy)
+        }
+      case other =>
+        // should be unreachable due to GpuOverrides checks
+        throw new IllegalStateException(s"Unsupported dateFormat $other")
+    }
+  }
+
+  def tagTimestampFormatSupport(meta: RapidsMeta[_, _, _],
+    timestampFormat: Option[String]): Unit = {}
+
   def castJsonStringToTimestamp(input: ColumnVector,
       options: Map[String, String]): ColumnVector = {
     withResource(Scalar.fromString(" ")) { space =>

sql-plugin/src/main/spark320/scala/com/nvidia/spark/rapids/shims/GpuJsonToStructsShim.scala

@@ -0,0 +1,93 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed 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.
+ */
+/*** spark-rapids-shim-json-lines
+{"spark": "320"}
+{"spark": "321"}
+{"spark": "321cdh"}
+{"spark": "321db"}
+{"spark": "322"}
+{"spark": "323"}
+{"spark": "324"}
+{"spark": "330"}
+{"spark": "330cdh"}
+{"spark": "330db"}
+{"spark": "331"}
+{"spark": "332"}
+{"spark": "332cdh"}
+{"spark": "332db"}
+{"spark": "333"}
+spark-rapids-shim-json-lines ***/
+package com.nvidia.spark.rapids.shims
+
+import ai.rapids.cudf.{ColumnVector, DType, Scalar}
+import com.nvidia.spark.rapids.Arm.withResource
+import com.nvidia.spark.rapids.{DateUtils, GpuCast, GpuOverrides, RapidsMeta}
+
+import org.apache.spark.sql.rapids.ExceptionTimeParserPolicy
+
+object GpuJsonToStructsShim {
+
+  def tagDateFormatSupport(meta: RapidsMeta[_, _, _], dateFormat: Option[String]): Unit = {
+  }
+
+  def castJsonStringToDate(input: ColumnVector, options: Map[String, String]): ColumnVector = {
+    // dateFormat is ignored in from_json in Spark 3.2
+    withResource(Scalar.fromString(" ")) { space =>
+      withResource(input.strip(space)) { trimmed =>
+        GpuCast.castStringToDate(trimmed)
+      }
+    }
+  }
+
+  def tagDateFormatSupportFromScan(meta: RapidsMeta[_, _, _], dateFormat: Option[String]): Unit = {
+  }
+
+  def castJsonStringToDateFromScan(input: ColumnVector, dt: DType,
+      dateFormat: Option[String]): ColumnVector = {
+    dateFormat match {
+      case None =>
+        // legacy behavior
+        withResource(input.strip()) { trimmed =>
+          GpuCast.castStringToDateAnsi(trimmed, ansiMode =
+            GpuOverrides.getTimeParserPolicy == ExceptionTimeParserPolicy)
+        }
+      case Some(f) =>
+        withResource(input.strip()) { trimmed =>
+          val regexRoot = dateFormatPattern
+            .replace("yyyy", raw"\d{4}")
+            .replace("MM", raw"\d{1,2}")
+            .replace("dd", raw"\d{1,2}")
+          val cudfFormat = DateUtils.toStrf(dateFormatPattern, parseString = true)
+          GpuCast.convertDateOrNull(input, "^" + regexRoot + "$", cudfFormat,
+            failOnInvalid = GpuOverrides.getTimeParserPolicy == ExceptionTimeParserPolicy)
+        }
+    }
+  }
+
+  def tagTimestampFormatSupport(meta: RapidsMeta[_, _, _],
+      timestampFormat: Option[String]): Unit = {}
+
+  def castJsonStringToTimestamp(input: ColumnVector,
+      options: Map[String, String]): ColumnVector = {
+    // legacy behavior
+    withResource(Scalar.fromString(" ")) { space =>
+      withResource(input.strip(space)) { trimmed =>
+        // from_json doesn't respect ansi mode
+        GpuCast.castStringToTimestamp(trimmed, ansiMode = false)
+      }
+    }
+  }
+}