From ef9f0c9dcbd127d247d67b78e09b3f06efbcb684 Mon Sep 17 00:00:00 2001
From: PHILO-HE <feilong.he@intel.com>
Date: Wed, 24 Apr 2024 14:46:04 +0800
Subject: [PATCH] Fix more timezone handling

---
 velox/expression/CastExpr-inl.h | 38 +++++++++++++++------------------
 velox/expression/CastExpr.cpp   |  6 +++---
 velox/functions/lib/TimeUtils.h |  3 +++
 3 files changed, 23 insertions(+), 24 deletions(-)

diff --git a/velox/expression/CastExpr-inl.h b/velox/expression/CastExpr-inl.h
index 258f9638aaa5d..6360b607cc78b 100644
--- a/velox/expression/CastExpr-inl.h
+++ b/velox/expression/CastExpr-inl.h
@@ -22,6 +22,7 @@
 #include "velox/core/CoreTypeSystem.h"
 #include "velox/expression/StringWriter.h"
 #include "velox/external/date/tz.h"
+#include "velox/functions/lib/TimeUtils.h"
 #include "velox/type/Type.h"
 #include "velox/vector/SelectivityVector.h"
 
@@ -732,27 +733,22 @@ void CastExpr::applyCastPrimitives(
   // current GMT timezone to the user provided session timezone.
   if constexpr (ToKind == TypeKind::TIMESTAMP) {
     const auto& queryConfig = context.execCtx()->queryCtx()->queryConfig();
-    // If user explicitly asked us to adjust the timezone.
-    if (queryConfig.adjustTimestampToTimezone()) {
-      auto sessionTzName = queryConfig.sessionTimezone();
-      if (!sessionTzName.empty()) {
-        // When context.throwOnError is false, some rows will be marked as
-        // 'failed'. These rows should not be processed further. 'remainingRows'
-        // will contain a subset of 'rows' that have passed all the checks (e.g.
-        // keys are not nulls and number of keys and values is the same).
-        exec::LocalSelectivityVector remainingRows(context, rows);
-        context.deselectErrors(*remainingRows);
-
-        // locate_zone throws runtime_error if the timezone couldn't be found
-        // (so we're safe to dereference the pointer).
-        auto* timeZone = date::locate_zone(sessionTzName);
-        auto rawTimestamps = resultFlatVector->mutableRawValues();
-
-        applyToSelectedNoThrowLocal(
-            context, *remainingRows, result, [&](int row) {
-              rawTimestamps[row].toGMT(*timeZone);
-            });
-      }
+    auto* timeZone =
+        facebook::velox::functions::getTimeZoneFromConfig(queryConfig);
+    if (timeZone) {
+      // When context.throwOnError is false, some rows will be marked as
+      // 'failed'. These rows should not be processed further. 'remainingRows'
+      // will contain a subset of 'rows' that have passed all the checks (e.g.
+      // keys are not nulls and number of keys and values is the same).
+      exec::LocalSelectivityVector remainingRows(context, rows);
+      context.deselectErrors(*remainingRows);
+
+      auto rawTimestamps = resultFlatVector->mutableRawValues();
+
+      applyToSelectedNoThrowLocal(
+          context, *remainingRows, result, [&](int row) {
+            rawTimestamps[row].toGMT(*timeZone);
+          });
     }
   }
 }
diff --git a/velox/expression/CastExpr.cpp b/velox/expression/CastExpr.cpp
index 970d09e11d9ee..edb87047d526f 100644
--- a/velox/expression/CastExpr.cpp
+++ b/velox/expression/CastExpr.cpp
@@ -231,9 +231,9 @@ VectorPtr CastExpr::castToDate(
     }
     case TypeKind::TIMESTAMP: {
       const auto& queryConfig = context.execCtx()->queryCtx()->queryConfig();
-      auto sessionTzName = queryConfig.sessionTimezone();
-      if (queryConfig.adjustTimestampToTimezone() && !sessionTzName.empty()) {
-        auto* timeZone = date::locate_zone(sessionTzName);
+      auto* timeZone =
+          facebook::velox::functions::getTimeZoneFromConfig(queryConfig);
+      if (timeZone) {
         castTimestampToDate<true>(rows, input, context, castResult, timeZone);
       } else {
         castTimestampToDate<false>(rows, input, context, castResult);
diff --git a/velox/functions/lib/TimeUtils.h b/velox/functions/lib/TimeUtils.h
index 0028bfab4da37..6164119387e24 100644
--- a/velox/functions/lib/TimeUtils.h
+++ b/velox/functions/lib/TimeUtils.h
@@ -41,11 +41,14 @@ FOLLY_ALWAYS_INLINE std::string formatTimezone(const std::string& timezone) {
   return timezone;
 }
 
+// If user doesn't explicitly ask us to adjust the timezone or configured
+// session timezone is empty, returns nullptr.
 FOLLY_ALWAYS_INLINE const date::time_zone* getTimeZoneFromConfig(
     const core::QueryConfig& config) {
   if (config.adjustTimestampToTimezone()) {
     auto sessionTzName = config.sessionTimezone();
     if (!sessionTzName.empty()) {
+      // locate_zone throws runtime_error if the timezone couldn't be found.
       return date::locate_zone(formatTimezone(sessionTzName));
     }
   }