fix: fix nullable filter mask in group_by (#11207)

pola-rs · Sep 20, 2023 · 9ff5140 · 9ff5140
1 parent 4390604
commit 9ff5140
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/crates/polars-lazy/src/physical_plan/expressions/filter.rs b/crates/polars-lazy/src/physical_plan/expressions/filter.rs
@@ -73,7 +73,7 @@ impl PhysicalExpr for FilterExpr {
             let predicate = predicate_s.bool()?;
 
             // All values true - don't do anything.
-            if predicate.all() {
+            if let Some(true) = predicate.all_kleene() {
                 return Ok(ac_s);
             }
             // All values false - create empty groups.

diff --git a/py-polars/tests/unit/operations/test_filter.py b/py-polars/tests/unit/operations/test_filter.py
@@ -37,6 +37,22 @@ def test_melt_values_predicate_pushdown() -> None:
     ).to_dict(False) == {"id": [1], "variable": ["asset_key_1"], "value": ["123"]}
 
 
+def test_group_by_filter_all_true() -> None:
+    df = pl.DataFrame(
+        {
+            "name": ["a", "a", "b", "b"],
+            "type": [None, 1, 1, None],
+            "order": [1, 2, 3, 4],
+        }
+    )
+    out = (
+        df.group_by("name")
+        .agg([pl.col("order").filter(pl.col("type") == 1).n_unique().alias("n_unique")])
+        .select("n_unique")
+    )
+    assert out.to_dict(False) == {"n_unique": [1, 1]}
+
+
 def test_filter_is_in_4572() -> None:
     df = pl.DataFrame({"id": [1, 2, 1, 2], "k": ["a"] * 2 + ["b"] * 2})
     expected = (