feat(rust, python): keep sorted flag on offset_by (pola-rs#9253)

c-peters · Jul 14, 2023 · b9769c8 · b9769c8
1 parent 0d4a81e
commit b9769c8
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
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
@@ -8,7 +8,7 @@ use super::*;
 
 #[cfg(feature = "date_offset")]
 pub(super) fn date_offset(s: Series, offset: Duration) -> PolarsResult<Series> {
-    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<Series> {
         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<Series> {

diff --git 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"),
     [