From 418484146c2912d6119b34c359615ad805d56e9e Mon Sep 17 00:00:00 2001 From: Pedro Eugenio Rocha Pedreira Date: Tue, 17 Sep 2024 08:42:00 -0700 Subject: [PATCH] Fix timestamp and interval arithmetic bug (#11017) Summary: Pull Request resolved: https://github.com/facebookincubator/velox/pull/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 | 51 +++++++++++++++++-- velox/functions/prestosql/DateTimeImpl.h | 18 +++++++ .../prestosql/tests/DateTimeFunctionsTest.cpp | 11 +++- 3 files changed, 76 insertions(+), 4 deletions(-) diff --git a/velox/functions/prestosql/DateTimeFunctions.h b/velox/functions/prestosql/DateTimeFunctions.h index bb09e7ebd3a4..4770cd58a229 100644 --- a/velox/functions/prestosql/DateTimeFunctions.h +++ b/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& /*inputTypes*/, + const core::QueryConfig& config, + const arg_type*, + const arg_type*) { + sessionTimeZone_ = getTimeZoneFromConfig(config); + } + FOLLY_ALWAYS_INLINE void call( out_type& result, const arg_type& timestamp, const arg_type& 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 @@ -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& /*inputTypes*/, + const core::QueryConfig& config, + const arg_type*, + const arg_type*) { + sessionTimeZone_ = getTimeZoneFromConfig(config); + } + FOLLY_ALWAYS_INLINE void call( out_type& result, const arg_type& interval, const arg_type& 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 @@ -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& /*inputTypes*/, + const core::QueryConfig& config, + const arg_type*, + const arg_type*) { + sessionTimeZone_ = getTimeZoneFromConfig(config); + } + FOLLY_ALWAYS_INLINE void call( out_type& result, const arg_type& timestamp, const arg_type& 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 diff --git a/velox/functions/prestosql/DateTimeImpl.h b/velox/functions/prestosql/DateTimeImpl.h index 4d3109a74d08..689294123d8f 100644 --- a/velox/functions/prestosql/DateTimeImpl.h +++ b/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, diff --git a/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp b/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp index 575206afcfd6..919dc1655bb4 100644 --- a/velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp +++ b/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) {