fix: support string, date, and datetime bounds in polars fill_gaps

tests/test_preprocessing.py

@@ -523,3 +523,109 @@ def test_polars_str_freq_with_int_data(self, polars_freq_data, freq_int, include
                 data = create_error_test_data(dates_int, N_PERIODS, include_start, include_end, "polars")
                 with pytest.raises(ValueError):
                     fill_gaps(data, freq_str, start=start, end=end)
+
+
+def test_fill_gaps_polars_custom_bounds():
+    df_date = pl.DataFrame(
+        {
+            "unique_id": [1, 1],
+            "ds": [date(2020, 1, 2), date(2020, 1, 4)],
+            "y": [1.0, 2.0],
+        }
+    )
+    res_str = fill_gaps(df_date, freq="1d", start="2020-01-01", end="2020-01-05")
+    assert res_str["ds"].to_list() == [
+        date(2020, 1, 1),
+        date(2020, 1, 2),
+        date(2020, 1, 3),
+        date(2020, 1, 4),
+        date(2020, 1, 5),
+    ]
+
+    res_dt = fill_gaps(
+        df_date,
+        freq="1d",
+        start=datetime(2020, 1, 1),
+        end=datetime(2020, 1, 5),
+    )
+    assert res_dt["ds"].to_list() == [
+        date(2020, 1, 1),
+        date(2020, 1, 2),
+        date(2020, 1, 3),
+        date(2020, 1, 4),
+        date(2020, 1, 5),
+    ]
+
+    df_datetime = pl.DataFrame(
+        {
+            "unique_id": [1, 1],
+            "ds": [datetime(2020, 1, 2), datetime(2020, 1, 4)],
+            "y": [1.0, 2.0],
+        }
+    )
+    res_dt_str = fill_gaps(df_datetime, freq="1d", start="2020-01-01", end="2020-01-05")
+    assert res_dt_str["ds"].to_list() == [
+        datetime(2020, 1, 1),
+        datetime(2020, 1, 2),
+        datetime(2020, 1, 3),
+        datetime(2020, 1, 4),
+        datetime(2020, 1, 5),
+    ]
+
+    res_dt_date = fill_gaps(
+        df_datetime,
+        freq="1d",
+        start=date(2020, 1, 1),
+        end=date(2020, 1, 5),
+    )
+    assert res_dt_date["ds"].to_list() == [
+        datetime(2020, 1, 1),
+        datetime(2020, 1, 2),
+        datetime(2020, 1, 3),
+        datetime(2020, 1, 4),
+        datetime(2020, 1, 5),
+    ]
+
+    res_date_dt_str = fill_gaps(df_date, freq="1d", start="2020-01-01 00:00:00")
+    assert res_date_dt_str["ds"].to_list()[0] == date(2020, 1, 1)
+
+
+@pytest.mark.parametrize(
+    "start", ["2020-01-01", datetime(2020, 1, 1), date(2020, 1, 1)]
+)
+def test_fill_gaps_polars_naive_bound_tz_aware_column(start):
+    df = pl.DataFrame(
+        {
+            "unique_id": [1, 1],
+            "ds": [datetime(2020, 1, 2), datetime(2020, 1, 4)],
+            "y": [1.0, 2.0],
+        }
+    ).with_columns(pl.col("ds").dt.replace_time_zone("America/New_York"))
+    res = fill_gaps(df, freq="1d", start=start)
+    expected = pl.datetime_range(
+        datetime(2020, 1, 1),
+        datetime(2020, 1, 4),
+        "1d",
+        time_unit="us",
+        time_zone="America/New_York",
+        eager=True,
+    )
+    assert res["ds"].to_list() == expected.to_list()
+    assert res["y"].drop_nulls().to_list() == [1.0, 2.0]
+
+
+@pytest.mark.parametrize(
+    "ds, start",
+    [
+        ([2, 4], date(2020, 1, 1)),
+        ([2, 4], "2020-01-01"),
+        ([date(2020, 1, 2), date(2020, 1, 4)], 5),
+        ([datetime(2020, 1, 2), datetime(2020, 1, 4)], 5),
+        ([datetime(2020, 1, 2), datetime(2020, 1, 4)], "2020-01-01T00:00:00+02:00"),
+    ],
+)
+def test_fill_gaps_polars_mismatched_bound(ds, start):
+    df = pl.DataFrame({"unique_id": [1, 1], "ds": ds, "y": [1.0, 2.0]})
+    freq = 1 if isinstance(ds[0], int) else "1d"
+    with pytest.raises(ValueError):
+        fill_gaps(df, freq=freq, start=start)

utilsforecast/preprocessing.py

@@ -39,19 +39,57 @@ def _determine_bound(bound, freq, times_by_id, agg) -> np.ndarray:
     return out
 
 
+def _validate_bound_pl(
+    bound: Union[str, int, date, datetime],
+    name: str,
+    times: pl_Series,
+) -> None:
+    if isinstance(bound, str) and bound in ("per_serie", "global"):
+        return
+    if times.dtype.is_temporal():
+        expected = "a string, date or datetime"
+        valid = isinstance(bound, (str, date))
+    else:
+        expected = "an integer"
+        valid = isinstance(bound, (int, np.integer))
+    if not valid:
+        raise ValueError(
+            f"Time column has type {times.dtype} but `{name}` is {bound!r}. "
+            f"When using polars dataframes `{name}` must be 'per_serie', 'global' "
+            f"or {expected}."
+        )
+
+
 def _determine_bound_pl(
     bound: Union[str, int, date, datetime],
     times_by_id: pl_DataFrame,
     agg: str,
 ) -> pl_Series:
+    target_dtype = times_by_id[agg].dtype
     if bound == "per_serie":
         out = times_by_id[agg]
     else:
         if bound == "global":
             val = getattr(times_by_id[agg], agg)()
+            val_s = pl_Series([val], dtype=target_dtype)
         else:
-            val = bound
-        out = repeat(pl_Series([val]), times_by_id.shape[0])
+            val_s = pl_Series([bound])
+            if isinstance(bound, str):
+                val_s = val_s.str.to_datetime()
+            target_tz = getattr(target_dtype, "time_zone", None)
+            bound_tz = getattr(val_s.dtype, "time_zone", None)
+            if bound_tz is not None and target_tz is None:
+                raise ValueError(
+                    f"Bound {bound!r} is timezone-aware but the time column "
+                    f"has type {target_dtype}, which has no timezone."
+                )
+            if target_tz is not None and bound_tz is None:
+                # naive bounds are wall times in the column's timezone
+                if val_s.dtype == pl.Date:
+                    val_s = val_s.cast(pl.Datetime(target_dtype.time_unit))
+                val_s = val_s.dt.replace_time_zone(target_tz)
+            val_s = val_s.cast(target_dtype)
+        out = repeat(val_s, times_by_id.shape[0])
     return out
 
 
@@ -85,6 +123,8 @@ def id_time_grid(
         pandas or polars DataFrame: Dataframe with expected ids and times.
     """
     if isinstance(df, pl_DataFrame):
+        _validate_bound_pl(start, "start", df[time_col])
+        _validate_bound_pl(end, "end", df[time_col])
         times_by_id = (
             group_by(df, id_col)
             .agg(