diff --git a/polars/polars-time/src/groupby/dynamic.rs b/polars/polars-time/src/groupby/dynamic.rs index b8592c0cff21..27ad69450446 100644 --- a/polars/polars-time/src/groupby/dynamic.rs +++ b/polars/polars-time/src/groupby/dynamic.rs @@ -7,7 +7,7 @@ use polars_core::series::IsSorted; use polars_core::utils::ensure_sorted_arg; use polars_core::utils::flatten::flatten_par; use polars_core::POOL; -use polars_utils::slice::SortedSlice; +use polars_utils::slice::{GetSaferUnchecked, SortedSlice}; #[cfg(feature = "serde")] use serde::{Deserialize, Serialize}; use smartstring::alias::String as SmartString; @@ -638,15 +638,19 @@ fn update_subgroups_idx( sub_groups .iter() .map(|&[first, len]| { - let new_first = unsafe { *base_g.1.get_unchecked(first as usize) }; + let new_first = if len == 0 { + // in case the group is empty + // keep the original first so that the + // groupby keys still point to the original group + base_g.0 + } else { + unsafe { *base_g.1.get_unchecked_release(first as usize) } + }; let first = first as usize; let len = len as usize; let idx = (first..first + len) - .map(|i| { - debug_assert!(i < base_g.1.len()); - unsafe { *base_g.1.get_unchecked(i) } - }) + .map(|i| unsafe { *base_g.1.get_unchecked_release(i) }) .collect_trusted::>(); (new_first, idx) }) diff --git a/py-polars/tests/unit/operations/test_groupby_rolling.py b/py-polars/tests/unit/operations/test_groupby_rolling.py index f628ec27cfcd..d06bdbed75ec 100644 --- a/py-polars/tests/unit/operations/test_groupby_rolling.py +++ b/py-polars/tests/unit/operations/test_groupby_rolling.py @@ -1,6 +1,6 @@ from __future__ import annotations -from datetime import datetime +from datetime import date, datetime from typing import TYPE_CHECKING, Any import pytest @@ -258,3 +258,42 @@ def test_groupby_rolling_dynamic_sortedness_check() -> None: match=r"argument in operation 'groupby_rolling' is not explicitly sorted", ): df.groupby_rolling("idx", period="2i").agg(pl.col("idx").alias("idx1")) + + +def test_groupby_rolling_empty_groups_9973() -> None: + dt1 = date(2001, 1, 1) + dt2 = date(2001, 1, 2) + + data = pl.DataFrame( + { + "id": ["A", "A", "B", "B", "C", "C"], + "date": [dt1, dt2, dt1, dt2, dt1, dt2], + "value": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0], + } + ).sort(by=["id", "date"]) + + expected = pl.DataFrame( + { + "id": ["A", "A", "B", "B", "C", "C"], + "date": [ + date(2001, 1, 1), + date(2001, 1, 2), + date(2001, 1, 1), + date(2001, 1, 2), + date(2001, 1, 1), + date(2001, 1, 2), + ], + "value": [[2.0], [], [4.0], [], [6.0], []], + } + ) + + out = data.groupby_rolling( + index_column="date", + by="id", + period="2d", + offset="1d", + closed="left", + check_sorted=True, + ).agg(pl.col("value")) + + assert_frame_equal(out, expected)