DeviousCardi commented on code in PR #53325:
URL: https://github.com/apache/arrow/pull/53325#discussion_r4364871384


##########
cpp/src/arrow/compute/kernels/scalar_arithmetic.cc:
##########
@@ -1126,6 +1135,311 @@ void 
AddArithmeticFunctionDurationTime(std::shared_ptr<ScalarFunction> func) {
                             std::move(exec_4)));
 }
 
+// ----------------------------------------------------------------------
+// Timestamp +/- interval arithmetic
+//
+// An interval is split into a calendar part (months and days) and a physical
+// part (the milliseconds of a day_time_interval or the nanoseconds of a
+// month_day_nano_interval).  The calendar part is applied to the local
+// (wall clock) time of the timestamp, so that for example adding one day
+// keeps the same wall clock time across a DST transition.  The physical part
+// is a fixed duration.
+//
+// `add` applies months, then days, then the physical part.  `subtract`
+// applies the physical part first, so that subtracting e.g. "1 day 24 hours"
+// undoes adding it.  Months are applied before days in both cases, as in
+// Pandas' DateOffset and PostgreSQL.
+//
+// Adding months clamps the day of month to the last day of the resulting
+// month (e.g. Jan 31 + 1 month = Feb 28).  An error is returned if the
+// resulting local time is ambiguous or does not exist, as when adding a
+// DateOffset to a timezone-aware Pandas Series.
+
+struct CalendarInterval {
+  int32_t months;
+  int32_t days;
+  int64_t nanoseconds;
+};
+
+CalendarInterval ToCalendarInterval(int32_t months) { return {months, 0, 0}; }
+
+CalendarInterval ToCalendarInterval(DayTimeIntervalType::DayMilliseconds v) {
+  return {0, v.days, static_cast<int64_t>(v.milliseconds) * 1000000};
+}
+
+CalendarInterval ToCalendarInterval(MonthDayNanoIntervalType::MonthDayNanos v) 
{
+  return {v.months, v.days, v.nanoseconds};
+}
+
+int64_t FloorDiv(int64_t x, int64_t y) {
+  const int64_t q = x / y;
+  return (x % y != 0 && ((x < 0) != (y < 0))) ? q - 1 : q;
+}
+
+// Proleptic Gregorian calendar conversions, using 64-bit years so that the
+// whole range of second-resolution timestamps is covered.
+// See https://howardhinnant.github.io/date_algorithms.html
+int64_t DaysFromCivil(int64_t y, int64_t m, int64_t d) {
+  y -= m <= 2;
+  const int64_t era = FloorDiv(y, 400);
+  const int64_t yoe = y - era * 400;
+  const int64_t doy = (153 * (m > 2 ? m - 3 : m + 9) + 2) / 5 + d - 1;
+  const int64_t doe = yoe * 365 + yoe / 4 - yoe / 100 + doy;
+  return era * 146097 + doe - 719468;
+}
+
+void CivilFromDays(int64_t z, int64_t* y, int64_t* m, int64_t* d) {
+  z += 719468;
+  const int64_t era = FloorDiv(z, 146097);
+  const int64_t doe = z - era * 146097;
+  const int64_t yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
+  const int64_t doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
+  const int64_t mp = (5 * doy + 2) / 153;
+  *d = doy - (153 * mp + 2) / 5 + 1;
+  *m = mp < 10 ? mp + 3 : mp - 9;
+  *y = yoe + era * 400 + (*m <= 2);
+}
+
+int64_t LastDayOfMonth(int64_t y, int64_t m) {
+  if (m == 2) {
+    const bool is_leap = y % 4 == 0 && (y % 100 != 0 || y % 400 == 0);
+    return is_leap ? 29 : 28;
+  }
+  return (m == 4 || m == 6 || m == 9 || m == 11) ? 30 : 31;
+}
+
+// Add a number of months to a day count, clamping the day of month.
+int64_t AddMonthsToDays(int64_t days, int64_t months) {
+  int64_t y, m, d;
+  CivilFromDays(days, &y, &m, &d);
+  const int64_t total_months = y * 12 + (m - 1) + months;
+  y = FloorDiv(total_months, 12);
+  m = total_months - y * 12 + 1;
+  return DaysFromCivil(y, m, std::min(d, LastDayOfMonth(y, m)));
+}
+
+// Converts between UTC and local time (in seconds) for a timestamp's timezone.
+class IntervalLocalizer {
+ public:
+  static Result<IntervalLocalizer> Make(const std::string& timezone) {
+    IntervalLocalizer localizer;
+    if (timezone.empty()) return localizer;
+    ARROW_ASSIGN_OR_RAISE(ArrowTimeZone tz, LocateZone(timezone));
+    if (auto zone = std::get_if<const time_zone*>(&tz)) {
+      localizer.zone_ = *zone;
+    } else {
+      const auto& offset_zone = std::get<OffsetZone>(tz);
+      localizer.fixed_offset_ =
+          
offset_zone.get_info(sys_time<std::chrono::seconds>{}).offset.count();
+    }
+    return localizer;
+  }
+
+  // UTC offset in effect at the given UTC instant
+  int64_t OffsetAtSys(int64_t sys_seconds) const {
+    if (zone_ == nullptr) return fixed_offset_;
+    const auto st = 
sys_time<std::chrono::seconds>{std::chrono::seconds{sys_seconds}};
+    return zone_->get_info(st).offset.count();
+  }
+
+  // UTC offset to use for converting the given local time (in `Duration`
+  // units) to UTC.  Like ZonedLocalizer::ConvertLocalToSys, this returns an
+  // error if the local time is ambiguous or does not exist.
+  template <typename Duration>
+  int64_t OffsetForLocal(int64_t local, Status* st) const {
+    if (zone_ == nullptr) return fixed_offset_;
+    const auto lt = local_time<Duration>{Duration{local}};
+    const auto info = zone_->get_info(floor<std::chrono::seconds>(lt));
+    switch (info.result) {
+      case chrono::local_info::unique:
+        break;
+      case chrono::local_info::nonexistent:
+        *st = Status::Invalid("Local time does not exist: ",
+                              chrono::nonexistent_local_time(lt, info).what());
+        break;
+      case chrono::local_info::ambiguous:
+        *st = Status::Invalid("Local time is ambiguous: ",
+                              chrono::ambiguous_local_time(lt, info).what());
+        break;
+    }
+    return info.first.offset.count();
+  }
+
+ private:
+  const time_zone* zone_ = nullptr;
+  int64_t fixed_offset_ = 0;
+};
+
+template <typename Duration, bool kChecked>
+struct TimestampIntervalArithmeticImpl {
+  static constexpr int64_t kUnitsPerSecond = Duration::period::den;
+  static constexpr int64_t kUnitsPerDay = 86400 * kUnitsPerSecond;
+  static constexpr int64_t kNanosPerUnit = 1000000000 / kUnitsPerSecond;
+
+  // Overflow-aware helpers: in unchecked mode, results wrap around.
+  static int64_t Add(int64_t a, int64_t b, Status* st) {
+    int64_t out;
+    if (ARROW_PREDICT_FALSE(AddWithOverflow(a, b, &out)) && kChecked) {
+      *st = Status::Invalid("overflow");
+    }
+    return out;
+  }
+
+  static int64_t Multiply(int64_t a, int64_t b, Status* st) {
+    int64_t out;
+    if (ARROW_PREDICT_FALSE(MultiplyWithOverflow(a, b, &out)) && kChecked) {
+      *st = Status::Invalid("overflow");
+    }
+    return out;
+  }
+
+  static int64_t Subtract(int64_t a, int64_t b, Status* st) {
+    int64_t out;
+    if (ARROW_PREDICT_FALSE(SubtractWithOverflow(a, b, &out)) && kChecked) {
+      *st = Status::Invalid("overflow");
+    }
+    return out;
+  }
+
+  // Convert the physical part of an interval to the timestamp unit
+  static int64_t PhysicalToUnit(int64_t nanoseconds, Status* st) {
+    if (ARROW_PREDICT_FALSE(nanoseconds % kNanosPerUnit != 0)) {
+      *st = Status::Invalid("Interval has a component of ", nanoseconds,
+                            " nanoseconds which cannot be represented in a 
timestamp ",
+                            "with unit ", UnitName());
+      return 0;
+    }
+    return nanoseconds / kNanosPerUnit;
+  }
+
+  static constexpr const char* UnitName() {
+    switch (kUnitsPerSecond) {
+      case 1:
+        return "s";
+      case 1000:
+        return "ms";
+      case 1000000:
+        return "us";
+      default:
+        return "ns";
+    }
+  }
+
+  // Add months, then days, to the local time of `t`.
+  static int64_t AddCalendar(const IntervalLocalizer& localizer, int64_t t,
+                             int64_t months, int64_t days, Status* st) {
+    if (months == 0 && days == 0) return t;
+    const int64_t offset = localizer.OffsetAtSys(FloorDiv(t, kUnitsPerSecond));
+    const int64_t local = Add(t, offset * kUnitsPerSecond, st);
+    int64_t day = FloorDiv(local, kUnitsPerDay);
+    const int64_t time_of_day = local - day * kUnitsPerDay;

Review Comment:
   Addressed in 46a2ad2b5. Calendar arithmetic now splits the timestamp into 
day and time before applying the timezone offset, then applies the calendar and 
offset deltas to the original timestamp. That avoids constructing an 
out-of-range local or day-to-unit intermediate when the final UTC value fits. 
Added checked and unchecked kernel tests at both INT64 bounds for UTC and fixed 
±05:30 offsets. The rebuilt temporal test target passes.



##########
cpp/src/arrow/compute/kernels/scalar_temporal_test.cc:
##########
@@ -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<Case> 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<Case> 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,

Review Comment:
   Added the paired DST test in 46a2ad2b5. Starting at 2025-03-10 09:30 EDT, 
subtracting 1 calendar day plus 24 hours returns 2025-03-08 09:30 EST, while 
adding the negative mixed interval returns 08:30 EST. The test also confirms 
that subtracting one calendar day equals adding negative one day, across all 
timestamp units and both checked and unchecked kernels. The rebuilt temporal 
test target passes.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to