python polars 0.13.37 (#3426)

* allow count in window expressions * python polars 0.13.37
pola-rs · May 18, 2022 · 1ca735d · 1ca735d
1 parent 8df2ba4
commit 1ca735d
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 2 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/Cargo.lock b/py-polars/Cargo.lock
diff --git a/py-polars/Cargo.toml b/py-polars/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "py-polars"
-version = "0.13.36"
+version = "0.13.37"
 authors = ["ritchie46 <ritchie46@gmail.com>"]
 documentation = "https://pola-rs.github.io/polars/py-polars/html/reference/index.html"
 edition = "2021"

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]
+    )