feat(rust, python): keep sorted flag after Expr::truncate (pola-rs#9275)

c-peters · Jul 14, 2023 · 063e3ae · 063e3ae
1 parent 35aa3ed
commit 063e3ae
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/polars/polars-lazy/polars-plan/src/dsl/function_expr/datetime.rs b/polars/polars-lazy/polars-plan/src/dsl/function_expr/datetime.rs
@@ -181,7 +181,7 @@ pub(super) fn timestamp(s: &Series, tu: TimeUnit) -> PolarsResult<Series> {
 pub(super) fn truncate(s: &Series, every: &str, offset: &str) -> PolarsResult<Series> {
     let every = Duration::parse(every);
     let offset = Duration::parse(offset);
-    Ok(match s.dtype() {
+    let mut out = match s.dtype() {
         DataType::Datetime(_, tz) => match tz {
             #[cfg(feature = "timezones")]
             Some(tz) => s
@@ -201,7 +201,9 @@ pub(super) fn truncate(s: &Series, every: &str, offset: &str) -> PolarsResult<Se
             .truncate(every, offset, None)?
             .into_series(),
         dt => polars_bail!(opq = round, got = dt, expected = "date/datetime"),
-    })
+    };
+    out.set_sorted_flag(s.is_sorted_flag());
+    Ok(out)
 }
 
 #[cfg(feature = "date_offset")]

diff --git a/py-polars/tests/unit/namespaces/test_datetime.py b/py-polars/tests/unit/namespaces/test_datetime.py
@@ -540,14 +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:
+def test_offset_by_truncate_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"]
+    s2 = s1.dt.truncate("1mo")
+    assert s2.flags["SORTED_ASC"]
 
 
 @pytest.mark.parametrize(