Initial commit

facebookincubator · May 6, 2024 · 1fb948d · 1fb948d
1 parent 9fde0a2
commit 1fb948d
Show file tree

Hide file tree

Showing 6 changed files with 66 additions and 15 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"
 

diff --git a/velox/expression/CastExpr.cpp b/velox/expression/CastExpr.cpp
@@ -26,6 +26,7 @@
 #include "velox/expression/ScopedVarSetter.h"
 #include "velox/external/date/tz.h"
 #include "velox/functions/lib/RowsTranslationUtil.h"
+#include "velox/functions/lib/TimeUtils.h"
 #include "velox/type/Type.h"
 #include "velox/vector/ComplexVector.h"
 #include "velox/vector/FunctionVector.h"
@@ -181,11 +182,7 @@ VectorPtr CastExpr::castFromDate(
     case TypeKind::TIMESTAMP: {
       static const int64_t kMillisPerDay{86'400'000};
       const auto& queryConfig = context.execCtx()->queryCtx()->queryConfig();
-      const auto sessionTzName = queryConfig.sessionTimezone();
-      const auto* timeZone =
-          (queryConfig.adjustTimestampToTimezone() && !sessionTzName.empty())
-          ? date::locate_zone(sessionTzName)
-          : nullptr;
+      const auto* timeZone = functions::getTimeZoneFromConfig(queryConfig);
       auto* resultFlatVector = castResult->as<FlatVector<Timestamp>>();
       applyToSelectedNoThrowLocal(context, rows, castResult, [&](int row) {
         auto timestamp = Timestamp::fromMillis(
@@ -233,9 +230,8 @@ 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);
+      const auto* timeZone = functions::getTimeZoneFromConfig(queryConfig);
+      if (timeZone) {
         castTimestampToDate<true>(rows, input, context, castResult, timeZone);
       } else {
         castTimestampToDate<false>(rows, input, context, castResult);

diff --git a/velox/expression/PrestoCastHooks.cpp b/velox/expression/PrestoCastHooks.cpp
@@ -16,6 +16,7 @@
 
 #include "velox/expression/PrestoCastHooks.h"
 #include "velox/external/date/tz.h"
+#include "velox/functions/lib/TimeUtils.h"
 #include "velox/type/TimestampConversion.h"
 
 namespace facebook::velox::exec {
@@ -25,10 +26,7 @@ PrestoCastHooks::PrestoCastHooks(const core::QueryConfig& config)
   if (!legacyCast_) {
     options_.zeroPaddingYear = true;
     options_.dateTimeSeparator = ' ';
-    const auto sessionTzName = config.sessionTimezone();
-    if (config.adjustTimestampToTimezone() && !sessionTzName.empty()) {
-      options_.timeZone = date::locate_zone(sessionTzName);
-    }
+    options_.timeZone = functions::getTimeZoneFromConfig(config);
   }
 }
 

diff --git a/velox/expression/tests/CastExprTest.cpp b/velox/expression/tests/CastExprTest.cpp
@@ -609,6 +609,26 @@ TEST_F(CastExprTest, stringToTimestamp) {
       Timestamp(946729316, 0),
   };
   testCast<std::string, Timestamp>("timestamp", input, expected);
+
+  setTimezone("GMT+8");
+  testCast<std::string, Timestamp>(
+      "timestamp",
+      {
+          "1970-01-01",
+          "1970-01-01 08:00:00",
+          "1970-01-01 00:00:00",
+          "1970-01-01 08:00:11",
+          "1970-01-01 09:00:00",
+          std::nullopt,
+      },
+      {
+          Timestamp(-28800, 0),
+          Timestamp(0, 0),
+          Timestamp(-28800, 0),
+          Timestamp(11, 0),
+          Timestamp(3600, 0),
+          std::nullopt,
+      });
 }
 
 TEST_F(CastExprTest, timestampToString) {
@@ -2397,8 +2417,8 @@ class TestingDictionaryToFewerRowsFunction : public exec::VectorFunction {
       exec::EvalCtx& context,
       VectorPtr& result) const override {
     const auto size = rows.size();
-    auto indices =
-        makeIndices(size, [](auto /*row*/) { return 0; }, context.pool());
+    auto indices = makeIndices(
+        size, [](auto /*row*/) { return 0; }, context.pool());
 
     result = BaseVector::wrapInDictionary(nullptr, indices, size, args[0]);
   }

diff --git a/velox/functions/lib/TimeUtils.h b/velox/functions/lib/TimeUtils.h
@@ -26,12 +26,29 @@ inline constexpr int64_t kSecondsInDay = 86'400;
 inline constexpr int64_t kDaysInWeek = 7;
 extern const folly::F14FastMap<std::string, int8_t> kDayOfWeekNames;
 
+/// Format timezone to make it compatible with date::locate_zone.
+/// For example, converts "GMT+8" to "Etc/GMT-8".
+/// Here is the list of IANA timezone names:
+/// https://en.wikipedia.org/wiki/List_of_tz_database_time_zones
+FOLLY_ALWAYS_INLINE std::string formatTimezone(const std::string& timezone) {
+  if (timezone.find("GMT") == 0 && timezone.size() > 4) {
+    if (timezone[3] == '+') {
+      return "Etc/GMT-" + timezone.substr(4);
+    }
+    if (timezone[3] == '-') {
+      return "Etc/GMT+" + timezone.substr(4);
+    }
+  }
+  return timezone;
+}
+
 FOLLY_ALWAYS_INLINE const date::time_zone* getTimeZoneFromConfig(
     const core::QueryConfig& config) {
   if (config.adjustTimestampToTimezone()) {
     auto sessionTzName = config.sessionTimezone();
     if (!sessionTzName.empty()) {
-      return date::locate_zone(sessionTzName);
+      // locate_zone throws runtime_error if the timezone couldn't be found.
+      return date::locate_zone(formatTimezone(sessionTzName));
     }
   }
   return nullptr;

diff --git a/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp b/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp
@@ -3341,6 +3341,25 @@ TEST_F(DateTimeFunctionsTest, dateFormat) {
           fromTimestampString("-2000-02-29 00:00:00.987"),
           "%Y-%m-%d %H:%i:%s.%f"));
 
+  setQueryTimeZone("GMT-8");
+
+  EXPECT_EQ(
+      "1969-12-31", dateFormat(fromTimestampString("1970-01-01"), "%Y-%m-%d"));
+  EXPECT_EQ(
+      "2000-02-28 04:00:00 PM",
+      dateFormat(
+          fromTimestampString("2000-02-29 00:00:00.987"), "%Y-%m-%d %r"));
+  EXPECT_EQ(
+      "2000-02-28 16:00:00.987000",
+      dateFormat(
+          fromTimestampString("2000-02-29 00:00:00.987"),
+          "%Y-%m-%d %H:%i:%s.%f"));
+  EXPECT_EQ(
+      "-2000-02-28 16:00:00.987000",
+      dateFormat(
+          fromTimestampString("-2000-02-29 00:00:00.987"),
+          "%Y-%m-%d %H:%i:%s.%f"));
+
   // User format errors or unsupported errors.
   const auto timestamp = fromTimestampString("-2000-02-29 00:00:00.987");
   VELOX_ASSERT_THROW(