fix: offset=-0i was being treated differently to offset=0i in rolling (…

…pola-rs#16184)
Wouittone · Jun 22, 2024 · 5311f44 · 5311f44
1 parent 0137b75
commit 5311f44
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 3 deletions.
diff --git a/crates/polars-time/src/windows/group_by.rs b/crates/polars-time/src/windows/group_by.rs
@@ -576,7 +576,7 @@ pub fn group_by_values(
     let run_parallel = !POOL.current_thread_has_pending_tasks().unwrap_or(false);
 
     // we have a (partial) lookbehind window
-    if offset.negative {
+    if offset.negative && offset.duration_ns() > 0 {
         // lookbehind
         if offset.duration_ns() == period.duration_ns() {
             // t is right at the end of the window
@@ -647,7 +647,7 @@ pub fn group_by_values(
             iter.map(|result| result.map(|(offset, len)| [offset, len]))
                 .collect::<PolarsResult<_>>()
         }
-    } else if offset != Duration::parse("0ns")
+    } else if offset.duration_ns() != 0
         || closed_window == ClosedWindow::Right
         || closed_window == ClosedWindow::None
     {

diff --git a/py-polars/tests/parametric/test_groupby_rolling.py b/py-polars/tests/parametric/test_groupby_rolling.py
@@ -29,7 +29,7 @@ def interval_defs() -> SearchStrategy[ClosedInterval]:
         min_value=timedelta(microseconds=0), max_value=timedelta(days=1000)
     ).map(parse_as_duration_string),
     offset=st.timedeltas(
-        min_value=timedelta(microseconds=0), max_value=timedelta(days=1000)
+        min_value=timedelta(days=-1000), max_value=timedelta(days=1000)
     ).map(parse_as_duration_string),
     closed=interval_defs(),
     data=st.data(),

diff --git a/py-polars/tests/unit/operations/test_rolling.py b/py-polars/tests/unit/operations/test_rolling.py
@@ -321,3 +321,15 @@ def test_multiple_rolling_in_single_expression() -> None:
         front_count.alias("front"),
         (back_count - front_count).alias("back - front"),
     )["back - front"].to_list() == [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 4, 5]
+
+
+def test_negative_zero_offset_16168() -> None:
+    df = pl.DataFrame({"foo": [1] * 3}).sort("foo").with_row_index()
+    result = df.rolling(index_column="foo", period="1i", offset="0i").agg("index")
+    expected = pl.DataFrame(
+        {"foo": [1, 1, 1], "index": [[], [], []]},
+        schema_overrides={"index": pl.List(pl.UInt32)},
+    )
+    assert_frame_equal(result, expected)
+    result = df.rolling(index_column="foo", period="1i", offset="-0i").agg("index")
+    assert_frame_equal(result, expected)