diff --git a/crates/polars-time/src/truncate.rs b/crates/polars-time/src/truncate.rs index 42cbc16645bf..991ce50b547a 100644 --- a/crates/polars-time/src/truncate.rs +++ b/crates/polars-time/src/truncate.rs @@ -37,7 +37,7 @@ impl PolarsTruncate for DatetimeChunked { return Ok(self .apply_values(|t| { let remainder = t % every; - t - remainder + every * (remainder < 0) as i64 + t - (remainder + every * (remainder < 0) as i64) }) .into_datetime(self.time_unit(), time_zone.clone())); } else { diff --git a/py-polars/polars/expr/datetime.py b/py-polars/polars/expr/datetime.py index 59e481d2995d..cdf6ccb6516f 100644 --- a/py-polars/polars/expr/datetime.py +++ b/py-polars/polars/expr/datetime.py @@ -270,9 +270,8 @@ def truncate(self, every: str | dt.timedelta | Expr) -> Expr: │ 2001-01-01 01:00:00 ┆ 2001-01-01 01:00:00 │ └─────────────────────┴─────────────────────┘ """ - if not isinstance(every, pl.Expr): + if isinstance(every, dt.timedelta): every = parse_as_duration_string(every) - every = parse_into_expression(every, str_as_lit=True) return wrap_expr(self._pyexpr.dt_truncate(every)) diff --git a/py-polars/polars/series/datetime.py b/py-polars/polars/series/datetime.py index c65e822e1044..8c8bfb32bad8 100644 --- a/py-polars/polars/series/datetime.py +++ b/py-polars/polars/series/datetime.py @@ -1642,7 +1642,7 @@ def offset_by(self, by: str | Expr) -> Series: ] """ - def truncate(self, every: str | dt.timedelta | Expr) -> Series: + def truncate(self, every: str | dt.timedelta | IntoExprColumn) -> Series: """ Divide the date/ datetime range into buckets. diff --git a/py-polars/tests/unit/operations/namespaces/temporal/test_truncate.py b/py-polars/tests/unit/operations/namespaces/temporal/test_truncate.py index b679d3859a4c..f56d356b0457 100644 --- a/py-polars/tests/unit/operations/namespaces/temporal/test_truncate.py +++ b/py-polars/tests/unit/operations/namespaces/temporal/test_truncate.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import date, datetime +from datetime import date, datetime, timedelta from typing import TYPE_CHECKING import hypothesis.strategies as st @@ -8,6 +8,7 @@ from hypothesis import given import polars as pl +from polars._utils.convert import parse_as_duration_string from polars.testing import assert_series_equal if TYPE_CHECKING: @@ -92,3 +93,29 @@ def test_truncate_datetime_w_expression(time_unit: TimeUnit) -> None: result = df.select(pl.col("a").dt.truncate(pl.col("b")))["a"] assert result[0] == datetime(2020, 1, 1) assert result[1] == datetime(2020, 1, 3) + + +def test_pre_epoch_truncate_17581() -> None: + s = pl.Series([datetime(1980, 1, 1), datetime(1969, 1, 1, 1)]) + result = s.dt.truncate("1d") + expected = pl.Series([datetime(1980, 1, 1), datetime(1969, 1, 1)]) + assert_series_equal(result, expected) + + +@given( + datetimes=st.lists( + st.datetimes(min_value=datetime(1960, 1, 1), max_value=datetime(1980, 1, 1)), + min_size=1, + max_size=3, + ), + every=st.timedeltas( + min_value=timedelta(microseconds=1), max_value=timedelta(days=1) + ).map(parse_as_duration_string), +) +def test_fast_path_vs_slow_path(datetimes: list[datetime], every: str) -> None: + s = pl.Series(datetimes) + # Might use fastpath: + result = s.dt.truncate(every) + # Definitely uses slowpath: + expected = s.dt.truncate(pl.Series([every] * len(datetimes))) + assert_series_equal(result, expected)