fix: rolling_* aggs were behaving as if they return scalars in group-…

…by (#15657)
pola-rs · Apr 15, 2024 · 4e17d02 · 4e17d02
1 parent 45620ef
commit 4e17d02
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/crates/polars-plan/src/dsl/mod.rs b/crates/polars-plan/src/dsl/mod.rs
@@ -1208,7 +1208,7 @@ impl Expr {
             self.apply_many_private(
                 FunctionExpr::RollingExpr(rolling_function_by(options)),
                 &[col(&name)],
-                true,
+                false,
                 false,
             )
         } else {

diff --git a/py-polars/tests/unit/operations/rolling/test_rolling.py b/py-polars/tests/unit/operations/rolling/test_rolling.py
@@ -917,3 +917,16 @@ def test_rolling_min_periods(
         )
     )["value"]
     assert_series_equal(result, pl.Series("value", expected, pl.Int64))
+
+
+def test_rolling_returns_scalar_15656() -> None:
+    df = pl.DataFrame(
+        {
+            "a": [date(2020, 1, 1), date(2020, 1, 2), date(2020, 1, 3)],
+            "b": [4, 5, 6],
+            "c": [1, 2, 3],
+        }
+    )
+    result = df.group_by("c").agg(pl.col("b").rolling_mean("2d", by="a")).sort("c")
+    expected = pl.DataFrame({"c": [1, 2, 3], "b": [[4.0], [5.0], [6.0]]})
+    assert_frame_equal(result, expected)