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