dmitry-chirkov-dremio commented on code in PR #53325:
URL: https://github.com/apache/arrow/pull/53325#discussion_r4334600989


##########
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:
   The intermediate calendar values can overflow even when the final timestamp 
fits.
   
   For timestamp[ns], adding one day to INT64_MIN causes signed overflow here. 
Subtracting one day from INT64_MIN + 86400000000000 should return INT64_MIN, 
but the multiplication below reports overflow. Nonzero timezone offsets cause 
similar failures near the bounds.
   
   Please avoid these intermediate overflows and add kernel regression tests 
for both bounds, including a nonzero timezone offset.



##########
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:
   Could you add a paired test showing the difference between subtracting a 
mixed interval and adding its negative?
   
   Starting at 2025-03-10 09:30 EDT:
   - Subtracting 1 day + 24 hours gives 2025-03-08 09:30 EST.
   - Adding -1 day -24 hours gives 2025-03-08 08:30 EST.
   
   Please also show that subtracting one calendar day and adding -1 calendar 
day give the same result. This would make the component order and its effect 
across DST clear.



-- 
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