Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix timestamp and interval arithmetic bug #11017

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 48 additions & 3 deletions velox/functions/prestosql/DateTimeFunctions.h
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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>
Expand All @@ -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(
Expand All @@ -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>
Expand All @@ -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(
Expand All @@ -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>
Expand Down
18 changes: 18 additions & 0 deletions velox/functions/prestosql/DateTimeImpl.h
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
11 changes: 10 additions & 1 deletion velox/functions/prestosql/tests/DateTimeFunctionsTest.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -919,7 +924,6 @@ TEST_F(DateTimeFunctionsTest, timestampPlusIntervalYearMonth) {

// They should be the same.
EXPECT_EQ(result1, result2);

return result1;
};

Expand All @@ -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) {
Expand Down
Loading