allow count in window expressions

pola-rs · May 18, 2022 · 44ca891 · 44ca891
1 parent 8df2ba4
commit 44ca891
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/polars/polars-lazy/src/physical_plan/planner/expr.rs b/polars/polars-lazy/src/physical_plan/planner/expr.rs
@@ -35,6 +35,8 @@ impl DefaultPlanner {
                 if apply_columns.is_empty() {
                     if has_aexpr(function, expr_arena, |e| matches!(e, AExpr::Literal(_))) {
                         apply_columns.push(Arc::from("literal"))
+                    } else if has_aexpr(function, expr_arena, |e| matches!(e, AExpr::Count)) {
+                        apply_columns.push(Arc::from("count"))
                     } else {
                         let e = node_to_expr(function, expr_arena);
                         return Err(PolarsError::ComputeError(

diff --git a/py-polars/tests/test_window.py b/py-polars/tests/test_window.py
@@ -158,3 +158,16 @@ def test_cumulative_eval_window_functions() -> None:
         "val": [20, 40, 30, 2, 4, 3],
         "cumulative_eval_max": [20, 40, 40, 2, 4, 4],
     }
+
+
+def test_count_window() -> None:
+    assert (
+        pl.DataFrame(
+            {
+                "a": [1, 1, 2],
+            }
+        )
+        .with_column(pl.count().over("a"))["count"]
+        .to_list()
+        == [2, 2, 1]
+    )