From 096a41ff8084cdc561c80d428ae604e6717ad31a Mon Sep 17 00:00:00 2001 From: Aarav Date: Mon, 12 Oct 2026 20:51:43 +0530 Subject: [PATCH 1/3] GH-47400: [C++] Support adding/subtracting intervals to/from timestamps Add add, add_checked, subtract and subtract_checked kernels for timestamp +/- {month_interval, day_time_interval, month_day_nano_interval}, with and without a timezone (add also accepts the interval first). Months and days are calendar units applied to the local time of the timestamp, while the sub-day part of the interval is a fixed duration, following pandas.DateOffset. add applies months, days, then the sub-day part; subtract applies the sub-day part first. Adding months clamps the day of month. Ambiguous or nonexistent resulting local times return an error, and the checked variants detect overflow. Implicit temporal promotion is no longer applied in arithmetic dispatch when an argument is an interval, so that e.g. date32 + interval is not silently cast to a timestamp. --- .../compute/kernels/scalar_arithmetic.cc | 332 +++++++++++++++++- .../compute/kernels/scalar_temporal_test.cc | 310 ++++++++++++++++ docs/source/cpp/compute.rst | 15 + python/pyarrow/tests/test_compute.py | 63 ++++ 4 files changed, 718 insertions(+), 2 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc b/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc index 14dbd5b7e8bd..8219fc4b9dac 100644 --- a/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc +++ b/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc @@ -20,6 +20,7 @@ #include #include #include +#include #include #include "arrow/compare.h" @@ -28,6 +29,7 @@ #include "arrow/compute/kernels/base_arithmetic_internal.h" #include "arrow/compute/kernels/codegen_internal.h" #include "arrow/compute/kernels/common_internal.h" +#include "arrow/compute/kernels/temporal_internal.h" #include "arrow/compute/kernels/util_internal.h" #include "arrow/compute/registry_internal.h" #include "arrow/type.h" @@ -747,6 +749,11 @@ ArrayKernelExec GenerateArithmeticWithFixedIntOutType(detail::GetTypeId get_id) } } +bool HasInterval(const std::vector& types) { + return std::any_of(types.begin(), types.end(), + [](const TypeHolder& type) { return is_interval(type.id()); }); +} + struct ArithmeticFunction : ScalarFunction { using ScalarFunction::ScalarFunction; @@ -771,8 +778,10 @@ struct ArithmeticFunction : ScalarFunction { EnsureDictionaryDecoded(types); - // Only promote types for binary functions - if (types->size() == 2) { + // Only promote types for binary functions. Temporal arguments are not + // promoted when combined with an interval, as an interval does not + // have a unit (e.g. date32 + interval should not become a timestamp). + if (types->size() == 2 && !HasInterval(*types)) { ReplaceNullWithOtherType(types); TimeUnit::type finest_unit; if (CommonTemporalResolution(types->data(), types->size(), &finest_unit)) { @@ -1126,6 +1135,311 @@ void AddArithmeticFunctionDurationTime(std::shared_ptr 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(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 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(&tz)) { + localizer.zone_ = *zone; + } else { + const auto& offset_zone = std::get(tz); + localizer.fixed_offset_ = + offset_zone.get_info(sys_time{}).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{sys_seconds}}; + return zone_->get_info(st).offset.count(); + } + + // 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 + int64_t OffsetForLocal(int64_t local, Status* st) const { + if (zone_ == nullptr) return fixed_offset_; + const auto lt = local_time{Duration{local}}; + const auto info = zone_->get_info(floor(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 +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)); + const int64_t local = Add(t, offset * kUnitsPerSecond, st); + int64_t day = FloorDiv(local, kUnitsPerDay); + const int64_t time_of_day = local - day * kUnitsPerDay; + if (months != 0) { + day = AddMonthsToDays(day, months); + } + day += days; + const int64_t new_local = Add(Multiply(day, kUnitsPerDay, st), time_of_day, st); + const int64_t new_offset = localizer.OffsetForLocal(new_local, st); + return Subtract(new_local, new_offset * kUnitsPerSecond, st); + } + + static int64_t Call(const IntervalLocalizer& localizer, int64_t t, + const CalendarInterval& interval, bool subtract, Status* st) { + const int64_t physical = PhysicalToUnit(interval.nanoseconds, st); + if (!subtract) { + t = AddCalendar(localizer, t, interval.months, interval.days, st); + return Add(t, physical, st); + } + t = Subtract(t, physical, st); + return AddCalendar(localizer, t, -static_cast(interval.months), + -static_cast(interval.days), st); + } +}; + +template +struct TimestampIntervalOp { + IntervalLocalizer localizer; + + template + T Call(KernelContext*, Arg0 t, Arg1 interval, Status* st) const { + return TimestampIntervalArithmeticImpl::Call( + localizer, t, ToCalendarInterval(interval), kSubtract, st); + } +}; + +// Same as TimestampIntervalOp, but with the interval as first argument +template +struct IntervalTimestampOp { + IntervalLocalizer localizer; + + template + T Call(KernelContext*, Arg0 interval, Arg1 t, Status* st) const { + return TimestampIntervalArithmeticImpl::Call( + localizer, t, ToCalendarInterval(interval), /*subtract=*/false, st); + } +}; + +template +struct TimestampIntervalKernel { + static Status Exec(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) { + const auto& timezone = GetInputTimezone(*batch[kIntervalFirst ? 1 : 0].type()); + ARROW_ASSIGN_OR_RAISE(auto localizer, IntervalLocalizer::Make(timezone)); + if constexpr (kIntervalFirst) { + using Op = IntervalTimestampOp; + applicator::ScalarBinaryNotNullStateful + kernel{Op{localizer}}; + return kernel.Exec(ctx, batch, out); + } else { + using Op = TimestampIntervalOp; + applicator::ScalarBinaryNotNullStateful + kernel{Op{localizer}}; + return kernel.Exec(ctx, batch, out); + } + } +}; + +template +void AddTimestampIntervalKernel(ScalarFunction* func, TimeUnit::type unit, + const std::shared_ptr& interval_type) { + InputType ts_type(match::TimestampTypeUnit(unit)); + DCHECK_OK( + func->AddKernel({ts_type, interval_type}, OutputType(FirstType), + TimestampIntervalKernel::Exec)); + if constexpr (!kSubtract) { + DCHECK_OK(func->AddKernel( + {interval_type, ts_type}, OutputType(LastType), + TimestampIntervalKernel::Exec)); + } +} + +template +void AddTimestampIntervalKernels(ScalarFunction* func, TimeUnit::type unit) { + AddTimestampIntervalKernel( + func, unit, month_interval()); + AddTimestampIntervalKernel( + func, unit, day_time_interval()); + AddTimestampIntervalKernel( + func, unit, month_day_nano_interval()); +} + +// Add {add,subtract}[_checked](timestamp, interval) -> timestamp kernels, and +// add[_checked](interval, timestamp) -> timestamp kernels. +template +void AddTimestampIntervalKernels(ScalarFunction* func) { + AddTimestampIntervalKernels( + func, TimeUnit::SECOND); + AddTimestampIntervalKernels( + func, TimeUnit::MILLI); + AddTimestampIntervalKernels( + func, TimeUnit::MICRO); + AddTimestampIntervalKernels( + func, TimeUnit::NANO); +} + const FunctionDoc absolute_value_doc{ "Calculate the absolute value of the argument element-wise", ("Results will wrap around on integer overflow.\n" @@ -1550,6 +1864,9 @@ void RegisterScalarArithmetic(FunctionRegistry* registry) { AddArithmeticFunctionTimeDuration(add); AddArithmeticFunctionDurationTime(add); + // Add add(timestamp, interval) -> timestamp and add(interval, timestamp) + AddTimestampIntervalKernels(add.get()); + DCHECK_OK(registry->AddFunction(std::move(add))); // ---------------------------------------------------------------------- @@ -1578,6 +1895,10 @@ void RegisterScalarArithmetic(FunctionRegistry* registry) { AddArithmeticFunctionTimeDuration(add_checked); AddArithmeticFunctionDurationTime(add_checked); + // Add add_checked(timestamp, interval) -> timestamp and + // add_checked(interval, timestamp) -> timestamp + AddTimestampIntervalKernels(add_checked.get()); + DCHECK_OK(registry->AddFunction(std::move(add_checked))); // ---------------------------------------------------------------------- @@ -1636,6 +1957,9 @@ void RegisterScalarArithmetic(FunctionRegistry* registry) { AddArithmeticFunctionTimeDuration(subtract); + // Add subtract(timestamp, interval) -> timestamp + AddTimestampIntervalKernels(subtract.get()); + DCHECK_OK(registry->AddFunction(std::move(subtract))); // ---------------------------------------------------------------------- @@ -1702,6 +2026,10 @@ void RegisterScalarArithmetic(FunctionRegistry* registry) { AddArithmeticFunctionTimeDuration(subtract_checked); + // Add subtract_checked(timestamp, interval) -> timestamp + AddTimestampIntervalKernels( + subtract_checked.get()); + DCHECK_OK(registry->AddFunction(std::move(subtract_checked))); // ---------------------------------------------------------------------- diff --git a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc index 86a81ffdd384..f83982ade202 100644 --- a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +#include +#include #include #include @@ -1343,6 +1345,314 @@ TEST_F(ScalarTemporalTest, TestTemporalSubtractTimestampAndDuration) { } } +TEST_F(ScalarTemporalTest, TestTemporalAddTimestampAndInterval) { + // Without a timezone, calendar and physical units only differ in that months + // have a variable length. + const char* times = R"(["2024-01-31T09:00:00", "2024-02-29T09:00:00", + "2025-01-31T23:00:00", "1969-12-31T12:00:00", null])"; + const char* times_plus_month = R"(["2024-02-29T09:00:00", "2024-03-29T09:00:00", + "2025-02-28T23:00:00", "1970-01-31T12:00:00", + null])"; + const char* times_minus_year = R"(["2023-01-31T09:00:00", "2023-02-28T09:00:00", + "2024-01-31T23:00:00", "1968-12-31T12:00:00", + null])"; + const char* times_plus_day_2h = R"(["2024-02-01T11:00:00", "2024-03-01T11:00:00", + "2025-02-02T01:00:00", "1970-01-01T14:00:00", + null])"; + const char* times_plus_mdn = R"(["2024-03-01T09:00:01", "2024-03-30T09:00:01", + "2025-03-01T23:00:01", "1970-02-01T12:00:01", + null])"; + for (auto op : {"add", "add_checked"}) { + for (auto unit : TimeUnit::values()) { + auto ty = timestamp(unit); + auto arr = ArrayFromJSON(ty, times); + CheckScalarBinaryCommutative(op, arr, ScalarFromJSON(month_interval(), "1"), + ArrayFromJSON(ty, times_plus_month)); + CheckScalarBinaryCommutative(op, arr, ScalarFromJSON(month_interval(), "-12"), + ArrayFromJSON(ty, times_minus_year)); + CheckScalarBinaryCommutative(op, arr, + ScalarFromJSON(day_time_interval(), "[1, 7200000]"), + ArrayFromJSON(ty, times_plus_day_2h)); + CheckScalarBinaryCommutative( + op, arr, ScalarFromJSON(month_day_nano_interval(), "[1, 1, 1000000000]"), + ArrayFromJSON(ty, times_plus_mdn)); + CheckScalarBinaryCommutative( + op, arr, + ArrayFromJSON(month_day_nano_interval(), + "[[1, 0, 0], [0, 1, 0], [-1, -1, 0], null, [1, 1, 0]]"), + ArrayFromJSON(ty, R"(["2024-02-29T09:00:00", "2024-03-01T09:00:00", + "2024-12-30T23:00:00", null, null])")); + CheckScalarBinaryCommutative(op, arr, MakeNullScalar(month_day_nano_interval()), + ArrayFromJSON(ty, "[null, null, null, null, null]")); + } + } +} + +TEST_F(ScalarTemporalTest, TestTemporalSubtractTimestampAndInterval) { + for (auto op : {"subtract", "subtract_checked"}) { + for (auto unit : TimeUnit::values()) { + auto ty = timestamp(unit); + // Month end clamping, leap years, negative intervals and nulls + auto arr = ArrayFromJSON(ty, R"(["2025-03-31T09:00:00", "2024-03-31T09:00:00", + "2025-03-01T12:00:00", null])"); + CheckScalarBinary(op, arr, ScalarFromJSON(month_interval(), "1"), + ArrayFromJSON(ty, R"(["2025-02-28T09:00:00", + "2024-02-29T09:00:00", + "2025-02-01T12:00:00", null])")); + CheckScalarBinary(op, arr, ScalarFromJSON(month_interval(), "-11"), + ArrayFromJSON(ty, R"(["2026-02-28T09:00:00", + "2025-02-28T09:00:00", + "2026-02-01T12:00:00", null])")); + CheckScalarBinary(op, arr, ScalarFromJSON(day_time_interval(), "[1, 3600000]"), + ArrayFromJSON(ty, R"(["2025-03-30T08:00:00", + "2024-03-30T08:00:00", + "2025-02-28T11:00:00", null])")); + // Months are subtracted before days + CheckScalarBinary(op, arr, + ScalarFromJSON(month_day_nano_interval(), "[1, 1, -1000000000]"), + ArrayFromJSON(ty, R"(["2025-02-27T09:00:01", + "2024-02-28T09:00:01", + "2025-01-31T12:00:01", null])")); + CheckScalarBinary(op, arr, MakeNullScalar(month_interval()), + ArrayFromJSON(ty, "[null, null, null, null]")); + } + } +} + +TEST_F(ScalarTemporalTest, TestTemporalTimestampAndIntervalAcrossDST) { + // Expected values were checked against pandas.DateOffset. Days and months + // are calendar units that keep the local time of day, while the physical + // part of an interval is a fixed duration. + constexpr int64_t kHour = 3600LL * 1000000000LL; + struct Case { + const char* timezone; + const char* input; + int32_t months; + int32_t days; + int64_t nanoseconds; + const char* expected; + }; + const std::vector add_cases = { + // 2025-03-08 09:30 EST + 1 day = 2025-03-09 09:30 EDT + {"America/New_York", "2025-03-08T14:30:00", 0, 1, 0, "2025-03-09T13:30:00"}, + // 2025-03-08 09:30 EST + 24 hours = 2025-03-09 10:30 EDT + {"America/New_York", "2025-03-08T14:30:00", 0, 0, 24 * kHour, + "2025-03-09T14:30:00"}, + // Mixed: 1 day, then 24 hours + {"America/New_York", "2025-03-08T14:30:00", 0, 1, 24 * kHour, + "2025-03-10T13:30:00"}, + // 2025-03-08 03:30 EST + 1 day = 2025-03-09 03:30 EDT (just after the gap) + {"America/New_York", "2025-03-08T08:30:00", 0, 1, 0, "2025-03-09T07:30:00"}, + // 2025-11-01 02:30 EDT + 1 day = 2025-11-02 02:30 EST (just after the overlap) + {"America/New_York", "2025-11-01T06:30:00", 0, 1, 0, "2025-11-02T07:30:00"}, + // 2025-02-10 09:00 EST + 1 month = 2025-03-10 09:00 EDT + {"America/New_York", "2025-02-10T14:00:00", 1, 0, 0, "2025-03-10T13:00:00"}, + // 2025-01-31 23:00 EST + 1 month 1 day 2 hours = 2025-03-01 01:00 EST + {"America/New_York", "2025-02-01T04:00:00", 1, 1, 2 * kHour, "2025-03-02T06:00:00"}, + // 2025-03-09 09:30 EDT - 1 day = 2025-03-08 09:30 EST + {"America/New_York", "2025-03-09T13:30:00", 0, -1, 0, "2025-03-08T14:30:00"}, + // 2025-03-09 03:30 EDT - 1 hour = 2025-03-09 01:30 EST + {"America/New_York", "2025-03-09T07:30:00", 0, 0, -kHour, "2025-03-09T06:30:00"}, + // 2025-10-25 02:30 BST + 1 day = 2025-10-26 02:30 GMT + {"Europe/London", "2025-10-25T01:30:00", 0, 1, 0, "2025-10-26T02:30:00"}, + // 2025-03-29 12:00 GMT + 1 day = 2025-03-30 12:00 BST + {"Europe/London", "2025-03-29T12:00:00", 0, 1, 0, "2025-03-30T11:00:00"}, + {"Europe/London", "2025-03-29T12:00:00", 0, 0, 24 * kHour, "2025-03-30T12:00:00"}, + // 2025-02-28 01:30 GMT + 1 month 1 day = 2025-03-29 01:30 GMT + {"Europe/London", "2025-02-28T01:30:00", 1, 1, 0, "2025-03-29T01:30:00"}, + // Lord Howe Island has a 30 minute DST shift. + // 2025-04-05 02:15 +11 + 1 day = 2025-04-06 02:15 +1030 + {"Australia/Lord_Howe", "2025-04-04T15:15:00", 0, 1, 0, "2025-04-05T15:45:00"}, + // 2025-10-04 02:45 +1030 + 1 day = 2025-10-05 02:45 +11 + {"Australia/Lord_Howe", "2025-10-03T16:15:00", 0, 1, 0, "2025-10-04T15:45:00"}, + // 2025-10-04 12:00 +1030 + 1 day = 2025-10-05 12:00 +11 + {"Australia/Lord_Howe", "2025-10-04T01:30:00", 0, 1, 0, "2025-10-05T01:00:00"}, + {"Australia/Lord_Howe", "2025-10-04T01:30:00", 0, 0, 24 * kHour, + "2025-10-05T01:30:00"}, + // Fixed offsets have no DST + {"+05:30", "2025-03-08T14:30:00", 1, 1, kHour, "2025-04-09T15:30:00"}, + }; + const std::vector subtract_cases = { + // 2025-11-02 01:30 EST - 1 day = 2025-11-01 01:30 EDT + {"America/New_York", "2025-11-02T06:30:00", 0, 1, 0, "2025-11-01T05:30:00"}, + // Subtraction undoes the mixed addition above: 24 hours, then 1 day + {"America/New_York", "2025-03-10T13:30:00", 0, 1, 24 * kHour, + "2025-03-08T14:30:00"}, + // 2025-03-31 09:00 EDT - 1 month = 2025-02-28 09:00 EST + {"America/New_York", "2025-03-31T13:00:00", 1, 0, 0, "2025-02-28T14:00:00"}, + // 2025-03-01 12:00 EST - 1 month 1 day = 2025-01-31 12:00 EST + {"America/New_York", "2025-03-01T17:00:00", 1, 1, 0, "2025-01-31T17:00:00"}, + {"America/New_York", "2025-03-09T13:30:00", 0, -1, 0, "2025-03-10T13:30:00"}, + {"Europe/London", "2025-03-30T11:00:00", 0, 1, 0, "2025-03-29T12:00:00"}, + {"Australia/Lord_Howe", "2025-10-05T01:00:00", 0, 1, 0, "2025-10-04T01:30:00"}, + }; + + auto check = [&](const std::vector& cases, + std::initializer_list functions, bool commutative) { + for (const auto& c : cases) { + ARROW_SCOPED_TRACE(c.timezone, " ", c.input, " [", c.months, ", ", c.days, ", ", + c.nanoseconds, "]"); + const auto interval = std::make_shared( + MonthDayNanoIntervalType::MonthDayNanos{c.months, c.days, c.nanoseconds}); + for (auto unit : TimeUnit::values()) { + auto ty = timestamp(unit, c.timezone); + auto input = ArrayFromJSON(ty, std::string("[\"") + c.input + "\", null]"); + auto expected = ArrayFromJSON(ty, std::string("[\"") + c.expected + "\", null]"); + for (auto func : functions) { + if (commutative) { + CheckScalarBinaryCommutative(func, input, interval, expected); + } else { + CheckScalarBinary(func, input, interval, expected); + } + } + } + } + }; + check(add_cases, {"add", "add_checked"}, /*commutative=*/true); + check(subtract_cases, {"subtract", "subtract_checked"}, /*commutative=*/false); + + // day_time_interval and month_interval follow the same rules + for (auto unit : TimeUnit::values()) { + auto ty = timestamp(unit, "America/New_York"); + auto input = ArrayFromJSON(ty, R"(["2025-03-08T14:30:00", "2025-02-10T14:00:00"])"); + CheckScalarBinaryCommutative( + "add", input, ScalarFromJSON(day_time_interval(), "[1, 0]"), + ArrayFromJSON(ty, R"(["2025-03-09T13:30:00", "2025-02-11T14:00:00"])")); + CheckScalarBinaryCommutative( + "add", input, ScalarFromJSON(day_time_interval(), "[0, 86400000]"), + ArrayFromJSON(ty, R"(["2025-03-09T14:30:00", "2025-02-11T14:00:00"])")); + CheckScalarBinaryCommutative( + "add_checked", input, ScalarFromJSON(month_interval(), "1"), + ArrayFromJSON(ty, R"(["2025-04-08T13:30:00", "2025-03-10T13:00:00"])")); + CheckScalarBinary( + "subtract_checked", input, ScalarFromJSON(month_interval(), "-1"), + ArrayFromJSON(ty, R"(["2025-04-08T13:30:00", "2025-03-10T13:00:00"])")); + } + + // Local times that fall in a DST gap or overlap are an error, as in Pandas + struct ErrorCase { + const char* timezone; + const char* input; + int32_t months; + int32_t days; + const char* function; + const char* message; + }; + const std::vector error_cases = { + // 2025-03-08 02:30 EST + 1 day = 2025-03-09 02:30 (nonexistent) + {"America/New_York", "2025-03-08T07:30:00", 0, 1, "add", "does not exist"}, + // 2025-11-01 01:30 EDT + 1 day = 2025-11-02 01:30 (ambiguous) + {"America/New_York", "2025-11-01T05:30:00", 0, 1, "add_checked", "is ambiguous"}, + // 2025-03-10 02:30 EDT - 1 day = 2025-03-09 02:30 (nonexistent) + {"America/New_York", "2025-03-10T06:30:00", 0, 1, "subtract", "does not exist"}, + // 2025-03-29 01:30 GMT + 1 day = 2025-03-30 01:30 (nonexistent) + {"Europe/London", "2025-03-29T01:30:00", 0, 1, "add", "does not exist"}, + // 2025-09-26 01:30 BST + 1 month = 2025-10-26 01:30 (ambiguous) + {"Europe/London", "2025-09-26T00:30:00", 1, 0, "add", "is ambiguous"}, + // 2025-04-05 01:45 +11 + 1 day = 2025-04-06 01:45 (ambiguous) + {"Australia/Lord_Howe", "2025-04-04T14:45:00", 0, 1, "add", "is ambiguous"}, + // 2025-11-05 02:15 +11 - 1 month = 2025-10-05 02:15 (nonexistent) + {"Australia/Lord_Howe", "2025-11-04T15:15:00", 1, 0, "subtract_checked", + "does not exist"}, + }; + for (const auto& c : error_cases) { + ARROW_SCOPED_TRACE(c.timezone, " ", c.input, " ", c.function); + for (auto unit : TimeUnit::values()) { + auto input = ArrayFromJSON(timestamp(unit, c.timezone), + std::string("[\"") + c.input + "\", null]"); + auto interval = std::make_shared( + MonthDayNanoIntervalType::MonthDayNanos{c.months, c.days, 0}); + EXPECT_RAISES_WITH_MESSAGE_THAT(Invalid, ::testing::HasSubstr(c.message), + CallFunction(c.function, {input, interval})); + } + } + + // Unknown timezone + auto unknown_tz = ArrayFromJSON(timestamp(TimeUnit::SECOND, "Mars/Mariner"), "[0]"); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("Cannot locate or parse timezone 'Mars/Mariner'"), + CallFunction("add", {unknown_tz, ScalarFromJSON(month_interval(), "1")})); +} + +TEST_F(ScalarTemporalTest, TestTemporalTimestampAndIntervalErrors) { + auto mdn = month_day_nano_interval(); + + // The physical part of the interval must be representable in the timestamp unit + for (auto func : {"add", "add_checked", "subtract", "subtract_checked"}) { + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, + ::testing::HasSubstr("Interval has a component of 1 nanoseconds which cannot " + "be represented in a timestamp with unit s"), + CallFunction(func, {ArrayFromJSON(timestamp(TimeUnit::SECOND), "[0]"), + ScalarFromJSON(mdn, "[0, 0, 1]")})); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("timestamp with unit us"), + CallFunction(func, {ArrayFromJSON(timestamp(TimeUnit::MICRO, "UTC"), "[0]"), + ScalarFromJSON(mdn, "[1, 0, 1]")})); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("timestamp with unit s"), + CallFunction(func, {ArrayFromJSON(timestamp(TimeUnit::SECOND), "[0]"), + ScalarFromJSON(day_time_interval(), "[0, 1]")})); + } + CheckScalarBinary("add", ArrayFromJSON(timestamp(TimeUnit::MILLI), "[0]"), + ScalarFromJSON(day_time_interval(), "[0, 1]"), + ArrayFromJSON(timestamp(TimeUnit::MILLI), "[1]")); + + // Overflow + constexpr int64_t kMax = std::numeric_limits::max(); + constexpr int64_t kMin = std::numeric_limits::min(); + for (auto tz : {"", "UTC"}) { + auto ty = timestamp(TimeUnit::NANO, tz); + auto near_max = ArrayFromJSON(ty, R"(["2262-04-11T00:00:00"])"); + auto near_min = ArrayFromJSON(ty, R"(["1677-09-22T00:00:00"])"); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("overflow"), + CallFunction("add_checked", {near_max, ScalarFromJSON(month_interval(), "1")})); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("overflow"), + CallFunction("add_checked", + {ScalarFromJSON(day_time_interval(), "[1, 0]"), near_max})); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("overflow"), + CallFunction("subtract_checked", + {near_min, ScalarFromJSON(month_interval(), "1")})); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("overflow"), + CallFunction("subtract_checked", {near_min, ScalarFromJSON(mdn, "[0, 1, 0]")})); + + // Physical part + auto max = ArrayFromJSON(ty, "[" + std::to_string(kMax) + "]"); + auto min = ArrayFromJSON(ty, "[" + std::to_string(kMin) + "]"); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("overflow"), + CallFunction("add_checked", {max, ScalarFromJSON(mdn, "[0, 0, 1]")})); + EXPECT_RAISES_WITH_MESSAGE_THAT( + Invalid, ::testing::HasSubstr("overflow"), + CallFunction("subtract_checked", {min, ScalarFromJSON(mdn, "[0, 0, 1]")})); + // The unchecked variants wrap around + CheckScalarBinary("add", max, ScalarFromJSON(mdn, "[0, 0, 1]"), min); + CheckScalarBinary("subtract", min, ScalarFromJSON(mdn, "[0, 0, 1]"), max); + ASSERT_OK(CallFunction("add", {near_max, ScalarFromJSON(month_interval(), "1")})); + } + + // Extreme intervals don't overflow second-resolution timestamps + auto ty = timestamp(TimeUnit::SECOND); + CheckScalarBinary("add_checked", ArrayFromJSON(ty, "[0]"), + ScalarFromJSON(mdn, "[2147483647, 2147483647, 0]"), + ArrayFromJSON(ty, "[5832879117840000]")); + CheckScalarBinary("subtract_checked", ArrayFromJSON(ty, "[0]"), + ScalarFromJSON(mdn, "[2147483647, 2147483647, 0]"), + ArrayFromJSON(ty, "[-5832879117926400]")); + + // Dates are not supported yet, and are not implicitly cast to timestamps + for (auto date_ty : {date32(), date64()}) { + EXPECT_RAISES_WITH_MESSAGE_THAT( + NotImplemented, ::testing::HasSubstr("Function 'add' has no kernel matching"), + CallFunction("add", {ArrayFromJSON(date_ty, "[0]"), + ScalarFromJSON(month_interval(), "1")})); + } +} + TEST_F(ScalarTemporalTest, TestTemporalSubtractDate) { for (auto op : {"subtract", "subtract_checked"}) { auto arr_date32s = ArrayFromJSON(date32(), date32s); diff --git a/docs/source/cpp/compute.rst b/docs/source/cpp/compute.rst index 15d426dca6c6..5bc35ac780d0 100644 --- a/docs/source/cpp/compute.rst +++ b/docs/source/cpp/compute.rst @@ -494,6 +494,21 @@ arguments to floating-point, while mixed decimal and integer arguments will cast all arguments to decimals. Mixed time resolution temporal inputs will be cast to finest input resolution. +``add``, ``subtract`` and their checked variants also accept a timestamp and an +interval (``month_interval``, ``day_time_interval`` or +``month_day_nano_interval``), returning a timestamp of the same type +(``add`` also accepts the interval as first argument). Months and days are +calendar units applied to the local time of the timestamp, so adding one day +keeps the same wall clock time even across a daylight saving time transition. +The sub-day part of the interval is a fixed duration, and returns an +``Invalid`` :class:`Status` if it is not a multiple of the timestamp unit. +``add`` applies months, then days, then the sub-day part; ``subtract`` +applies the sub-day part first, then months, then days. When adding months, +the day of month is clamped to the last day of the resulting month (e.g. +January 31st plus one month is February 28th or 29th). An ``Invalid`` +:class:`Status` is returned if the resulting local time is ambiguous or does +not exist. + +-------------------+--------+-------------------------+-------------------------------+-------+ | Function name | Arity | Input types | Output type | Notes | +===================+========+=========================+===============================+=======+ diff --git a/python/pyarrow/tests/test_compute.py b/python/pyarrow/tests/test_compute.py index d4d44e9347e0..2da1a7ee74fb 100644 --- a/python/pyarrow/tests/test_compute.py +++ b/python/pyarrow/tests/test_compute.py @@ -3034,6 +3034,69 @@ def test_round_temporal_duration(unit, base_frequency, round_frequency): np.testing.assert_array_equal(result, expected) +@pytest.mark.pandas +@pytest.mark.timezone_data +@pytest.mark.parametrize("unit", ("s", "ms", "us", "ns")) +@pytest.mark.parametrize("tz", ("America/New_York", "Europe/London", + "Australia/Lord_Howe", "UTC", None)) +def test_add_subtract_timestamp_interval(unit, tz): + # Days and months are added to the local time (like pandas.DateOffset), + # while the sub-day part of an interval is a fixed duration. + times = pd.to_datetime([ + "2025-03-08 09:30", "2025-03-29 12:00", "2025-10-04 12:00", + "2025-01-31 23:00", "2024-02-29 12:00", "2025-11-01 18:00", + None]).as_unit(unit) + if tz is not None: + times = times.tz_localize(tz) + arr = pa.array(times, type=pa.timestamp(unit, tz)) + hour = 3600 * 10**9 + for months, days, nanos in [(0, 1, 0), (0, -1, 0), (1, 0, 0), (-12, 0, 0), + (1, 1, 2 * hour), (0, 1, 24 * hour), + (0, 0, 24 * hour), (0, 2, -hour)]: + interval = pa.scalar(pa.MonthDayNano([months, days, nanos])) + offset = pd.DateOffset(months=months, days=days) + delta = pd.Timedelta(nanos, "ns") + + expected = pa.array((times + offset) + delta, + type=pa.timestamp(unit, tz)) + for func in (pc.add, pc.add_checked): + assert func(arr, interval).equals(expected) + assert func(interval, arr).equals(expected) + + expected = pa.array((times - delta) - offset, + type=pa.timestamp(unit, tz)) + for func in (pc.subtract, pc.subtract_checked): + assert func(arr, interval).equals(expected) + + +@pytest.mark.pandas +@pytest.mark.timezone_data +@pytest.mark.parametrize("tz,local,message", [ + ("America/New_York", "2025-03-08 02:30", "does not exist"), + ("America/New_York", "2025-11-01 01:30", "is ambiguous"), + ("Europe/London", "2025-03-29 01:30", "does not exist"), + ("Australia/Lord_Howe", "2025-04-05 01:45", "is ambiguous"), +]) +def test_add_timestamp_interval_dst_errors(tz, local, message): + # Adding a day lands in a DST gap or overlap: like pandas, raise. + times = pd.DatetimeIndex([local]).tz_localize(tz) + with pytest.raises(ValueError): + times + pd.DateOffset(days=1) + arr = pa.array(times) + with pytest.raises(pa.ArrowInvalid, match=message): + pc.add(arr, pa.scalar(pa.MonthDayNano([0, 1, 0]))) + + +def test_add_timestamp_interval_errors(): + ts = pa.array([0], type=pa.timestamp("s")) + with pytest.raises(pa.ArrowInvalid, match="cannot be represented"): + pc.add(ts, pa.scalar(pa.MonthDayNano([0, 0, 1]))) + + ts = pa.array([datetime.datetime(2262, 4, 11)], type=pa.timestamp("ns")) + with pytest.raises(pa.ArrowInvalid, match="overflow"): + pc.add_checked(ts, pa.scalar(pa.MonthDayNano([1, 0, 0]))) + + def test_count(): arr = pa.array([1, 2, 3, None, None]) assert pc.count(arr).as_py() == 3 From 46a2ad2b5950211a93aa916eb62cf2ce3638db3c Mon Sep 17 00:00:00 2001 From: Aarav Date: Sat, 24 Oct 2026 04:21:35 +0530 Subject: [PATCH 2/3] Avoid intermediate overflow in timestamp interval arithmetic --- .../compute/kernels/scalar_arithmetic.cc | 36 +++++++++++++--- .../compute/kernels/scalar_temporal_test.cc | 41 +++++++++++++++++++ 2 files changed, 71 insertions(+), 6 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc b/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc index 8219fc4b9dac..7405a77b6c11 100644 --- a/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc +++ b/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc @@ -1242,6 +1242,8 @@ class IntervalLocalizer { 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. @@ -1330,16 +1332,38 @@ struct TimestampIntervalArithmeticImpl { int64_t months, int64_t days, Status* st) { if (months == 0 && days == 0) return t; const int64_t offset = localizer.OffsetAtSys(FloorDiv(t, kUnitsPerSecond)); - const int64_t local = Add(t, offset * kUnitsPerSecond, st); - int64_t day = FloorDiv(local, kUnitsPerDay); - const int64_t time_of_day = local - day * kUnitsPerDay; + // 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; - const int64_t new_local = Add(Multiply(day, kUnitsPerDay, st), time_of_day, st); - const int64_t new_offset = localizer.OffsetForLocal(new_local, st); - return Subtract(new_local, new_offset * kUnitsPerSecond, st); + 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(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), + Multiply(offset_delta, kUnitsPerSecond, st), st); + return Add(t, delta, st); } static int64_t Call(const IntervalLocalizer& localizer, int64_t t, diff --git a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc index f83982ade202..ed236d67ee0e 100644 --- a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc @@ -1511,6 +1511,29 @@ TEST_F(ScalarTemporalTest, TestTemporalTimestampAndIntervalAcrossDST) { check(add_cases, {"add", "add_checked"}, /*commutative=*/true); check(subtract_cases, {"subtract", "subtract_checked"}, /*commutative=*/false); + // Subtraction applies the fixed duration before the calendar day, while + // adding the negative interval applies the calendar day first. The two + // differ across the spring DST transition; calendar-only operations agree. + for (auto unit : TimeUnit::values()) { + auto ty = timestamp(unit, "America/New_York"); + auto start = ArrayFromJSON(ty, R"(["2025-03-10T13:30:00"])"); + auto mixed = ScalarFromJSON(month_day_nano_interval(), "[0, 1, 86400000000000]"); + auto negative_mixed = + ScalarFromJSON(month_day_nano_interval(), "[0, -1, -86400000000000]"); + auto day = ScalarFromJSON(day_time_interval(), "[1, 0]"); + auto negative_day = ScalarFromJSON(day_time_interval(), "[-1, 0]"); + for (auto suffix : {"", "_checked"}) { + CheckScalarBinary(std::string("subtract") + suffix, start, mixed, + ArrayFromJSON(ty, R"(["2025-03-08T14:30:00"])")); + CheckScalarBinary(std::string("add") + suffix, start, negative_mixed, + ArrayFromJSON(ty, R"(["2025-03-08T13:30:00"])")); + auto one_day_earlier = ArrayFromJSON(ty, R"(["2025-03-09T13:30:00"])"); + CheckScalarBinary(std::string("subtract") + suffix, start, day, one_day_earlier); + CheckScalarBinary(std::string("add") + suffix, start, negative_day, + one_day_earlier); + } + } + // day_time_interval and month_interval follow the same rules for (auto unit : TimeUnit::values()) { auto ty = timestamp(unit, "America/New_York"); @@ -1635,6 +1658,24 @@ TEST_F(ScalarTemporalTest, TestTemporalTimestampAndIntervalErrors) { ASSERT_OK(CallFunction("add", {near_max, ScalarFromJSON(month_interval(), "1")})); } + // Calendar arithmetic must not overflow an intermediate local timestamp or + // day-to-unit multiplication when the final UTC value fits. + constexpr int64_t kDay = 86400000000000LL; + for (auto tz : {"", "+05:30", "-05:30"}) { + auto ty = timestamp(TimeUnit::NANO, tz); + auto min = ArrayFromJSON(ty, "[" + std::to_string(kMin) + "]"); + auto min_plus_day = ArrayFromJSON(ty, "[" + std::to_string(kMin + kDay) + "]"); + auto max = ArrayFromJSON(ty, "[" + std::to_string(kMax) + "]"); + auto max_minus_day = ArrayFromJSON(ty, "[" + std::to_string(kMax - kDay) + "]"); + auto day = ScalarFromJSON(day_time_interval(), "[1, 0]"); + for (auto suffix : {"", "_checked"}) { + CheckScalarBinary(std::string("add") + suffix, min, day, min_plus_day); + CheckScalarBinary(std::string("subtract") + suffix, min_plus_day, day, min); + CheckScalarBinary(std::string("subtract") + suffix, max, day, max_minus_day); + CheckScalarBinary(std::string("add") + suffix, max_minus_day, day, max); + } + } + // Extreme intervals don't overflow second-resolution timestamps auto ty = timestamp(TimeUnit::SECOND); CheckScalarBinary("add_checked", ArrayFromJSON(ty, "[0]"), From 66a6b80c45c92ff3bfe42dbbffb0974924da5761 Mon Sep 17 00:00:00 2001 From: Aarav Date: Sun, 25 Oct 2026 10:38:12 +0530 Subject: [PATCH 3/3] Avoid overflow scaling large calendar deltas to the timestamp unit Rebuild the result from the shifted day and time of day instead of adding days * units_per_day to the original timestamp, so that a calendar delta longer than int64 in the timestamp unit does not overflow when the final timestamp fits. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../compute/kernels/scalar_arithmetic.cc | 31 +++++--- .../compute/kernels/scalar_temporal_test.cc | 73 +++++++++++++++++++ 2 files changed, 94 insertions(+), 10 deletions(-) diff --git a/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc b/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc index 7405a77b6c11..6ddaaf627729 100644 --- a/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc +++ b/cpp/src/arrow/compute/kernels/scalar_arithmetic.cc @@ -1303,6 +1303,18 @@ struct TimestampIntervalArithmeticImpl { return out; } + // Compute day * units_per_day + rem, where 0 <= rem < units_per_day, + // without overflowing when only day * units_per_day is out of range. + // For negative days, the product is formed with one day less in magnitude + // and the remainder is shifted by -units_per_day instead. + static int64_t ComposeDays(int64_t day, int64_t rem, int64_t units_per_day, + Status* st) { + if (day < 0) { + return Add(Multiply(day + 1, units_per_day, st), rem - units_per_day, st); + } + return Add(Multiply(day, units_per_day, st), rem, st); + } + // 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)) { @@ -1341,11 +1353,10 @@ struct TimestampIntervalArithmeticImpl { 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; + day = Add(day, days, st); int64_t new_offset; if (localizer.IsFixed()) { new_offset = offset; @@ -1353,17 +1364,17 @@ struct TimestampIntervalArithmeticImpl { // 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); + ComposeDays(day, time_of_day / kUnitsPerSecond, 86400, st); if (!st->ok()) return 0; new_offset = localizer.OffsetForLocal(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), - Multiply(offset_delta, kUnitsPerSecond, st), st); - return Add(t, delta, st); + // Convert back to UTC on the (day, time of day) pair, so that neither the + // local timestamp nor day * kUnitsPerDay has to be representable on its + // own: only the final timestamp must fit. + int64_t utc_time_of_day = time_of_day - new_offset * kUnitsPerSecond; + const int64_t utc_carry = FloorDiv(utc_time_of_day, kUnitsPerDay); + utc_time_of_day -= utc_carry * kUnitsPerDay; + return ComposeDays(Add(day, utc_carry, st), utc_time_of_day, kUnitsPerDay, st); } static int64_t Call(const IntervalLocalizer& localizer, int64_t t, diff --git a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc index ed236d67ee0e..b93009e18356 100644 --- a/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc +++ b/cpp/src/arrow/compute/kernels/scalar_temporal_test.cc @@ -1676,6 +1676,79 @@ TEST_F(ScalarTemporalTest, TestTemporalTimestampAndIntervalErrors) { } } + // A calendar delta whose length in the timestamp unit exceeds int64 must not + // overflow when the final timestamp fits. E.g. 106752 days is more than + // INT64_MAX nanoseconds, but INT64_MIN ns + 106752 days == 763145224192 ns. + ASSERT_EQ(kMin + 106751 * kDay + kDay, 763145224192LL); + for (auto unit : {TimeUnit::NANO, TimeUnit::MICRO}) { + const int64_t units_per_day = unit == TimeUnit::NANO ? kDay : kDay / 1000; + // Smallest day count whose length in `unit` does not fit in int64 + const int64_t n = kMax / units_per_day + 1; + // Step in two parts so that the expected values are computed without overflow + const int64_t from_min = kMin + (n - 1) * units_per_day + units_per_day; + const int64_t from_max = kMax - (n - 1) * units_per_day - units_per_day; + const auto n_str = std::to_string(n); + for (auto tz : {"", "UTC", "+05:30", "-05:30"}) { + auto ty = timestamp(unit, tz); + auto ts = [&](int64_t v) { + return ArrayFromJSON(ty, "[" + std::to_string(v) + "]"); + }; + for (auto days : {ScalarFromJSON(day_time_interval(), "[" + n_str + ", 0]"), + ScalarFromJSON(mdn, "[0, " + n_str + ", 0]")}) { + for (auto suffix : {"", "_checked"}) { + const auto add = std::string("add") + suffix; + const auto sub = std::string("subtract") + suffix; + CheckScalarBinary(add, ts(kMin), days, ts(from_min)); + CheckScalarBinary(add, days, ts(kMin), ts(from_min)); + CheckScalarBinary(sub, ts(from_min), days, ts(kMin)); + CheckScalarBinary(sub, ts(kMax), days, ts(from_max)); + CheckScalarBinary(add, ts(from_max), days, ts(kMax)); + } + } + } + } + // Same for a month delta: 3509 months is more than 106752 days. + for (auto tz : {"", "UTC", "+05:30", "-05:30"}) { + auto ty = timestamp(TimeUnit::NANO, tz); + // 1677-09-21T00:12:43.145224192 and 2262-04-11T23:47:16.854775807 + auto min = ArrayFromJSON(ty, "[" + std::to_string(kMin) + "]"); + auto max = ArrayFromJSON(ty, "[" + std::to_string(kMax) + "]"); + // 1970-02-21T00:12:43.145224192 and 1969-11-11T23:47:16.854775807 + auto min_plus = ArrayFromJSON(ty, "[4407163145224192]"); + auto max_minus = ArrayFromJSON(ty, "[-4320763145224193]"); + for (auto months : {ScalarFromJSON(month_interval(), "3509"), + ScalarFromJSON(mdn, "[3509, 0, 0]")}) { + for (auto suffix : {"", "_checked"}) { + CheckScalarBinary(std::string("add") + suffix, min, months, min_plus); + CheckScalarBinary(std::string("subtract") + suffix, max, months, max_minus); + } + } + } + // A calendar delta whose final timestamp does not fit is still an overflow + // in the checked variants, and wraps around in the unchecked ones. + for (auto tz : {"", "+05:30"}) { + auto ty = timestamp(TimeUnit::NANO, tz); + const int64_t n = 2 * (kMax / kDay + 1); + auto days = ScalarFromJSON(day_time_interval(), "[" + std::to_string(n) + ", 0]"); + auto min = ArrayFromJSON(ty, "[" + std::to_string(kMin) + "]"); + auto max = ArrayFromJSON(ty, "[" + std::to_string(kMax) + "]"); + EXPECT_RAISES_WITH_MESSAGE_THAT(Invalid, ::testing::HasSubstr("overflow"), + CallFunction("add_checked", {min, days})); + EXPECT_RAISES_WITH_MESSAGE_THAT(Invalid, ::testing::HasSubstr("overflow"), + CallFunction("subtract_checked", {max, days})); + const uint64_t span = static_cast(n) * static_cast(kDay); + CheckScalarBinary("add", min, days, + ArrayFromJSON(ty, "[" + + std::to_string(static_cast( + static_cast(kMin) + span)) + + "]")); + CheckScalarBinary("subtract", max, days, + ArrayFromJSON(ty, "[" + + std::to_string(static_cast( + static_cast(kMax) - span)) + + "]")); + } + // Extreme intervals don't overflow second-resolution timestamps auto ty = timestamp(TimeUnit::SECOND); CheckScalarBinary("add_checked", ArrayFromJSON(ty, "[0]"),