fix count union predicate (#3969)

polars/polars-lazy/src/logical_plan/optimizer/predicate_pushdown/mod.rs

@@ -500,11 +500,7 @@ impl PredicatePushDown {
                     }
                 });
                 let lp = self.pushdown_and_continue(lp, acc_predicates, lp_arena, expr_arena, false)?;
-                Ok(if local_predicates.is_empty() {
-                    self.optional_apply_predicate(lp, local_predicates, lp_arena, expr_arena)
-                } else {
-                    lp
-                })
+                Ok(self.optional_apply_predicate(lp, local_predicates, lp_arena, expr_arena))
             }
             // Pushed down passed these nodes
             lp @ Cache { .. } |  lp @ Sort { .. } => {

polars/tests/it/lazy/predicate_queries.rs

@@ -186,3 +186,33 @@ fn test_predicate_pushdown_blocked_by_outer_join() -> Result<()> {
     assert!(out.frame_equal_missing(&expected));
     Ok(())
 }
+
+#[test]
+fn test_count_blocked_at_union_3963() -> Result<()> {
+    let lf1 = df![
+        "k" => ["x", "x", "y"],
+        "v" => [3, 2, 6,]
+    ]?
+    .lazy();
+
+    let lf2 = df![
+        "k" => ["a", "a", "b"],
+        "v" => [1, 8, 5]
+    ]?
+    .lazy();
+
+    let expected = df![
+        "k" => ["x", "x", "a", "a"],
+        "v" => [3, 2, 1, 8]
+    ]?;
+
+    for rechunk in [true, false] {
+        let out = concat([lf1.clone(), lf2.clone()], rechunk)?
+            .filter(count().over([col("k")]).gt(lit(1)))
+            .collect()?;
+
+        assert!(out.frame_equal(&expected));
+    }
+
+    Ok(())
+}