From 2caf870fe98f55ab5eea087935151a29438b2c94 Mon Sep 17 00:00:00 2001
From: Liangcai Li <firestarmanllc@gmail.com>
Date: Wed, 6 Dec 2023 18:16:56 +0800
Subject: [PATCH] Correct the fallback message for formats supported by
 conversion (#10)

---
 .../scala/com/nvidia/spark/rapids/DateUtils.scala  | 14 ++++++++------
 .../spark/sql/rapids/datetimeExpressions.scala     |  2 +-
 2 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/DateUtils.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/DateUtils.scala
index ccb2e91f57a..eaa5c2c22a0 100644
--- a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/DateUtils.scala
+++ b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/DateUtils.scala
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2020-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.
@@ -217,15 +217,17 @@ object DateUtils {
   def tagAndGetCudfFormat(
       meta: RapidsMeta[_, _, _],
       sparkFormat: String,
-      parseString: Boolean): String = {
+      parseString: Boolean,
+      inputFormat: String = null): String = {
+    val formatToConvert = if (inputFormat != null) inputFormat else sparkFormat
     var strfFormat: String = null
     if (GpuOverrides.getTimeParserPolicy == LegacyTimeParserPolicy) {
       try {
         // try and convert the format to cuDF format - this will throw an exception if
         // the format contains unsupported characters or words
-        strfFormat = toStrf(sparkFormat, parseString)
+        strfFormat = toStrf(formatToConvert, parseString)
         // format parsed ok but we have no 100% compatible formats in LEGACY mode
-        if (GpuToTimestamp.LEGACY_COMPATIBLE_FORMATS.contains(sparkFormat)) {
+        if (GpuToTimestamp.LEGACY_COMPATIBLE_FORMATS.contains(formatToConvert)) {
           // LEGACY support has a number of issues that mean we cannot guarantee
           // compatibility with CPU
           // - we can only support 4 digit years but Spark supports a wider range
@@ -249,9 +251,9 @@ object DateUtils {
       try {
         // try and convert the format to cuDF format - this will throw an exception if
         // the format contains unsupported characters or words
-        strfFormat = toStrf(sparkFormat, parseString)
+        strfFormat = toStrf(formatToConvert, parseString)
         // format parsed ok, so it is either compatible (tested/certified) or incompatible
-        if (!GpuToTimestamp.CORRECTED_COMPATIBLE_FORMATS.contains(sparkFormat) &&
+        if (!GpuToTimestamp.CORRECTED_COMPATIBLE_FORMATS.contains(formatToConvert) &&
           !meta.conf.incompatDateFormats) {
           meta.willNotWorkOnGpu(s"CORRECTED format '$sparkFormat' on the GPU is not guaranteed " +
             s"to produce the same results as Spark on CPU. Set " +
diff --git a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/datetimeExpressions.scala b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/datetimeExpressions.scala
index 9b59c1772c4..36093569075 100644
--- a/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/datetimeExpressions.scala
+++ b/sql-plugin/src/main/scala/org/apache/spark/sql/rapids/datetimeExpressions.scala
@@ -1033,7 +1033,7 @@ class FromUnitTimeMeta(a: FromUnixTime,
             tempFormat
           }.getOrElse(sparkFormat)
         strfFormat = DateUtils.tagAndGetCudfFormat(this,
-          inputFormat, a.left.dataType == DataTypes.StringType)
+          sparkFormat, a.left.dataType == DataTypes.StringType, inputFormat)
       case None =>
         willNotWorkOnGpu("format has to be a string literal")
     }