diff --git a/polars/polars-lazy/polars-plan/src/dsl/function_expr/temporal.rs b/polars/polars-lazy/polars-plan/src/dsl/function_expr/temporal.rs index c9d3359ce319..ca01ecfeb52e 100644 --- a/polars/polars-lazy/polars-plan/src/dsl/function_expr/temporal.rs +++ b/polars/polars-lazy/polars-plan/src/dsl/function_expr/temporal.rs @@ -8,7 +8,7 @@ use super::*; #[cfg(feature = "date_offset")] pub(super) fn date_offset(s: Series, offset: Duration) -> PolarsResult { - match s.dtype().clone() { + let out = match s.dtype().clone() { DataType::Date => { let s = s .cast(&DataType::Datetime(TimeUnit::Milliseconds, None)) @@ -42,7 +42,11 @@ pub(super) fn date_offset(s: Series, offset: Duration) -> PolarsResult { dt => polars_bail!( ComputeError: "cannot use 'date_offset' on Series of datatype {}", dt, ), - } + }; + out.map(|mut out| { + out.set_sorted_flag(s.is_sorted_flag()); + out + }) } pub(super) fn combine(s: &[Series], tu: TimeUnit) -> PolarsResult { diff --git a/py-polars/tests/unit/namespaces/test_datetime.py b/py-polars/tests/unit/namespaces/test_datetime.py index baddc791e6bc..eaeddfa4d99b 100644 --- a/py-polars/tests/unit/namespaces/test_datetime.py +++ b/py-polars/tests/unit/namespaces/test_datetime.py @@ -540,6 +540,16 @@ def test_negative_offset_by_err_msg_8464() -> None: pl.Series([datetime(2022, 3, 30)]).dt.offset_by("-1mo") +def test_offset_by_sorted_flag() -> None: + s = pl.Series([datetime(2001, 1, 1), datetime(2001, 1, 2)]) + s = s.set_sorted() + + assert s.flags["SORTED_ASC"] + s1 = s.dt.offset_by("1d") + assert s1.to_list() == [datetime(2001, 1, 2), datetime(2001, 1, 3)] + assert s1.flags["SORTED_ASC"] + + @pytest.mark.parametrize( ("duration", "input_date", "expected"), [