Support hour minute second for non-UTC time zone (#10024)

* Support hour minute second for non-UTC time zone Signed-off-by: Haoyang Li <[email protected]> * Address comments Signed-off-by: Haoyang Li <[email protected]> * combine tests Signed-off-by: Haoyang Li <[email protected]> --------- Signed-off-by: Haoyang Li <[email protected]>
NVIDIA · Dec 13, 2023 · 4d7ec5b · 4d7ec5b
1 parent 3720faf
commit 4d7ec5b
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 15 deletions.
diff --git a/integration_tests/src/main/python/date_time_test.py b/integration_tests/src/main/python/date_time_test.py
@@ -135,20 +135,25 @@ def test_datediff(data_gen):
             'datediff(a, date(null))',
             'datediff(a, \'2016-03-02\')'))
 
-@allow_non_gpu(*non_utc_allow)
+hms_fallback = ['ProjectExec'] if not is_supported_time_zone() else []
+
+@allow_non_gpu(*hms_fallback)
 def test_hour():
     assert_gpu_and_cpu_are_equal_collect(
-        lambda spark : unary_op_df(spark, timestamp_gen).selectExpr('hour(a)'))
+        lambda spark : unary_op_df(spark, timestamp_gen).selectExpr('hour(a)'),
+        conf = {'spark.rapids.sql.nonUTC.enabled': True})
 
-@allow_non_gpu(*non_utc_allow)
+@allow_non_gpu(*hms_fallback)
 def test_minute():
     assert_gpu_and_cpu_are_equal_collect(
-        lambda spark : unary_op_df(spark, timestamp_gen).selectExpr('minute(a)'))
+        lambda spark : unary_op_df(spark, timestamp_gen).selectExpr('minute(a)'),
+        conf = {'spark.rapids.sql.nonUTC.enabled': True})
 
-@allow_non_gpu(*non_utc_allow)
+@allow_non_gpu(*hms_fallback)
 def test_second():
     assert_gpu_and_cpu_are_equal_collect(
-        lambda spark : unary_op_df(spark, timestamp_gen).selectExpr('second(a)'))
+        lambda spark : unary_op_df(spark, timestamp_gen).selectExpr('second(a)'),
+        conf = {'spark.rapids.sql.nonUTC.enabled': True})
 
 def test_quarter():
     assert_gpu_and_cpu_are_equal_collect(

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
@@ -1739,26 +1739,26 @@ object GpuOverrides extends Logging {
       ExprChecks.unaryProject(TypeSig.INT, TypeSig.INT,
         TypeSig.TIMESTAMP, TypeSig.TIMESTAMP),
       (hour, conf, p, r) => new UnaryExprMeta[Hour](hour, conf, p, r) {
-
-        override def convertToGpu(expr: Expression): GpuExpression = GpuHour(expr)
+        override def isTimeZoneSupported = true
+        override def convertToGpu(expr: Expression): GpuExpression = GpuHour(expr, hour.timeZoneId)
       }),
     expr[Minute](
       "Returns the minute component of the string/timestamp",
       ExprChecks.unaryProject(TypeSig.INT, TypeSig.INT,
         TypeSig.TIMESTAMP, TypeSig.TIMESTAMP),
       (minute, conf, p, r) => new UnaryExprMeta[Minute](minute, conf, p, r) {
-
+        override def isTimeZoneSupported = true
         override def convertToGpu(expr: Expression): GpuExpression =
-          GpuMinute(expr)
+          GpuMinute(expr, minute.timeZoneId)
       }),
     expr[Second](
       "Returns the second component of the string/timestamp",
       ExprChecks.unaryProject(TypeSig.INT, TypeSig.INT,
         TypeSig.TIMESTAMP, TypeSig.TIMESTAMP),
       (second, conf, p, r) => new UnaryExprMeta[Second](second, conf, p, r) {
-
+        override def isTimeZoneSupported = true
         override def convertToGpu(expr: Expression): GpuExpression =
-          GpuSecond(expr)
+          GpuSecond(expr, second.timeZoneId)
       }),
     expr[WeekDay](
       "Returns the day of the week (0 = Monday...6=Sunday)",

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
@@ -91,7 +91,14 @@ case class GpuMinute(child: Expression, timeZoneId: Option[String] = None)
     copy(timeZoneId = Option(timeZoneId))
 
   override protected def doColumnar(input: GpuColumnVector): ColumnVector =
-    input.getBase.minute()
+    if (GpuOverrides.isUTCTimezone(zoneId)) {
+      input.getBase.minute()
+    } else {
+      // Non-UTC time zone
+      withResource(GpuTimeZoneDB.fromUtcTimestampToTimestamp(input.getBase, zoneId)) {
+        shifted => shifted.minute()
+      }
+    }
 }
 
 case class GpuSecond(child: Expression, timeZoneId: Option[String] = None)
@@ -101,7 +108,14 @@ case class GpuSecond(child: Expression, timeZoneId: Option[String] = None)
     copy(timeZoneId = Option(timeZoneId))
 
   override protected def doColumnar(input: GpuColumnVector): ColumnVector =
-    input.getBase.second()
+    if (GpuOverrides.isUTCTimezone(zoneId)) {
+      input.getBase.second()
+    } else {
+      // Non-UTC time zone
+      withResource(GpuTimeZoneDB.fromUtcTimestampToTimestamp(input.getBase, zoneId)) {
+        shifted => shifted.second()
+      }
+    }
 }
 
 case class GpuHour(child: Expression, timeZoneId: Option[String] = None)
@@ -111,7 +125,14 @@ case class GpuHour(child: Expression, timeZoneId: Option[String] = None)
     copy(timeZoneId = Option(timeZoneId))
 
   override protected def doColumnar(input: GpuColumnVector): ColumnVector =
-    input.getBase.hour()
+    if (GpuOverrides.isUTCTimezone(zoneId)) {
+      input.getBase.hour()
+    } else {
+      // Non-UTC time zone
+      withResource(GpuTimeZoneDB.fromUtcTimestampToTimestamp(input.getBase, zoneId)) {
+        shifted => shifted.hour()
+      }
+    }
 }
 
 case class GpuYear(child: Expression) extends GpuDateUnaryExpression {