DeviousCardi commented on code in PR #53325:
URL: https://github.com/apache/arrow/pull/53325#discussion_r4369731210
##########
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:
Fixed in 66a6b80c4. AddCalendar no longer adds day_delta * kUnitsPerDay to
the original timestamp. It now rebuilds the result from the shifted day and the
UTC time of day, and for negative days it multiplies (day + 1) and adds
(time_of_day - kUnitsPerDay), so day * kUnitsPerDay never has to fit on its
own. The tz lookup's local-seconds computation uses the same helper. All steps
still use the checked helpers, so a result that doesn't fit is still reported
as overflow. INT64_MIN ns + 106752 days now returns 763145224192 in both add
and add_checked.
New tests in TestTemporalTimestampAndIntervalErrors use the smallest day
count longer than INT64_MAX in the unit (106752 for ns, 106751992 for us), with
no tz, UTC, +05:30 and -05:30, as both day_time and month_day_nano, in checked
and unchecked kernels: INT64_MIN + n days, interval + INT64_MIN, the subtract
back to INT64_MIN, and the mirror cases from INT64_MAX. A month case also
covers 3509 months from INT64_MIN and back from INT64_MAX in ns. A delta of 2 *
106752 days still raises overflow in add_checked/subtract_checked, and the
unchecked kernels wrap. The new tests fail on 46a2ad2b5 and pass now.
arrow-compute-scalar-temporal-test and arrow-compute-scalar-math-test both pass.
--
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]
