Support CPU path for from_utc_timestamp function with timezone

integration_tests/src/main/python/date_time_test.py

@@ -278,6 +278,14 @@ def test_from_utc_timestamp_unsupported_timezone_fallback(data_gen, time_zone):
     'FromUTCTimestamp')
 
 
+@pytest.mark.parametrize('time_zone', ["UTC", "Asia/Shanghai"], ids=idfn)
+@pytest.mark.parametrize('data_gen', [timestamp_gen], ids=idfn)
+def test_from_utc_timestamp_supported_timezones(data_gen, time_zone):
+    # Remove spark.rapids.test.CPU.timezone configuration when GPU kernel is ready to really test on GPU
+    assert_gpu_and_cpu_are_equal_collect(
+        lambda spark: unary_op_df(spark, data_gen).select(f.from_utc_timestamp(f.col('a'), time_zone)), conf = {"spark.rapids.test.CPU.timezone": "true"})
+
+
 @allow_non_gpu('ProjectExec')
 @pytest.mark.parametrize('data_gen', [timestamp_gen], ids=idfn)
 def test_unsupported_fallback_from_utc_timestamp(data_gen):

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

@@ -2055,6 +2055,12 @@ object RapidsConf {
     .booleanConf
     .createOptional
 
+  val TEST_USE_TIMEZONE_CPU_BACKEND = conf("spark.rapids.test.CPU.timezone")
+    .doc("Only for tests: verify for timezone related functions")
+    .internal()
+    .booleanConf
+    .createOptional
+
   private def printSectionHeader(category: String): Unit =
     println(s"\n### $category")
 

tests/src/test/scala/com/nvidia/spark/rapids/timezone/TimeZoneDB.scala → sql-plugin/src/main/scala/org/apache/spark/sql/rapids/TimeZoneDB.scala

@@ -14,17 +14,37 @@
  * limitations under the License.
  */
 
-package com.nvidia.spark.rapids.timezone
+package org.apache.spark.sql.rapids
 
 import java.time.ZoneId
 
 import ai.rapids.cudf.{ColumnVector, DType, HostColumnVector}
 import com.nvidia.spark.rapids.Arm.withResource
+import com.nvidia.spark.rapids.RapidsConf.TEST_USE_TIMEZONE_CPU_BACKEND
 
+import org.apache.spark.SparkEnv
 import org.apache.spark.sql.catalyst.util.DateTimeUtils
 
 object TimeZoneDB {
 
+  // Copied from Spark. Used to format time zone ID string with (+|-)h:mm and (+|-)hh:m
+  def getZoneId(timezoneId: String): ZoneId = {
+    val formattedZoneId = timezoneId
+      // To support the (+|-)h:mm format because it was supported before Spark 3.0.
+      .replaceFirst("(\\+|\\-)(\\d):", "$10$2:")
+      // To support the (+|-)hh:m format because it was supported before Spark 3.0.
+      .replaceFirst("(\\+|\\-)(\\d\\d):(\\d)$", "$1$2:0$3")
+    DateTimeUtils.getZoneId(formattedZoneId)
+  }
+
+  // Support fixed offset or no transition rule case
+  def isSupportedTimezone(timezoneId: String): Boolean = {
+    val rules = getZoneId(timezoneId).getRules
+    // CPU backend is just for test purpose
+    SparkEnv.get.conf.getBoolean(TEST_USE_TIMEZONE_CPU_BACKEND.key, false) ||
+      (rules.isFixedOffset || rules.getTransitionRules.isEmpty)
+  }
+
   def cacheDatabase(): Unit = {}
 
   /**
@@ -42,10 +62,14 @@ object TimeZoneDB {
       withResource(HostColumnVector.builder(DType.TIMESTAMP_MICROSECONDS, rowCount)) { builder =>
         var currRow = 0
         while (currRow < rowCount) {
-          val origin = input.getLong(currRow)
-          // Spark implementation
-          val dist = DateTimeUtils.toUTCTime(origin, zoneStr)
-          builder.append(dist)
+          if (input.isNull(currRow)) {
+            builder.appendNull()
+          } else {
+            val origin = input.getLong(currRow)
+            // Spark implementation
+            val dist = DateTimeUtils.toUTCTime(origin, zoneStr)
+            builder.append(dist)
+          }
           currRow += 1
         }
         withResource(builder.build()) { b =>
@@ -72,10 +96,14 @@ object TimeZoneDB {
       withResource(HostColumnVector.builder(DType.TIMESTAMP_MICROSECONDS, rowCount)) { builder =>
         var currRow = 0
         while (currRow < rowCount) {
-          val origin = input.getLong(currRow)
-          // Spark implementation
-          val dist = DateTimeUtils.fromUTCTime(origin, zoneStr)
-          builder.append(dist)
+          if(input.isNull(currRow)) {
+            builder.appendNull()
+          } else {
+            val origin = input.getLong(currRow)
+            // Spark implementation
+            val dist = DateTimeUtils.fromUTCTime(origin, zoneStr)
+            builder.append(dist)
+          }
           currRow += 1
         }
         withResource(builder.build()) { b =>
@@ -97,10 +125,14 @@ object TimeZoneDB {
       withResource(HostColumnVector.builder(DType.TIMESTAMP_DAYS, rowCount)) { builder =>
         var currRow = 0
         while (currRow < rowCount) {
-          val origin = input.getLong(currRow)
-          // Spark implementation
-          val dist = DateTimeUtils.microsToDays(origin, currentTimeZone)
-          builder.append(dist)
+          if (input.isNull(currRow)) {
+            builder.appendNull()
+          } else {
+            val origin = input.getLong(currRow)
+            // Spark implementation
+            val dist = DateTimeUtils.microsToDays(origin, currentTimeZone)
+            builder.append(dist)
+          }
           currRow += 1
         }
         withResource(builder.build()) { b =>
@@ -124,10 +156,14 @@ object TimeZoneDB {
       withResource(HostColumnVector.builder(DType.INT64, rowCount)) { builder =>
         var currRow = 0
         while (currRow < rowCount) {
-          val origin = input.getInt(currRow)
-          // Spark implementation
-          val dist = DateTimeUtils.daysToMicros(origin, desiredTimeZone)
-          builder.append(dist)
+          if (input.isNull(currRow)) {
+            builder.appendNull()
+          } else {
+            val origin = input.getInt(currRow)
+            // Spark implementation
+            val dist = DateTimeUtils.daysToMicros(origin, desiredTimeZone)
+            builder.append(dist)
+          }
           currRow += 1
         }
         withResource(builder.build()) { b =>

sql-plugin/src/main/scala/org/apache/spark/sql/rapids/datetimeExpressions.scala

@@ -1044,29 +1044,28 @@ class FromUTCTimestampExprMeta(
     rule: DataFromReplacementRule)
   extends BinaryExprMeta[FromUTCTimestamp](expr, conf, parent, rule) {
 
+  private[this] var timezoneId: ZoneId = null
+
   override def tagExprForGpu(): Unit = {
     extractStringLit(expr.right) match {
       case None =>
         willNotWorkOnGpu("timezone input must be a literal string")
       case Some(timezoneShortID) =>
         if (timezoneShortID != null) {
-          val utc = ZoneId.of("UTC").normalized
-          // This is copied from Spark, to convert `(+|-)h:mm` into `(+|-)0h:mm`.
-          val timezone = ZoneId.of(timezoneShortID.replaceFirst("(\\+|\\-)(\\d):", "$10$2:"),
-            ZoneId.SHORT_IDS).normalized
-
-          if (timezone != utc) {
-            willNotWorkOnGpu("only timezones equivalent to UTC are supported")
+          if (TimeZoneDB.isSupportedTimezone(timezoneShortID)) {
+            timezoneId = TimeZoneDB.getZoneId(timezoneShortID)
+          } else {
+            willNotWorkOnGpu(s"Not supported timezone type $timezoneShortID.")
           }
         }
     }
   }
 
   override def convertToGpu(timestamp: Expression, timezone: Expression): GpuExpression =
-    GpuFromUTCTimestamp(timestamp, timezone)
+    GpuFromUTCTimestamp(timestamp, timezone, timezoneId)
 }
 
-case class GpuFromUTCTimestamp(timestamp: Expression, timezone: Expression)
+case class GpuFromUTCTimestamp(timestamp: Expression, timezone: Expression, zoneId: ZoneId)
   extends GpuBinaryExpressionArgsAnyScalar
       with ImplicitCastInputTypes
       with NullIntolerant {
@@ -1078,8 +1077,7 @@ case class GpuFromUTCTimestamp(timestamp: Expression, timezone: Expression)
 
   override def doColumnar(lhs: GpuColumnVector, rhs: GpuScalar): ColumnVector = {
     if (rhs.getBase.isValid) {
-      // Just a no-op.
-      lhs.getBase.incRefCount()
+      TimeZoneDB.fromUtcTimestampToTimestamp(lhs.getBase, zoneId)
     } else {
       // All-null output column.
       GpuColumnVector.columnVectorFromNull(lhs.getRowCount.toInt, dataType)

tests/src/test/scala/com/nvidia/spark/rapids/timezone/TimeZoneSuite.scala

@@ -28,6 +28,7 @@ import com.nvidia.spark.rapids.SparkQueryCompareTestSuite
 
 import org.apache.spark.SparkConf
 import org.apache.spark.sql.{DataFrame, Row, SparkSession}
+import org.apache.spark.sql.rapids.TimeZoneDB
 import org.apache.spark.sql.types._
 
 class TimeZoneSuite extends SparkQueryCompareTestSuite {