Fix timestamp and interval arithmetic bug (#11017)

Summary:
Pull Request resolved: #11017

For timestamps plus or minus intervals, the arithmetic needs to be
done at the local timezone (not UTC) to account for intervals than span a
daylight savings time boundary. This is already handled for timestamp with
timezone types, and not needed for IntervalDayTime.

Reviewed By: spershin

Differential Revision: D62811693

velox/functions/prestosql/DateTimeFunctions.h

@@ -533,11 +533,22 @@ struct TimestampPlusInterval {
     result = Timestamp::fromMillisNoError(a.toMillis() + b);
   }
 
+  // We only need to capture the time zone session config if we are operating on
+  // a timestamp and a IntervalYearMonth.
+  FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
+      const core::QueryConfig& config,
+      const arg_type<Timestamp>*,
+      const arg_type<IntervalYearMonth>*) {
+    sessionTimeZone_ = getTimeZoneFromConfig(config);
+  }
+
   FOLLY_ALWAYS_INLINE void call(
       out_type<Timestamp>& result,
       const arg_type<Timestamp>& timestamp,
       const arg_type<IntervalYearMonth>& interval) {
-    result = addToTimestamp(timestamp, DateTimeUnit::kMonth, interval);
+    result = addToTimestamp(
+        timestamp, DateTimeUnit::kMonth, interval, sessionTimeZone_);
   }
 
   FOLLY_ALWAYS_INLINE void call(
@@ -555,6 +566,10 @@ struct TimestampPlusInterval {
     result = addToTimestampWithTimezone(
         timestampWithTimezone, DateTimeUnit::kMonth, interval);
   }
+
+ private:
+  // Only set if the parameters are timestamp and IntervalYearMonth.
+  const tz::TimeZone* sessionTimeZone_ = nullptr;
 };
 
 template <typename T>
@@ -574,11 +589,22 @@ struct IntervalPlusTimestamp {
     result = Timestamp::fromMillisNoError(a + b.toMillis());
   }
 
+  // We only need to capture the time zone session config if we are operating on
+  // a timestamp and a IntervalYearMonth.
+  FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
+      const core::QueryConfig& config,
+      const arg_type<IntervalYearMonth>*,
+      const arg_type<Timestamp>*) {
+    sessionTimeZone_ = getTimeZoneFromConfig(config);
+  }
+
   FOLLY_ALWAYS_INLINE void call(
       out_type<Timestamp>& result,
       const arg_type<IntervalYearMonth>& interval,
       const arg_type<Timestamp>& timestamp) {
-    result = addToTimestamp(timestamp, DateTimeUnit::kMonth, interval);
+    result = addToTimestamp(
+        timestamp, DateTimeUnit::kMonth, interval, sessionTimeZone_);
   }
 
   FOLLY_ALWAYS_INLINE void call(
@@ -596,6 +622,10 @@ struct IntervalPlusTimestamp {
     result = addToTimestampWithTimezone(
         timestampWithTimezone, DateTimeUnit::kMonth, interval);
   }
+
+ private:
+  // Only set if the parameters are timestamp and IntervalYearMonth.
+  const tz::TimeZone* sessionTimeZone_ = nullptr;
 };
 
 template <typename T>
@@ -615,11 +645,22 @@ struct TimestampMinusInterval {
     result = Timestamp::fromMillisNoError(a.toMillis() - b);
   }
 
+  // We only need to capture the time zone session config if we are operating on
+  // a timestamp and a IntervalYearMonth.
+  FOLLY_ALWAYS_INLINE void initialize(
+      const std::vector<TypePtr>& /*inputTypes*/,
+      const core::QueryConfig& config,
+      const arg_type<Timestamp>*,
+      const arg_type<IntervalYearMonth>*) {
+    sessionTimeZone_ = getTimeZoneFromConfig(config);
+  }
+
   FOLLY_ALWAYS_INLINE void call(
       out_type<Timestamp>& result,
       const arg_type<Timestamp>& timestamp,
       const arg_type<IntervalYearMonth>& interval) {
-    result = addToTimestamp(timestamp, DateTimeUnit::kMonth, -interval);
+    result = addToTimestamp(
+        timestamp, DateTimeUnit::kMonth, -interval, sessionTimeZone_);
   }
 
   FOLLY_ALWAYS_INLINE void call(
@@ -637,6 +678,10 @@ struct TimestampMinusInterval {
     result = addToTimestampWithTimezone(
         timestampWithTimezone, DateTimeUnit::kMonth, -interval);
   }
+
+ private:
+  // Only set if the parameters are timestamp and IntervalYearMonth.
+  const tz::TimeZone* sessionTimeZone_ = nullptr;
 };
 
 template <typename T>

velox/functions/prestosql/DateTimeImpl.h

@@ -199,6 +199,24 @@ FOLLY_ALWAYS_INLINE Timestamp addToTimestamp(
           timestamp.getNanos() % kNanosecondsInMillisecond);
 }
 
+// If time zone is provided, use it for the arithmetic operation (convert to it,
+// apply operation, then convert back to UTC).
+FOLLY_ALWAYS_INLINE Timestamp addToTimestamp(
+    const Timestamp& timestamp,
+    const DateTimeUnit unit,
+    const int32_t value,
+    const tz::TimeZone* timeZone) {
+  if (timeZone == nullptr) {
+    return addToTimestamp(timestamp, unit, value);
+  } else {
+    Timestamp zonedTimestamp = timestamp;
+    zonedTimestamp.toTimezone(*timeZone);
+    auto resultTimestamp = addToTimestamp(zonedTimestamp, unit, value);
+    resultTimestamp.toGMT(*timeZone);
+    return resultTimestamp;
+  }
+}
+
 FOLLY_ALWAYS_INLINE int64_t addToTimestampWithTimezone(
     int64_t timestampWithTimezone,
     const DateTimeUnit unit,

velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp

@@ -896,6 +896,11 @@ TEST_F(DateTimeFunctionsTest, timestampMinusIntervalYearMonth) {
   EXPECT_EQ("2001-02-28 04:05:06", minus("2001-03-30 04:05:06", 1));
   EXPECT_EQ("2000-02-29 04:05:06", minus("2000-03-30 04:05:06", 1));
   EXPECT_EQ("2000-01-29 04:05:06", minus("2000-02-29 04:05:06", 1));
+
+  // Check if it does the right thing if we cross daylight saving boundaries.
+  setQueryTimeZone("America/Los_Angeles");
+  EXPECT_EQ("2024-01-01 00:00:00", minus("2024-07-01 00:00:00", 6));
+  EXPECT_EQ("2023-07-01 00:00:00", minus("2024-01-01 00:00:00", 6));
 }
 
 TEST_F(DateTimeFunctionsTest, timestampPlusIntervalYearMonth) {
@@ -919,7 +924,6 @@ TEST_F(DateTimeFunctionsTest, timestampPlusIntervalYearMonth) {
 
     // They should be the same.
     EXPECT_EQ(result1, result2);
-
     return result1;
   };
 
@@ -933,6 +937,11 @@ TEST_F(DateTimeFunctionsTest, timestampPlusIntervalYearMonth) {
   EXPECT_EQ("2001-02-28 04:05:06", plus("2001-01-31 04:05:06", 1));
   EXPECT_EQ("2000-02-29 04:05:06", plus("2000-01-31 04:05:06", 1));
   EXPECT_EQ("2000-02-29 04:05:06", plus("2000-01-29 04:05:06", 1));
+
+  // Check if it does the right thing if we cross daylight saving boundaries.
+  setQueryTimeZone("America/Los_Angeles");
+  EXPECT_EQ("2025-01-01 00:00:00", plus("2024-07-01 00:00:00", 6));
+  EXPECT_EQ("2024-07-01 00:00:00", plus("2024-01-01 00:00:00", 6));
 }
 
 TEST_F(DateTimeFunctionsTest, plusMinusTimestampIntervalDayTime) {