fix: keep the offset's resolution with an array n and in time_ranges

tests/test_processing.py

@@ -324,6 +324,16 @@ def test_offset_times_with_array_n_pyarrow(freq):
     assert res.dtype == times.dtype
 
 
+def test_offset_times_with_array_n_pyarrow_date():
+    freq = pd.offsets.MonthBegin()
+    n = np.array([1, 0, 2, 1, 0])
+    times = pd.Series(_ARRAY_N_TIMES).astype("date32[pyarrow]")
+    res = offset_times(times, freq, n)
+    expected = _offset_times_elementwise(pd.Series(_ARRAY_N_TIMES), freq, n)
+    assert list(res) == [t.date() for t in expected]
+    assert res.dtype == times.dtype
+
+
 @pytest.mark.parametrize(
     "n", [np.zeros(5, dtype=np.int64), np.full(5, 3), np.array([2])]
 )
@@ -371,6 +381,78 @@ def test_offset_times_with_float_array_n(freq):
     )
 
 
+@pytest.mark.parametrize(
+    "freq",
+    [
+        pd.DateOffset(milliseconds=500),
+        pd.DateOffset(microseconds=1),
+        pd.offsets.Milli(1000),
+        pd.offsets.Micro(10**6),
+        pd.offsets.Nano(10**9),
+    ],
+)
+@pytest.mark.parametrize("tz", [None, "Europe/Madrid"])
+def test_offset_times_with_array_n_finer_than_times(freq, tz):
+    n = np.array([1, 0, -2, 3, -1])
+    times = pd.Series(_ARRAY_N_TIMES.astype("datetime64[s]").tz_localize(tz))
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        res = offset_times(times, freq, n)
+    expected = pd.concat(
+        [offset_times(times.iloc[[i]], freq, int(k)) for i, k in enumerate(n)]
+    )
+    pd.testing.assert_series_equal(res, expected)
+
+
+@pytest.mark.parametrize("freq", [pd.offsets.Second(), pd.offsets.Second(-1)])
+def test_offset_times_with_array_n_overflow(freq):
+    times = pd.Series(_ARRAY_N_TIMES)
+    n = np.full(times.size, 2**62 // 10**9 * 4)
+    with pytest.raises(pd.errors.OutOfBoundsTimedelta):
+        offset_times(times, freq, int(n[0]))
+    with pytest.raises(pd.errors.OutOfBoundsTimedelta):
+        offset_times(times, freq, n)
+
+
+@pytest.mark.parametrize("freq", [pd.offsets.Second(0), pd.offsets.Second(-1)])
+def test_offset_times_with_array_n_non_positive_tick(freq):
+    times = pd.Series(_ARRAY_N_TIMES)
+    res = offset_times(times, freq, np.full(times.size, 2))
+    pd.testing.assert_series_equal(res, offset_times(times, freq, 2))
+
+
+@pytest.mark.parametrize("freq", [pd.offsets.Second(), pd.offsets.MonthBegin()])
+def test_offset_times_with_uint64_array_n_overflow(freq):
+    times = pd.Series(_ARRAY_N_TIMES)
+    n = np.zeros(times.size, dtype=np.uint64)
+    n[0] = 2**63 + 5
+    with pytest.raises(OverflowError):
+        offset_times(times, freq, int(n[0]))
+    with pytest.raises(OverflowError):
+        offset_times(times, freq, n)
+
+
+@pytest.mark.parametrize(
+    "freq",
+    [
+        pd.offsets.Day(),
+        pd.offsets.MonthBegin(),
+        pd.offsets.BusinessHour(),
+        pd.DateOffset(milliseconds=500),
+    ],
+)
+@pytest.mark.parametrize("times", [_ARRAY_N_TIMES[:0], pd.Series(_ARRAY_N_TIMES[:0])])
+def test_offset_times_with_empty_array_n(freq, times):
+    times = times.astype("datetime64[s]")
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        res = offset_times(times, freq, np.array([], dtype=np.int64))
+    assert type(res) is type(times)
+    assert res is not times
+    assert len(res) == 0
+    assert res.dtype == times.dtype
+
+
 def test_offset_times_pl():
     pl.testing.assert_series_equal(
         offset_times(
@@ -486,6 +568,31 @@ def test_time_ranges_with_year_end_offset_pd():
     )
 
 
+@pytest.mark.parametrize(
+    "freq", ["D", "h", "MS", "W-MON", "500ms", pd.DateOffset(milliseconds=500)]
+)
+@pytest.mark.parametrize("tz", [None, "America/New_York"])
+def test_time_ranges_matches_offset_times_pd(freq, tz):
+    starts = pd.Series(
+        _ARRAY_N_TIMES.astype("datetime64[s]").tz_localize(tz), name="ds"
+    )
+    periods = 4
+    res = time_ranges(starts, freq, periods)
+    per_period = [offset_times(starts, freq, i) for i in range(periods)]
+    expected = pd.Series(
+        [per_period[i][j] for j in range(len(starts)) for i in range(periods)],
+        dtype=per_period[0].dtype,
+    )
+    pd.testing.assert_series_equal(res, expected)
+
+
+@pytest.mark.parametrize("tz", [None, "America/New_York"])
+def test_time_ranges_zero_periods_pd(tz):
+    starts = _ARRAY_N_TIMES.astype("datetime64[s]").tz_localize(tz)
+    res = time_ranges(starts, "D", 0)
+    pd.testing.assert_series_equal(res, pd.Series([], dtype=starts.dtype))
+
+
 # ints
 def test_time_ranges_with_integers_pd():
     dates = pd.Series([1, 10])

utilsforecast/processing.py

@@ -451,18 +451,24 @@ def _offset_times_array_n(
         raise ValueError(f"Expected n to have shape ({len(times)},), got {n.shape}.")
     if not np.issubdtype(n.dtype, np.integer):
         raise TypeError(f"Expected n to have an integer dtype, got {n.dtype}.")
+    if len(times) == 0:
+        return times.copy()
+    if n.dtype.kind == "u" and n.max() > np.iinfo(np.int64).max:
+        raise OverflowError("Expected n to fit in int64.")
     # a python int multiplier keeps n's dtype, so narrower ints would overflow
     n = n.astype(np.int64, copy=False)
     if isinstance(freq, pd.offsets.Tick):
-        # the coarsest resolution that can hold both the times and the offset,
-        # which is the one pandas uses with a scalar n
-        times_nanos = _NANOS_PER_UNIT[_dt_unit(times)]
-        unit, unit_nanos = next(
-            (u, k)
-            for u, k in _NANOS_PER_UNIT.items()
-            if k <= times_nanos and freq.nanos % k == 0
+        # pandas uses the finer of the times' and the offset's resolutions
+        unit = min(
+            _dt_unit(times), pd.Timedelta(freq).unit, key=_NANOS_PER_UNIT.__getitem__
         )
-        deltas = (n * (freq.nanos // unit_nanos)).astype(f"timedelta64[{unit}]")
+        step = freq.nanos // _NANOS_PER_UNIT[unit]
+        bound = np.iinfo(np.int64).max // max(abs(step), 1)
+        if n.max() > bound or n.min() < -bound:
+            raise pd.errors.OutOfBoundsTimedelta(
+                f"{freq} multiplied by n overflows timedelta64[{unit}]."
+            )
+        deltas = (n * step).astype(f"timedelta64[{unit}]")
         return times + deltas
     # pandas only vectorizes non-fixed offsets for a scalar n, so we apply each
     # distinct n to its own run of times and then restore the original order
@@ -478,14 +484,14 @@ def _offset_times_array_n(
         ]
     inv_order = np.empty_like(order)
     inv_order[order] = np.arange(order.size)
-    # the empty slice preserves the dtype in case there are no times
-    out = dt_index[:0].append(runs).take(inv_order)
+    out = runs[0].append(runs[1:]).take(inv_order)
+    if isinstance(times.dtype, pd.ArrowDtype):
+        # DatetimeIndex converts arrow times to numpy
+        out = out.astype(times.dtype)
     if isinstance(times, pd.Series):
         out = pd.Series(out, index=times.index, name=times.name)
     else:
         out = out.rename(times.name)
-    if out.dtype != times.dtype:
-        out = out.astype(times.dtype)
     return out
 
 
@@ -547,21 +553,20 @@ def time_ranges(
             starts_np = starts.to_numpy(copy=False)  # may be pyarrow
             steps = (np.arange(periods) * freq).astype(starts_np.dtype)
             out = np.repeat(starts_np, periods) + np.tile(steps, starts_np.size)
+            out = pd.Series(out, dtype=starts.dtype)
         elif _is_dt_dtype(starts):
-            if isinstance(freq, str):
-                freq = pd.tseries.frequencies.to_offset(freq)
-            out = []
-            with warnings.catch_warnings():
-                warnings.filterwarnings("ignore", category=pd.errors.PerformanceWarning)
-                for i in range(periods):
-                    out.append([starts + i * freq])
-            # pyarrow timestamps don't seem to work with offsets yet, keeping np.vstack
-            out = np.vstack(out).ravel(order="F")
+            per_period = [offset_times(starts, freq, i) for i in range(periods)]
+            if per_period:
+                series_major = np.arange(starts.size * periods).reshape(periods, -1)
+                out = per_period[0].append(per_period[1:])
+                out = out.take(series_major.ravel(order="F"))
+            else:
+                out = starts[:0]
+            out = pd.Series(out.rename(None))
         else:
             raise ValueError(
                 f"`starts` must be integers or timestamps, got '{starts.dtype}'."
             )
-        out = pd.Series(out, dtype=starts.dtype)
     else:
         try:
             is_int = starts.dtype.is_integer()