python: pl.all()/pl.any() accept expr

py-polars/polars/internals/lazy_functions.py

@@ -871,9 +871,9 @@ def fold(
     return pli.wrap_expr(pyfold(acc._pyexpr, f, exprs))
 
 
-def any(name: str | list[pli.Expr]) -> pli.Expr:
+def any(name: str | list[pli.Expr] | pli.Expr) -> pli.Expr:
     """Evaluate columnwise or elementwise with a bitwise OR operation."""
-    if isinstance(name, list):
+    if isinstance(name, (list, pli.Expr)):
         return fold(lit(False), lambda a, b: a.cast(bool) | b.cast(bool), name).alias(
             "any"
         )
@@ -982,7 +982,7 @@ def exclude(
     return col("*").exclude(columns)
 
 
-def all(name: str | list[pli.Expr] | None = None) -> pli.Expr:
+def all(name: str | list[pli.Expr] | pli.Expr | None = None) -> pli.Expr:
     """
     Do one of two things.
 
@@ -1014,7 +1014,7 @@ def all(name: str | list[pli.Expr] | None = None) -> pli.Expr:
     """
     if name is None:
         return col("*")
-    if isinstance(name, list):
+    if isinstance(name, (list, pli.Expr)):
         return fold(lit(True), lambda a, b: a.cast(bool) & b.cast(bool), name).alias(
             "all"
         )

py-polars/tests/test_lazy.py

@@ -1404,3 +1404,21 @@ def test_type_coercion_unknown_4190() -> None:
         .lazy()
         .with_columns([pl.col("a") & pl.col("a").fill_null(True)])
     ).collect().shape == (3, 2)
+
+
+def test_all_any_accept_expr() -> None:
+    df = pl.DataFrame(
+        {
+            "a": [1, None, 2],
+            "b": [1, 2, None],
+        }
+    )
+    assert df.select(
+        [
+            pl.any(pl.all().is_null()).alias("null_in_row"),
+            pl.all(pl.all().is_null()).alias("all_null_in_row"),
+        ]
+    ).to_dict(False) == {
+        "null_in_row": [False, True, True],
+        "all_null_in_row": [False, False, False],
+    }