dmitry-chirkov-dremio commented on code in PR #53325:
URL: https://github.com/apache/arrow/pull/53325#discussion_r4369381749
##########
cpp/src/arrow/compute/kernels/scalar_arithmetic.cc:
##########
@@ -1126,6 +1135,335 @@ void
AddArithmeticFunctionDurationTime(std::shared_ptr<ScalarFunction> func) {
std::move(exec_4)));
}
+// ----------------------------------------------------------------------
+// Timestamp +/- interval arithmetic
+//
+// An interval is split into a calendar part (months and days) and a physical
+// part (the milliseconds of a day_time_interval or the nanoseconds of a
+// month_day_nano_interval). The calendar part is applied to the local
+// (wall clock) time of the timestamp, so that for example adding one day
+// keeps the same wall clock time across a DST transition. The physical part
+// is a fixed duration.
+//
+// `add` applies months, then days, then the physical part. `subtract`
+// applies the physical part first, so that subtracting e.g. "1 day 24 hours"
+// undoes adding it. Months are applied before days in both cases, as in
+// Pandas' DateOffset and PostgreSQL.
+//
+// Adding months clamps the day of month to the last day of the resulting
+// month (e.g. Jan 31 + 1 month = Feb 28). An error is returned if the
+// resulting local time is ambiguous or does not exist, as when adding a
+// DateOffset to a timezone-aware Pandas Series.
+
+struct CalendarInterval {
+ int32_t months;
+ int32_t days;
+ int64_t nanoseconds;
+};
+
+CalendarInterval ToCalendarInterval(int32_t months) { return {months, 0, 0}; }
+
+CalendarInterval ToCalendarInterval(DayTimeIntervalType::DayMilliseconds v) {
+ return {0, v.days, static_cast<int64_t>(v.milliseconds) * 1000000};
+}
+
+CalendarInterval ToCalendarInterval(MonthDayNanoIntervalType::MonthDayNanos v)
{
+ return {v.months, v.days, v.nanoseconds};
+}
+
+int64_t FloorDiv(int64_t x, int64_t y) {
+ const int64_t q = x / y;
+ return (x % y != 0 && ((x < 0) != (y < 0))) ? q - 1 : q;
+}
+
+// Proleptic Gregorian calendar conversions, using 64-bit years so that the
+// whole range of second-resolution timestamps is covered.
+// See https://howardhinnant.github.io/date_algorithms.html
+int64_t DaysFromCivil(int64_t y, int64_t m, int64_t d) {
+ y -= m <= 2;
+ const int64_t era = FloorDiv(y, 400);
+ const int64_t yoe = y - era * 400;
+ const int64_t doy = (153 * (m > 2 ? m - 3 : m + 9) + 2) / 5 + d - 1;
+ const int64_t doe = yoe * 365 + yoe / 4 - yoe / 100 + doy;
+ return era * 146097 + doe - 719468;
+}
+
+void CivilFromDays(int64_t z, int64_t* y, int64_t* m, int64_t* d) {
+ z += 719468;
+ const int64_t era = FloorDiv(z, 146097);
+ const int64_t doe = z - era * 146097;
+ const int64_t yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
+ const int64_t doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
+ const int64_t mp = (5 * doy + 2) / 153;
+ *d = doy - (153 * mp + 2) / 5 + 1;
+ *m = mp < 10 ? mp + 3 : mp - 9;
+ *y = yoe + era * 400 + (*m <= 2);
+}
+
+int64_t LastDayOfMonth(int64_t y, int64_t m) {
+ if (m == 2) {
+ const bool is_leap = y % 4 == 0 && (y % 100 != 0 || y % 400 == 0);
+ return is_leap ? 29 : 28;
+ }
+ return (m == 4 || m == 6 || m == 9 || m == 11) ? 30 : 31;
+}
+
+// Add a number of months to a day count, clamping the day of month.
+int64_t AddMonthsToDays(int64_t days, int64_t months) {
+ int64_t y, m, d;
+ CivilFromDays(days, &y, &m, &d);
+ const int64_t total_months = y * 12 + (m - 1) + months;
+ y = FloorDiv(total_months, 12);
+ m = total_months - y * 12 + 1;
+ return DaysFromCivil(y, m, std::min(d, LastDayOfMonth(y, m)));
+}
+
+// Converts between UTC and local time (in seconds) for a timestamp's timezone.
+class IntervalLocalizer {
+ public:
+ static Result<IntervalLocalizer> Make(const std::string& timezone) {
+ IntervalLocalizer localizer;
+ if (timezone.empty()) return localizer;
+ ARROW_ASSIGN_OR_RAISE(ArrowTimeZone tz, LocateZone(timezone));
+ if (auto zone = std::get_if<const time_zone*>(&tz)) {
+ localizer.zone_ = *zone;
+ } else {
+ const auto& offset_zone = std::get<OffsetZone>(tz);
+ localizer.fixed_offset_ =
+
offset_zone.get_info(sys_time<std::chrono::seconds>{}).offset.count();
+ }
+ return localizer;
+ }
+
+ // UTC offset in effect at the given UTC instant
+ int64_t OffsetAtSys(int64_t sys_seconds) const {
+ if (zone_ == nullptr) return fixed_offset_;
+ const auto st =
sys_time<std::chrono::seconds>{std::chrono::seconds{sys_seconds}};
+ return zone_->get_info(st).offset.count();
+ }
+
+ bool IsFixed() const { return zone_ == nullptr; }
+
+ // UTC offset to use for converting the given local time (in `Duration`
+ // units) to UTC. Like ZonedLocalizer::ConvertLocalToSys, this returns an
+ // error if the local time is ambiguous or does not exist.
+ template <typename Duration>
+ int64_t OffsetForLocal(int64_t local, Status* st) const {
+ if (zone_ == nullptr) return fixed_offset_;
+ const auto lt = local_time<Duration>{Duration{local}};
+ const auto info = zone_->get_info(floor<std::chrono::seconds>(lt));
+ switch (info.result) {
+ case chrono::local_info::unique:
+ break;
+ case chrono::local_info::nonexistent:
+ *st = Status::Invalid("Local time does not exist: ",
+ chrono::nonexistent_local_time(lt, info).what());
+ break;
+ case chrono::local_info::ambiguous:
+ *st = Status::Invalid("Local time is ambiguous: ",
+ chrono::ambiguous_local_time(lt, info).what());
+ break;
+ }
+ return info.first.offset.count();
+ }
+
+ private:
+ const time_zone* zone_ = nullptr;
+ int64_t fixed_offset_ = 0;
+};
+
+template <typename Duration, bool kChecked>
+struct TimestampIntervalArithmeticImpl {
+ static constexpr int64_t kUnitsPerSecond = Duration::period::den;
+ static constexpr int64_t kUnitsPerDay = 86400 * kUnitsPerSecond;
+ static constexpr int64_t kNanosPerUnit = 1000000000 / kUnitsPerSecond;
+
+ // Overflow-aware helpers: in unchecked mode, results wrap around.
+ static int64_t Add(int64_t a, int64_t b, Status* st) {
+ int64_t out;
+ if (ARROW_PREDICT_FALSE(AddWithOverflow(a, b, &out)) && kChecked) {
+ *st = Status::Invalid("overflow");
+ }
+ return out;
+ }
+
+ static int64_t Multiply(int64_t a, int64_t b, Status* st) {
+ int64_t out;
+ if (ARROW_PREDICT_FALSE(MultiplyWithOverflow(a, b, &out)) && kChecked) {
+ *st = Status::Invalid("overflow");
+ }
+ return out;
+ }
+
+ static int64_t Subtract(int64_t a, int64_t b, Status* st) {
+ int64_t out;
+ if (ARROW_PREDICT_FALSE(SubtractWithOverflow(a, b, &out)) && kChecked) {
+ *st = Status::Invalid("overflow");
+ }
+ return out;
+ }
+
+ // Convert the physical part of an interval to the timestamp unit
+ static int64_t PhysicalToUnit(int64_t nanoseconds, Status* st) {
+ if (ARROW_PREDICT_FALSE(nanoseconds % kNanosPerUnit != 0)) {
+ *st = Status::Invalid("Interval has a component of ", nanoseconds,
+ " nanoseconds which cannot be represented in a
timestamp ",
+ "with unit ", UnitName());
+ return 0;
+ }
+ return nanoseconds / kNanosPerUnit;
+ }
+
+ static constexpr const char* UnitName() {
+ switch (kUnitsPerSecond) {
+ case 1:
+ return "s";
+ case 1000:
+ return "ms";
+ case 1000000:
+ return "us";
+ default:
+ return "ns";
+ }
+ }
+
+ // Add months, then days, to the local time of `t`.
+ static int64_t AddCalendar(const IntervalLocalizer& localizer, int64_t t,
+ int64_t months, int64_t days, Status* st) {
+ if (months == 0 && days == 0) return t;
+ const int64_t offset = localizer.OffsetAtSys(FloorDiv(t, kUnitsPerSecond));
+ // Split before applying the offset. A local timestamp can lie just
outside
+ // int64's range even though the final UTC timestamp is representable.
+ const int64_t original_day = FloorDiv(t, kUnitsPerDay);
+ int64_t time_of_day = t % kUnitsPerDay;
+ if (time_of_day < 0) time_of_day += kUnitsPerDay;
+ time_of_day += offset * kUnitsPerSecond;
+ const int64_t carry = FloorDiv(time_of_day, kUnitsPerDay);
+ int64_t day = original_day + carry;
+ time_of_day -= carry * kUnitsPerDay;
+ const int64_t original_local_day = day;
+ if (months != 0) {
+ day = AddMonthsToDays(day, months);
+ }
+ day += days;
+ int64_t new_offset;
+ if (localizer.IsFixed()) {
+ new_offset = offset;
+ } else {
+ // Timezone transitions have second precision; the subsecond remainder
+ // does not affect which offset applies.
+ const int64_t local_seconds =
+ Add(Multiply(day, 86400, st), time_of_day / kUnitsPerSecond, st);
+ if (!st->ok()) return 0;
+ new_offset =
localizer.OffsetForLocal<std::chrono::seconds>(local_seconds, st);
+ }
+ // Apply only the calendar and offset deltas to the original timestamp.
+ // Reconstructing day * kUnitsPerDay first can overflow near INT64_MIN.
+ const int64_t day_delta = Subtract(day, original_local_day, st);
+ const int64_t offset_delta = Subtract(offset, new_offset, st);
+ const int64_t delta = Add(Multiply(day_delta, kUnitsPerDay, st),
Review Comment:
The intermediate overflow is still there. Adding 106752 days to
timestamp[[ns](app://-/INT64_MIN)](INT64_MIN) gives 763145224192, which fits,
but add_checked returns overflow because this multiplication happens first.
The new tests cover the reported examples. Please also cover a calendar
delta larger than int64 when the final timestamp still fits.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
