implement at_timezone

facebookincubator · Nov 25, 2023 · cc48163 · cc48163
1 parent fab2181
commit cc48163
Show file tree

Hide file tree

Showing 3 changed files with 97 additions and 0 deletions.
diff --git a/velox/functions/prestosql/DateTimeFunctions.h b/velox/functions/prestosql/DateTimeFunctions.h
@@ -15,7 +15,12 @@
  */
 #pragma once
 
+#include <time.h>
+#include <chrono>
+#include <string>
 #include <string_view>
+#include "velox/external/date/date.h"
+#include "velox/external/date/tz.h"
 #include "velox/functions/lib/DateTimeFormatter.h"
 #include "velox/functions/lib/TimeUtils.h"
 #include "velox/functions/prestosql/DateTimeImpl.h"
@@ -466,6 +471,36 @@ struct TimestampMinusIntervalDayTime {
   }
 };
 
+template <typename T>
+struct TimestampAtTimezoneFunction : public TimestampWithTimezoneSupport<T> {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE void call(
+      out_type<TimestampWithTimezone>& result,
+      const arg_type<Timestamp>& ts,
+      const arg_type<Varchar>& timeZone) {
+    auto timeZoneStr = std::string_view(timeZone.data(), timeZone.size());
+    auto zone = date::locate_zone(timeZoneStr);
+    std::chrono::system_clock::time_point tp{
+        std::chrono::seconds{ts.getSeconds()}};
+
+    // We are to interpret the input timestamp as being at GMT;
+    // Africa/Abidjan has offset +00:00; use this timezone to establish GMT as
+    // our starting offset.
+    auto startZonedTime = date::make_zoned("Africa/Abidjan", tp);
+
+    // Find corresponding time at input timezone.
+    auto adjustedZonedTime = date::make_zoned(zone, startZonedTime);
+    auto localTime = adjustedZonedTime.get_local_time();
+    auto adjustedSeconds = std::chrono::duration_cast<std::chrono::seconds>(
+                               localTime.time_since_epoch())
+                               .count();
+
+    result.template get_writer_at<0>() = adjustedSeconds;
+    result.template get_writer_at<1>() = util::getTimeZoneID(timeZoneStr);
+  }
+};
+
 template <typename T>
 struct DayOfWeekFunction : public InitSessionTimezone<T>,
                            public TimestampWithTimezoneSupport<T> {

diff --git a/velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp b/velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp
@@ -170,6 +170,11 @@ void registerSimpleFunctions(const std::string& prefix) {
   registerFunction<DateParseFunction, Timestamp, Varchar, Varchar>(
       {prefix + "date_parse"});
   registerFunction<CurrentDateFunction, Date>({prefix + "current_date"});
+  registerFunction<
+      TimestampAtTimezoneFunction,
+      TimestampWithTimezone,
+      Timestamp,
+      Varchar>({prefix + "timezone"});
 }
 } // namespace
 

diff --git a/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp b/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp
@@ -120,6 +120,26 @@ class DateTimeFunctionsTest : public functions::test::FunctionBaseTest {
         rowVector->children()[1]->as<SimpleVector<int16_t>>()->valueAt(0)};
   }
 
+  std::optional<TimestampWithTimezone> timestampAtTimezone(
+      const std::optional<Timestamp> t1,
+      const std::optional<std::string>& t2) {
+    auto resultVector = evaluate(
+        "timezone(c0, c1)",
+        makeRowVector(
+            {makeNullableFlatVector<Timestamp>({t1}),
+             makeNullableFlatVector<std::string>({t2})}));
+    EXPECT_EQ(1, resultVector->size());
+
+    if (resultVector->isNullAt(0)) {
+      return std::nullopt;
+    }
+
+    auto rowVector = resultVector->as<RowVector>();
+    return TimestampWithTimezone{
+        rowVector->children()[0]->as<SimpleVector<int64_t>>()->valueAt(0),
+        rowVector->children()[1]->as<SimpleVector<int16_t>>()->valueAt(0)};
+  }
+
   std::optional<Timestamp> dateParse(
       const std::optional<std::string>& input,
       const std::optional<std::string>& format) {
@@ -938,6 +958,43 @@ TEST_F(DateTimeFunctionsTest, minusTimestamp) {
       "Could not convert Timestamp(-9223372036854776, 0) to milliseconds");
 }
 
+TEST_F(DateTimeFunctionsTest, timestampAtTimezoneTest) {
+  // GMT -06:00, -21600 sec
+  EXPECT_EQ(
+      TimestampWithTimezone(
+          1500101514 - 21600, util::getTimeZoneID("America/Boise")),
+      timestampAtTimezone(Timestamp(1500101514, 0), "America/Boise"));
+
+  // GMT +03:00, +10800 sec
+  EXPECT_EQ(
+      TimestampWithTimezone(
+          1500101514 + 10800, util::getTimeZoneID("Asia/Baghdad")),
+      timestampAtTimezone(Timestamp(1500101514, 0), "Asia/Baghdad"));
+
+  // GMT +0:00
+  EXPECT_EQ(
+      TimestampWithTimezone(1500101514, util::getTimeZoneID("Africa/Bamako")),
+      timestampAtTimezone(Timestamp(1500101514, 0), "Africa/Bamako"));
+
+  // GMT -07:00, -25200 sec
+  EXPECT_EQ(
+      TimestampWithTimezone(
+          1513299114 - 25200, util::getTimeZoneID("America/Boise")),
+      timestampAtTimezone(Timestamp(1513299114, 0), "America/Boise"));
+
+  // GMT +01:00, +3600 sec
+  EXPECT_EQ(
+      TimestampWithTimezone(
+          507034293 + 3600, util::getTimeZoneID("Europe/Prague")),
+      timestampAtTimezone(Timestamp(507034293, 0), "Europe/Prague"));
+
+  // GMT +05:45, +20700 sec
+  EXPECT_EQ(
+      TimestampWithTimezone(
+          75325423914 + 20700, util::getTimeZoneID("Asia/Kathmandu")),
+      timestampAtTimezone(Timestamp(75325423914, 0), "Asia/Kathmandu"));
+}
+
 TEST_F(DateTimeFunctionsTest, dayOfMonthTimestampWithTimezone) {
   EXPECT_EQ(
       31,