Fix GpuCast losing the timezoneId during canonialization (#5713)

Signed-off-by: Thomas Graves <tgraves@nvidia.com>
NVIDIA · Jun 2, 2022 · 5db6ca9 · 5db6ca9
1 parent 34c9aea
commit 5db6ca9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuCanonicalize.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuCanonicalize.scala
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -57,8 +57,7 @@ object GpuCanonicalize {
   def ignoreTimeZone(e: Expression): Expression = e match {
     case c: CastBase if c.timeZoneId.nonEmpty && !c.needsTimeZone =>
       c.withTimeZone(null)
-    case c: GpuCast if c.timeZoneId.nonEmpty =>
-      // TODO when we start to support time zones check for `&& !c.needsTimeZone`
+    case c: GpuCast if c.timeZoneId.nonEmpty && !c.needsTimeZone =>
       c.withTimeZone(null)
     case _ => e
   }
@@ -124,4 +123,4 @@ object GpuCanonicalize {
 
     case _ => e
   }
-}
+}
diff --git a/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuCast.scala b/sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuCast.scala
@@ -1669,7 +1669,7 @@ case class GpuCast(
   override lazy val resolved: Boolean =
     childrenResolved && checkInputDataTypes().isSuccess && (!needsTimeZone || timeZoneId.isDefined)
 
-  private[this] def needsTimeZone: Boolean = Cast.needsTimeZone(child.dataType, dataType)
+  def needsTimeZone: Boolean = Cast.needsTimeZone(child.dataType, dataType)
 
   override def sql: String = dataType match {
     // HiveQL doesn't allow casting to complex types. For logical plans translated from HiveQL,