From 096a41ff8084cdc561c80d428ae604e6717ad31a Mon Sep 17 00:00:00 2001 From: Aarav Date: Mon, 12 Oct 2026 20:51:43 +0530 Subject: [PATCH] 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