Fix more timezone handling

facebookincubator · Apr 24, 2024 · ef9f0c9 · ef9f0c9
1 parent 7876a66
commit ef9f0c9
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 24 deletions.
diff --git 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
@@ -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
@@ -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));
     }
   }