[WIP] Move timezone check to each operator [databricks]

integration_tests/src/main/python/date_time_test.py

@@ -13,8 +13,9 @@
 # limitations under the License.
 
 import pytest
-from asserts import assert_gpu_and_cpu_are_equal_collect, assert_gpu_fallback_collect, assert_gpu_and_cpu_error
+
 from data_gen import *
+from asserts import assert_gpu_and_cpu_are_equal_collect, assert_gpu_and_cpu_are_equal_sql, assert_gpu_fallback_collect, assert_gpu_and_cpu_error, assert_spark_exception, with_gpu_session
 from datetime import date, datetime, timezone
 from marks import ignore_order, incompat, allow_non_gpu
 from pyspark.sql.types import *
@@ -558,3 +559,87 @@ def test_timestamp_millis_long_overflow():
 def test_timestamp_micros(data_gen):
     assert_gpu_and_cpu_are_equal_collect(
         lambda spark : unary_op_df(spark, data_gen).selectExpr("timestamp_micros(a)"))
+
+
+# used by timezone test cases
+def get_timezone_df(spark):
+    schema = StructType([
+        StructField("ts_str_col", StringType()),
+        StructField("long_col", LongType()),
+        StructField("ts_col", TimestampType()),
+        StructField("date_col", DateType()),
+        StructField("date_str_col", StringType()),
+    ])
+    data = [
+        ('1970-01-01 00:00:00', 0, datetime(1971, 1, 1), date(1971, 1, 1), '1971-01-01'),
+        ('1970-01-01 00:00:00', 0, datetime(1971, 1, 1), date(1971, 1, 1), '1971-01-01'),
+    ]
+    return spark.createDataFrame(SparkContext.getOrCreate().parallelize(data),schema)
+
+# used by timezone test cases, specify all the sqls that will be impacted by non-utc timezone
+time_zone_sql_conf_pairs = [
+    ("select minute(ts_col) from tab", {}),
+    ("select second(ts_col) from tab", {}),
+    ("select hour(ts_col) from tab", {}),
+    ("select date_col + (interval 10 days 3 seconds) from tab", {}),
+    ("select date_format(ts_col, 'yyyy-MM-dd HH:mm:ss') from tab", {}),
+    ("select unix_timestamp(ts_col) from tab", {"spark.rapids.sql.improvedTimeOps.enabled": "true"}),
+    ("select to_unix_timestamp(ts_str_col) from tab", {"spark.rapids.sql.improvedTimeOps.enabled": "false"}),
+    ("select to_unix_timestamp(ts_col) from tab", {"spark.rapids.sql.improvedTimeOps.enabled": "true"}), 
+    ("select to_date(date_str_col, 'yyyy-MM-dd') from tab", {}), # test GpuGetTimestamp
+    ("select to_date(date_str_col) from tab", {}),
+    ("select from_unixtime(long_col, 'yyyy-MM-dd HH:mm:ss') from tab", {}),
+    ("select cast(ts_col as string) from tab", {}), # cast
+    ("select cast(ts_col as date) from tab", {}), # cast
+    ("select cast(date_col as TIMESTAMP) from tab", {}), # cast
+    ("select to_timestamp(ts_str_col) from tab", {"spark.rapids.sql.improvedTimeOps.enabled": "false"}), 
+    ("select to_timestamp(ts_str_col) from tab", {"spark.rapids.sql.improvedTimeOps.enabled": "true"}), 
+    ]
+
+
+@allow_non_gpu("ProjectExec")
+@pytest.mark.parametrize('sql, extra_conf', time_zone_sql_conf_pairs)
+def test_timezone_for_operators_with_non_utc(sql, extra_conf):
+    # timezone is non-utc, should fallback to CPU
+    timezone_conf = {"spark.sql.session.timeZone": "+08:00", 
+            "spark.rapids.sql.hasExtendedYearValues": "false",
+            "spark.rapids.sql.castStringToTimestamp.enabled": "true"}
+    all_conf = copy_and_update(timezone_conf, extra_conf)
+    def gen_sql_df(spark):
+        df = get_timezone_df(spark)
+        df.createOrReplaceTempView("tab")
+        return spark.sql(sql)
+    assert_gpu_fallback_collect(gen_sql_df, "ProjectExec", all_conf)
+
+
+@pytest.mark.parametrize('sql, conf', time_zone_sql_conf_pairs)
+def test_timezone_for_operators_with_utc(sql, conf):
+    # timezone is utc, should be supported by GPU
+    timezone_conf = {"spark.sql.session.timeZone": "UTC", 
+            "spark.rapids.sql.hasExtendedYearValues": "false",
+            "spark.rapids.sql.castStringToTimestamp.enabled": "true",}
+    conf = copy_and_update(timezone_conf, conf)
+    def gen_sql_df(spark):
+        df = get_timezone_df(spark)
+        df.createOrReplaceTempView("tab")
+        return spark.sql(sql)
+    assert_gpu_and_cpu_are_equal_collect(gen_sql_df, conf)
+
+
+@allow_non_gpu("ProjectExec")
+def test_timezone_for_operator_from_utc_timestamp_with_non_utc():
+    # timezone is non-utc, should fallback to CPU
+    def gen_sql_df(spark):
+        df = get_timezone_df(spark)
+        df.createOrReplaceTempView("tab")
+        return spark.sql("select from_utc_timestamp(ts_col, '+08:00') from tab")
+    assert_gpu_fallback_collect(gen_sql_df, "ProjectExec")
+
+
+def test_timezone_for_operator_from_utc_timestamp_with_utc():
+    # timezone is utc, should be supported by GPU
+    def gen_sql_df(spark):
+        df = get_timezone_df(spark)
+        df.createOrReplaceTempView("tab")
+        return spark.sql("select from_utc_timestamp(ts_col, '+00:00') from tab").collect()
+    with_gpu_session(gen_sql_df)

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

@@ -90,6 +90,9 @@ abstract class CastExprMetaBase[INPUT <: UnaryExpression with TimeZoneAwareExpre
     recursiveTagExprForGpuCheck()
   }
 
+  // tag time zone by Cast itself, do not delegate to parent class
+  override def tagTimeZoneBySelf: Boolean = true
+
   protected def recursiveTagExprForGpuCheck(
       fromDataType: DataType = fromType,
       toDataType: DataType = toType,
@@ -177,9 +180,6 @@ abstract class CastExprMetaBase[INPUT <: UnaryExpression with TimeZoneAwareExpre
   def buildTagMessage(entry: ConfEntry[_]): String = {
     s"${entry.doc}. To enable this operation on the GPU, set ${entry.key} to true."
   }
-
-  // timezone tagging in type checks is good enough, so always false
-  override protected val needTimezoneTagging: Boolean = false
 }
 
 object CastOptions {

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

@@ -669,9 +669,7 @@ object GpuOverrides extends Logging {
       case FloatType => true
       case DoubleType => true
       case DateType => true
-      case TimestampType =>
-        TypeChecks.areTimestampsSupported(ZoneId.systemDefault()) &&
-        TypeChecks.areTimestampsSupported(SQLConf.get.sessionLocalTimeZone)
+      case TimestampType => true
       case StringType => true
       case dt: DecimalType if allowDecimal => dt.precision <= DType.DECIMAL64_MAX_PRECISION
       case NullType => allowNull

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

@@ -374,13 +374,12 @@ class RapidsExecutorPlugin extends ExecutorPlugin with Logging {
         case Some(value) => ZoneId.of(value)
         case None => throw new RuntimeException(s"Driver time zone cannot be determined.")
       }
-      if (TypeChecks.areTimestampsSupported(driverTimezone)) {
-        val executorTimezone = ZoneId.systemDefault()
-        if (executorTimezone.normalized() != driverTimezone.normalized()) {
-          throw new RuntimeException(s" Driver and executor timezone mismatch. " +
-              s"Driver timezone is $driverTimezone and executor timezone is " +
-              s"$executorTimezone. Set executor timezone to $driverTimezone.")
-        }
+
+      val executorTimezone = ZoneId.systemDefault()
+      if (executorTimezone.normalized() != driverTimezone.normalized()) {
+        throw new RuntimeException(s" Driver and executor timezone mismatch. " +
+            s"Driver timezone is $driverTimezone and executor timezone is " +
+            s"$executorTimezone. Set executor timezone to $driverTimezone.")
       }
 
       GpuCoreDumpHandler.executorInit(conf, pluginContext)

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

@@ -22,7 +22,7 @@ import scala.collection.mutable
 
 import com.nvidia.spark.rapids.shims.{DistributionUtil, SparkShimImpl}
 
-import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, BinaryExpression, ComplexTypeMergingExpression, Expression, QuaternaryExpression, String2TrimExpression, TernaryExpression, TimeZoneAwareExpression, UnaryExpression, WindowExpression, WindowFunction}
+import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateFunction, ImperativeAggregate, TypedImperativeAggregate}
 import org.apache.spark.sql.catalyst.plans.physical.Partitioning
 import org.apache.spark.sql.catalyst.trees.TreeNodeTag
@@ -1071,37 +1071,37 @@ abstract class BaseExprMeta[INPUT <: Expression](
 
   val isFoldableNonLitAllowed: Boolean = false
 
-  /**
-   * Whether to tag a TimeZoneAwareExpression for timezone after all the other tagging
-   * is done.
-   * By default a TimeZoneAwareExpression always requires the timezone tagging, but
-   * there are some exceptions, e.g. 'Cast', who requires timezone tagging only when it
-   * has timezone sensitive type as input or output.
-   *
-   * Override this to match special cases.
-   */
-  protected def needTimezoneTagging: Boolean = {
-    // A TimeZoneAwareExpression with no timezone sensitive types as input/output will
-    // escape from the timezone tagging in the prior type checks. So ask for tagging here.
-    // e.g. 'UnixTimestamp' with 'DateType' as the input, timezone will be taken into
-    // account when converting a Date to a Long.
-    !(dataType +: childExprs.map(_.dataType)).exists(TypeChecks.isTimezoneSensitiveType)
-  }
-
   final override def tagSelfForGpu(): Unit = {
     if (wrapped.foldable && !GpuOverrides.isLit(wrapped) && !isFoldableNonLitAllowed) {
       willNotWorkOnGpu(s"Cannot run on GPU. Is ConstantFolding excluded? Expression " +
         s"$wrapped is foldable and operates on non literals")
     }
     rule.getChecks.foreach(_.tag(this))
     tagExprForGpu()
-    wrapped match {
-      case tzAware: TimeZoneAwareExpression if needTimezoneTagging =>
-        checkTimeZoneId(tzAware.zoneId)
-      case _ => // do nothing
+
+    // if expr is time zone aware and GPU does not support non UTC tz for this expr yet,
+    // ensure it's in UTC tz
+    if (!tagTimeZoneBySelf && isTimeZoneAwareExpr && !supportsNonUTCTimeZone) {
+      checkTimeZoneId(expr.asInstanceOf[TimeZoneAwareExpression].zoneId)
     }
   }
 
+  // if the wrapped expression is time zone aware
+  private final def isTimeZoneAwareExpr: Boolean = expr.isInstanceOf[TimeZoneAwareExpression]
+
+  /**
+   * whether the GPU supports non UTC time zone, for each expression that supports non UTC time
+   * zone, should override this method to return true
+   */
+  def supportsNonUTCTimeZone: Boolean = false
+
+  /**
+   * For cast expr or might other exprs, it's time zone aware, but time zone check can be skipped
+   * for some input/output types, like cast(int as long). For this kind of expr, should override
+   * this method and return true which means this Expr Meta should check time zone itself
+   */
+  def tagTimeZoneBySelf: Boolean = false
+
   /**
    * Called to verify that this expression will work on the GPU. For most expressions without
    * extra checks all of the checks should have already been done.

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

@@ -363,8 +363,7 @@ final class TypeSig private(
       case FloatType => check.contains(TypeEnum.FLOAT)
       case DoubleType => check.contains(TypeEnum.DOUBLE)
       case DateType => check.contains(TypeEnum.DATE)
-      case TimestampType if check.contains(TypeEnum.TIMESTAMP) =>
-          TypeChecks.areTimestampsSupported()
+      case TimestampType => check.contains(TypeEnum.TIMESTAMP)
       case StringType => check.contains(TypeEnum.STRING)
       case dt: DecimalType =>
           check.contains(TypeEnum.DECIMAL) &&
@@ -841,10 +840,6 @@ object TypeChecks {
       areTimestampsSupported(SQLConf.get.sessionLocalTimeZone)
   }
 
-  def isTimezoneSensitiveType(dataType: DataType): Boolean = {
-    dataType == TimestampType
-  }
-
   def timezoneNotSupportedString(dataType: DataType): String = {
     s"$dataType is not supported with timezone settings: (JVM:" +
       s" ${ZoneId.systemDefault()}, session: ${SQLConf.get.sessionLocalTimeZone})." +
@@ -1502,7 +1497,18 @@ class CastChecks extends ExprChecks {
 
   def gpuCanCast(from: DataType, to: DataType): Boolean = {
     val (checks, _) = getChecksAndSigs(from)
-    checks.isSupportedByPlugin(to)
+    checks.isSupportedByPlugin(to) && gpuCanCastConsiderTimezone(from, to)
+  }
+
+  def gpuCanCastConsiderTimezone(from: DataType, to: DataType) = {
+    // remove this check after non UTC timezone is supported for cast
+    (from, to) match {
+      case (StringType, TimestampType | DateType) => TypeChecks.areTimestampsSupported()
+      case (TimestampType | DateType, StringType) => TypeChecks.areTimestampsSupported()
+      case (DateType, TimestampType) => TypeChecks.areTimestampsSupported()
+      case (TimestampType, DateType) => TypeChecks.areTimestampsSupported()
+      case _ => true
+    }
   }
 }
 

sql-plugin/src/main/scala/org/apache/spark/sql/catalyst/expressions/rapids/Timestamp.scala

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2021-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.

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

@@ -1045,6 +1045,7 @@ class FromUTCTimestampExprMeta(
   extends BinaryExprMeta[FromUTCTimestamp](expr, conf, parent, rule) {
 
   override def tagExprForGpu(): Unit = {
+    // remove this check after non-UTC timezone is supported
     extractStringLit(expr.right) match {
       case None =>
         willNotWorkOnGpu("timezone input must be a literal string")