Add support for TS_WITH_TZ +/- INTERVAL Presto functions (facebookinc…

…ubator#10164)

Summary:
Pull Request resolved: facebookincubator#10164

Add the following functions:

- plus(ts_with_tz, interval)
- plus(interval, ts_with_tz)
- minus(ts_with_tz, interval)

Interval can be either INTERVAL_DAY_TIME or INTERVAL_YEAR_MONTH.

These functions are equivalent to date_add and share it's implementation, i.e.

plus(ts, interval '1' hour) == date_add('hour', 1, ts).

Also, plus(ts, interval) == plus(interval, ts) == minus(ts, -interval).

Note that Presto doesn't define minus(interval, ts).

date_add has existing bugs described in facebookincubator#10163

Newly added function share implementation with date_add, hence, suffer from the same bugs.

Reviewed By: amitkdutta

Differential Revision: D58466111

fbshipit-source-id: d5e1c47be67636397f19e3d3d062e18aa91ec7ad

velox/functions/prestosql/DateTimeFunctions.h

@@ -466,6 +466,22 @@ struct TimestampPlusInterval {
       const arg_type<IntervalYearMonth>& interval) {
     result = addToTimestamp(timestamp, DateTimeUnit::kMonth, interval);
   }
+
+  FOLLY_ALWAYS_INLINE void call(
+      out_type<TimestampWithTimezone>& result,
+      const arg_type<TimestampWithTimezone>& timestampWithTimezone,
+      const arg_type<IntervalDayTime>& interval) {
+    result = addToTimestampWithTimezone(
+        timestampWithTimezone, DateTimeUnit::kMillisecond, interval);
+  }
+
+  FOLLY_ALWAYS_INLINE void call(
+      out_type<TimestampWithTimezone>& result,
+      const arg_type<TimestampWithTimezone>& timestampWithTimezone,
+      const arg_type<IntervalYearMonth>& interval) {
+    result = addToTimestampWithTimezone(
+        timestampWithTimezone, DateTimeUnit::kMonth, interval);
+  }
 };
 
 template <typename T>
@@ -491,6 +507,22 @@ struct IntervalPlusTimestamp {
       const arg_type<Timestamp>& timestamp) {
     result = addToTimestamp(timestamp, DateTimeUnit::kMonth, interval);
   }
+
+  FOLLY_ALWAYS_INLINE void call(
+      out_type<TimestampWithTimezone>& result,
+      const arg_type<IntervalDayTime>& interval,
+      const arg_type<TimestampWithTimezone>& timestampWithTimezone) {
+    result = addToTimestampWithTimezone(
+        timestampWithTimezone, DateTimeUnit::kMillisecond, interval);
+  }
+
+  FOLLY_ALWAYS_INLINE void call(
+      out_type<TimestampWithTimezone>& result,
+      const arg_type<IntervalYearMonth>& interval,
+      const arg_type<TimestampWithTimezone>& timestampWithTimezone) {
+    result = addToTimestampWithTimezone(
+        timestampWithTimezone, DateTimeUnit::kMonth, interval);
+  }
 };
 
 template <typename T>
@@ -516,6 +548,22 @@ struct TimestampMinusInterval {
       const arg_type<IntervalYearMonth>& interval) {
     result = addToTimestamp(timestamp, DateTimeUnit::kMonth, -interval);
   }
+
+  FOLLY_ALWAYS_INLINE void call(
+      out_type<TimestampWithTimezone>& result,
+      const arg_type<TimestampWithTimezone>& timestampWithTimezone,
+      const arg_type<IntervalDayTime>& interval) {
+    result = addToTimestampWithTimezone(
+        timestampWithTimezone, DateTimeUnit::kMillisecond, -interval);
+  }
+
+  FOLLY_ALWAYS_INLINE void call(
+      out_type<TimestampWithTimezone>& result,
+      const arg_type<TimestampWithTimezone>& timestampWithTimezone,
+      const arg_type<IntervalYearMonth>& interval) {
+    result = addToTimestampWithTimezone(
+        timestampWithTimezone, DateTimeUnit::kMonth, -interval);
+  }
 };
 
 template <typename T>
@@ -1081,11 +1129,8 @@ struct DateAddFunction : public TimestampWithTimezoneSupport<T> {
       VELOX_UNSUPPORTED("integer overflow");
     }
 
-    auto finalTimeStamp = addToTimestamp(
-        this->toTimestamp(timestampWithTimezone), unit, (int32_t)value);
-    auto tzID = unpackZoneKeyId(timestampWithTimezone);
-    finalTimeStamp.toGMT(tzID);
-    result = pack(finalTimeStamp.toMillis(), tzID);
+    result =
+        addToTimestampWithTimezone(timestampWithTimezone, unit, (int32_t)value);
   }
 
   FOLLY_ALWAYS_INLINE void call(

velox/functions/prestosql/DateTimeImpl.h

@@ -20,6 +20,7 @@
 
 #include "velox/common/base/Doubles.h"
 #include "velox/external/date/date.h"
+#include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h"
 #include "velox/type/Timestamp.h"
 #include "velox/type/TimestampConversion.h"
 
@@ -198,6 +199,20 @@ FOLLY_ALWAYS_INLINE Timestamp addToTimestamp(
           timestamp.getNanos() % kNanosecondsInMillisecond);
 }
 
+FOLLY_ALWAYS_INLINE int64_t addToTimestampWithTimezone(
+    int64_t timestampWithTimezone,
+    const DateTimeUnit unit,
+    const int32_t value) {
+  auto timestamp = unpackTimestampUtc(timestampWithTimezone);
+  auto tzID = unpackZoneKeyId(timestampWithTimezone);
+  timestamp.toTimezone(tzID);
+
+  auto finalTimestamp = addToTimestamp(timestamp, unit, value);
+
+  finalTimestamp.toGMT(tzID);
+  return pack(finalTimestamp.toMillis(), tzID);
+}
+
 FOLLY_ALWAYS_INLINE int64_t diffTimestamp(
     const DateTimeUnit unit,
     const Timestamp& fromTimestamp,

velox/functions/prestosql/registration/DateTimeFunctionsRegistration.cpp

@@ -20,6 +20,52 @@
 
 namespace facebook::velox::functions {
 namespace {
+
+// Register timestamp + interval and interval + timestamp
+// functions for specified TTimestamp type and 2 supported interval types
+// (IntervalDayTime and IntervalYearMonth).
+// @tparam TTimestamp Timestamp or TimestampWithTimezone.
+template <typename TTimestamp>
+void registerTimestampPlusInterval(const std::string& name) {
+  registerFunction<
+      TimestampPlusInterval,
+      TTimestamp,
+      TTimestamp,
+      IntervalDayTime>({name});
+  registerFunction<
+      TimestampPlusInterval,
+      TTimestamp,
+      TTimestamp,
+      IntervalYearMonth>({name});
+  registerFunction<
+      IntervalPlusTimestamp,
+      TTimestamp,
+      IntervalDayTime,
+      TTimestamp>({name});
+  registerFunction<
+      IntervalPlusTimestamp,
+      TTimestamp,
+      IntervalYearMonth,
+      TTimestamp>({name});
+}
+
+// Register timestamp - IntervalYearMonth and timestamp - IntervalDayTime
+// functions for specified TTimestamp type.
+// @tparam TTimestamp Timestamp or TimestampWithTimezone.
+template <typename TTimestamp>
+void registerTimestampMinusInterval(const std::string& name) {
+  registerFunction<
+      TimestampMinusInterval,
+      TTimestamp,
+      TTimestamp,
+      IntervalDayTime>({name});
+  registerFunction<
+      TimestampMinusInterval,
+      TTimestamp,
+      TTimestamp,
+      IntervalYearMonth>({name});
+}
+
 void registerSimpleFunctions(const std::string& prefix) {
   // Date time functions.
   registerFunction<ToUnixtimeFunction, double, Timestamp>(
@@ -69,41 +115,18 @@ void registerSimpleFunctions(const std::string& prefix) {
       {prefix + "plus"});
   registerFunction<DatePlusInterval, Date, Date, IntervalYearMonth>(
       {prefix + "plus"});
-  registerFunction<
-      TimestampMinusInterval,
-      Timestamp,
-      Timestamp,
-      IntervalDayTime>({prefix + "minus"});
-  registerFunction<
-      TimestampMinusInterval,
-      Timestamp,
-      Timestamp,
-      IntervalYearMonth>({prefix + "minus"});
-  registerFunction<
-      TimestampPlusInterval,
-      Timestamp,
-      Timestamp,
-      IntervalDayTime>({prefix + "plus"});
-  registerFunction<
-      TimestampPlusInterval,
-      Timestamp,
-      Timestamp,
-      IntervalYearMonth>({prefix + "plus"});
-  registerFunction<
-      IntervalPlusTimestamp,
-      Timestamp,
-      IntervalDayTime,
-      Timestamp>({prefix + "plus"});
-  registerFunction<
-      IntervalPlusTimestamp,
-      Timestamp,
-      IntervalYearMonth,
-      Timestamp>({prefix + "plus"});
+
+  registerTimestampPlusInterval<Timestamp>({prefix + "plus"});
+  registerTimestampMinusInterval<Timestamp>({prefix + "minus"});
+  registerTimestampPlusInterval<TimestampWithTimezone>({prefix + "plus"});
+  registerTimestampMinusInterval<TimestampWithTimezone>({prefix + "minus"});
+
   registerFunction<
       TimestampMinusFunction,
       IntervalDayTime,
       Timestamp,
       Timestamp>({prefix + "minus"});
+
   registerFunction<DayFunction, int64_t, TimestampWithTimezone>(
       {prefix + "day", prefix + "day_of_month"});
   registerFunction<DayOfWeekFunction, int64_t, Timestamp>(

velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp

@@ -1048,6 +1048,61 @@ TEST_F(DateTimeFunctionsTest, plusMinusTimestampIntervalDayTime) {
       Timestamp(-1, 0), kLongMin, Timestamp(9223372036854774, 808000000));
 }
 
+TEST_F(DateTimeFunctionsTest, timestampWithTimeZonePlusIntervalDayTime) {
+  auto test = [&](const std::string& timestamp, int64_t interval) {
+    // ts + interval == interval + ts == ts - (-interval) ==
+    // date_add('millisecond', interval, ts).
+    auto plusResult =
+        evaluateOnce<std::string>(
+            "cast(plus(cast(c0 as timestamp with time zone), c1) as varchar)",
+            {VARCHAR(), INTERVAL_DAY_TIME()},
+            std::optional(timestamp),
+            std::optional(interval))
+            .value();
+
+    auto minusResult =
+        evaluateOnce<std::string>(
+            "cast(minus(cast(c0 as timestamp with time zone), c1) as varchar)",
+            {VARCHAR(), INTERVAL_DAY_TIME()},
+            std::optional(timestamp),
+            std::optional(-interval))
+            .value();
+
+    auto otherPlusResult =
+        evaluateOnce<std::string>(
+            "cast(plus(c1, cast(c0 as timestamp with time zone)) as varchar)",
+            {VARCHAR(), INTERVAL_DAY_TIME()},
+            std::optional(timestamp),
+            std::optional(interval))
+            .value();
+
+    auto dateAddResult =
+        evaluateOnce<std::string>(
+            "cast(date_add('millisecond', c1, cast(c0 as timestamp with time zone)) as varchar)",
+            std::optional(timestamp),
+            std::optional(interval))
+            .value();
+
+    VELOX_CHECK_EQ(plusResult, minusResult);
+    VELOX_CHECK_EQ(plusResult, otherPlusResult);
+    VELOX_CHECK_EQ(plusResult, dateAddResult);
+    return plusResult;
+  };
+
+  EXPECT_EQ(
+      "2024-10-04 01:50:00.000 America/Los_Angeles",
+      test("2024-10-03 01:50 America/Los_Angeles", 1 * kMillisInDay));
+  EXPECT_EQ(
+      "2024-10-03 02:50:00.000 America/Los_Angeles",
+      test("2024-10-03 01:50 America/Los_Angeles", 1 * kMillisInHour));
+  EXPECT_EQ(
+      "2024-10-03 01:51:00.000 America/Los_Angeles",
+      test("2024-10-03 01:50 America/Los_Angeles", 1 * kMillisInMinute));
+
+  // TODO Add tests for transitioning to/from DST once
+  // https://github.com/facebookincubator/velox/issues/10163 is fixed.
+}
+
 TEST_F(DateTimeFunctionsTest, minusTimestamp) {
   const auto minus = [&](std::optional<int64_t> t1, std::optional<int64_t> t2) {
     const auto timestamp1 = (t1.has_value()) ? Timestamp(t1.value(), 0)