Add string parameter support for unix_timestamp for non-UTC time zones (

#10210) Signed-off-by: Ferdinand Xu <[email protected]>
NVIDIA · Jan 19, 2024 · 8e8062d · 8e8062d
1 parent 787a74c
commit 8e8062d
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 18 deletions.
diff --git a/integration_tests/src/main/python/date_time_test.py b/integration_tests/src/main/python/date_time_test.py
@@ -411,7 +411,7 @@ def invalid_date_string_df(spark):
 
 @pytest.mark.parametrize('ansi_enabled', [True, False], ids=['ANSI_ON', 'ANSI_OFF'])
 @pytest.mark.parametrize('data_gen,date_form', str_date_and_format_gen, ids=idfn)
-@allow_non_gpu(*non_utc_tz_allow)
+@allow_non_gpu(*non_supported_tz_allow)
 def test_string_to_unix_timestamp(data_gen, date_form, ansi_enabled):
     assert_gpu_and_cpu_are_equal_collect(
         lambda spark : unary_op_df(spark, data_gen, seed=1).selectExpr("to_unix_timestamp(a, '{}')".format(date_form)),
@@ -425,7 +425,7 @@ def test_string_to_unix_timestamp_ansi_exception():
 
 @pytest.mark.parametrize('ansi_enabled', [True, False], ids=['ANSI_ON', 'ANSI_OFF'])
 @pytest.mark.parametrize('data_gen,date_form', str_date_and_format_gen, ids=idfn)
-@allow_non_gpu(*non_utc_tz_allow)
+@allow_non_gpu(*non_supported_tz_allow)
 def test_string_unix_timestamp(data_gen, date_form, ansi_enabled):
     assert_gpu_and_cpu_are_equal_collect(
         lambda spark : unary_op_df(spark, data_gen, seed=1).select(f.unix_timestamp(f.col('a'), date_form)),

diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuOverrides.scala
@@ -1703,13 +1703,7 @@ object GpuOverrides extends Logging {
             TypeSig.STRING)),
       (a, conf, p, r) => new UnixTimeExprMeta[ToUnixTimestamp](a, conf, p, r) {
         // String type is not supported yet for non-UTC timezone.
-        override def isTimeZoneSupported: Boolean = a.timeZoneId.forall { zoneID =>
-          a.left.dataType match {
-            case _: StringType => GpuOverrides.isUTCTimezone(zoneID)
-            case _ => true
-          }
-        }
-
+        override def isTimeZoneSupported = true
         override def convertToGpu(lhs: Expression, rhs: Expression): GpuExpression = {
           GpuToUnixTimestamp(lhs, rhs, sparkFormat, strfFormat, a.timeZoneId)
         }
@@ -1724,14 +1718,7 @@ object GpuOverrides extends Logging {
             .withPsNote(TypeEnum.STRING, "A limited number of formats are supported"),
             TypeSig.STRING)),
       (a, conf, p, r) => new UnixTimeExprMeta[UnixTimestamp](a, conf, p, r) {
-        // String type is not supported yet for non-UTC timezone.
-        override def isTimeZoneSupported: Boolean = a.timeZoneId.forall { zoneID =>
-            a.left.dataType match {
-              case _: StringType => GpuOverrides.isUTCTimezone(zoneID)
-              case _ => true
-            }
-        }
-
+        override def isTimeZoneSupported = true
         override def convertToGpu(lhs: Expression, rhs: Expression): GpuExpression = {
           GpuUnixTimestamp(lhs, rhs, sparkFormat, strfFormat, a.timeZoneId)
         }

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
@@ -874,7 +874,9 @@ abstract class GpuToTimestamp
         if (GpuOverrides.isUTCTimezone(zoneId)) {
           res
         } else {
-          GpuTimeZoneDB.fromTimestampToUtcTimestamp(res, zoneId)
+          withResource(res) { _ =>
+            GpuTimeZoneDB.fromTimestampToUtcTimestamp(res, zoneId)
+          }
         }
       case _: DateType =>
         timeZoneId match {