Add timezone_hour Presto function (facebookincubator#4520)

Summary: Resolves facebookincubator#4469 Pull Request resolved: facebookincubator#4520 Reviewed By: mbasmanova Differential Revision: D45188569 Pulled By: Yuhta fbshipit-source-id: 5760c942ac6942042484d1754c6ce42f30871d2d
zuochunwei · Apr 26, 2023 · 300370f · 300370f
1 parent c125342
commit 300370f
Show file tree

Hide file tree

Showing 4 changed files with 62 additions and 1 deletion.
diff --git a/velox/docs/functions/presto/datetime.rst b/velox/docs/functions/presto/datetime.rst
@@ -202,6 +202,10 @@ This behavior is different from Presto Java that allows arbitrary large timestam
 
     Returns the second of the minute from ``x``.
 
+.. function:: timezone_hour(timestamp) -> bigint
+
+    Returns the hour of the time zone offset from ``timestamp``.
+
 .. function:: week(x) -> bigint
 
     Returns the `ISO-Week`_ of the year from x. The value ranges from ``1`` to ``53``.

diff --git a/velox/functions/prestosql/DateTimeFunctions.h b/velox/functions/prestosql/DateTimeFunctions.h
@@ -1080,4 +1080,29 @@ struct ParseDateTimeFunction {
   }
 };
 
+template <typename T>
+struct TimeZoneHourFunction : public TimestampWithTimezoneSupport<T> {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE void call(
+      int64_t& result,
+      const arg_type<TimestampWithTimezone>& input) {
+    // Convert timestampWithTimezone input to a timestamp representing the
+    // moment at the zone in timestampWithTimezone.
+    Timestamp inputTimeStamp = this->toTimestamp(input);
+
+    // Get the given timezone name
+    auto timezone = util::getTimeZoneName(*input.template at<1>());
+
+    auto* timezonePtr = date::locate_zone(timezone);
+
+    // Create a copy of inputTimeStamp and convert it to GMT
+    auto gmtTimeStamp = inputTimeStamp;
+    gmtTimeStamp.toGMT(*timezonePtr);
+
+    // Get offset in seconds with GMT and convert to hour
+    result = (inputTimeStamp.getSeconds() - gmtTimeStamp.getSeconds()) / 3600;
+  }
+};
+
 } // namespace facebook::velox::functions
diff --git a/velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp b/velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp
@@ -32,7 +32,8 @@ void registerSimpleFunctions(const std::string& prefix) {
   registerFunction<DateFunction, Date, Timestamp>({prefix + "date"});
   registerFunction<DateFunction, Date, TimestampWithTimezone>(
       {prefix + "date"});
-
+  registerFunction<TimeZoneHourFunction, int64_t, TimestampWithTimezone>(
+      {prefix + "timezone_hour"});
   registerFunction<YearFunction, int64_t, Timestamp>({prefix + "year"});
   registerFunction<YearFunction, int64_t, Date>({prefix + "year"});
   registerFunction<YearFunction, int64_t, TimestampWithTimezone>(

diff --git a/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp b/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp
@@ -3007,3 +3007,34 @@ TEST_F(DateTimeFunctionsTest, dateFunctionTimestampWithTimezone) {
       dateFunction(
           (-18297 * kSecondsInDay + 6 * 3'600) * 1'000, "America/Los_Angeles"));
 }
+
+TEST_F(DateTimeFunctionsTest, timeZoneHour) {
+  const auto timezone_hour = [&](const char* time, const char* timezone) {
+    Timestamp ts = util::fromTimestampString(time);
+    auto timestamp = ts.toMillis();
+    auto hour = evaluateWithTimestampWithTimezone<int64_t>(
+                    "timezone_hour(c0)", timestamp, timezone)
+                    .value();
+    return hour;
+  };
+
+  // Asia/Kolkata - should return 5 throughout the year
+  EXPECT_EQ(5, timezone_hour("2023-01-01 03:20:00", "Asia/Kolkata"));
+  EXPECT_EQ(5, timezone_hour("2023-06-01 03:20:00", "Asia/Kolkata"));
+  // America/Los_Angeles - Day light savings is from March 12 to Nov 5
+  EXPECT_EQ(-8, timezone_hour("2023-03-11 12:00:00", "America/Los_Angeles"));
+  EXPECT_EQ(-8, timezone_hour("2023-03-12 02:30:00", "America/Los_Angeles"));
+  EXPECT_EQ(-7, timezone_hour("2023-03-13 12:00:00", "America/Los_Angeles"));
+  EXPECT_EQ(-7, timezone_hour("2023-11-05 01:30:00", "America/Los_Angeles"));
+  EXPECT_EQ(-8, timezone_hour("2023-12-05 01:30:00", "America/Los_Angeles"));
+  // Different time with same date
+  EXPECT_EQ(-4, timezone_hour("2023-01-01 03:20:00", "Canada/Atlantic"));
+  EXPECT_EQ(-4, timezone_hour("2023-01-01 10:00:00", "Canada/Atlantic"));
+  // Invalid inputs
+  VELOX_ASSERT_THROW(
+      timezone_hour("invalid_date", "Canada/Atlantic"),
+      "Unable to parse timestamp value: \"invalid_date\", expected format is (YYYY-MM-DD HH:MM:SS[.MS])");
+  VELOX_ASSERT_THROW(
+      timezone_hour("123456", "Canada/Atlantic"),
+      "Unable to parse timestamp value: \"123456\", expected format is (YYYY-MM-DD HH:MM:SS[.MS])");
+}