Skip to content

Commit

Permalink
fix: offset=-0i was being treated differently to offset=0i in rolling
Browse files Browse the repository at this point in the history
  • Loading branch information
MarcoGorelli committed May 13, 2024
1 parent 3892c75 commit 54e0af0
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 2 deletions.
4 changes: 2 additions & 2 deletions crates/polars-time/src/windows/group_by.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
{
Expand Down
12 changes: 12 additions & 0 deletions py-polars/tests/unit/operations/test_rolling.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

0 comments on commit 54e0af0

Please sign in to comment.