fix cumulative_eval on window expressions (#3421)

pola-rs · May 18, 2022 · 7911011 · 7911011
1 parent 440681f
commit 7911011
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 1 deletion.
diff --git a/polars/polars-lazy/src/dsl/eval.rs b/polars/polars-lazy/src/dsl/eval.rs
@@ -84,7 +84,7 @@ impl Expr {
             Ok(Series::new(&name, avs))
         };
 
-        self.map(
+        self.apply(
             func,
             GetOutput::map_field(move |f| {
                 // dummy df to determine output dtype

diff --git a/py-polars/tests/test_window.py b/py-polars/tests/test_window.py
@@ -138,3 +138,23 @@ def test_quantile_as_window() -> None:
         .to_series()
         .series_equal(pl.Series("value", [1.0, 1.0, 2.0, 2.0]))
     )
+
+
+def test_cumulative_eval_window_functions() -> None:
+    df = pl.DataFrame(
+        {
+            "group": [0, 0, 0, 1, 1, 1],
+            "val": [20, 40, 30, 2, 4, 3],
+        }
+    )
+
+    assert df.with_column(
+        pl.col("val")
+        .cumulative_eval(pl.element().max())
+        .over("group")
+        .alias("cumulative_eval_max")
+    ).to_dict(False) == {
+        "group": [0, 0, 0, 1, 1, 1],
+        "val": [20, 40, 30, 2, 4, 3],
+        "cumulative_eval_max": [20, 40, 40, 2, 4, 4],
+    }