Fix timestamp and interval arithmetic bug (facebookincubator#11017)

Summary:
Pull Request resolved: facebookincubator#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) {