feat: expose array_filter higher-order function

Add array_filter, the remaining lambda-based higher-order array function
in DataFusion (alongside the already-exposed array_transform and
array_any_match). Includes the list_filter alias matching upstream, tests,
and documentation in the expressions guide and skill.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: timsaucer

crates/core/src/functions.rs

@@ -190,6 +190,13 @@ fn array_any_match(array: PyExpr, predicate: PyExpr) -> PyExpr {
     datafusion::functions_nested::expr_fn::array_any_match(array.into(), predicate.into()).into()
 }
 
+/// Higher-order function: keep the elements of `array` for which `predicate`
+/// (a lambda returning a boolean) is true, returning a new filtered array.
+#[pyfunction]
+fn array_filter(array: PyExpr, predicate: PyExpr) -> PyExpr {
+    datafusion::functions_nested::expr_fn::array_filter(array.into(), predicate.into()).into()
+}
+
 /// Computes a binary hash of the given data. type is the algorithm to use.
 /// Standard algorithms are md5, sha224, sha256, sha384, sha512, blake2s, blake2b, and blake3.
 // #[pyfunction(value, method)]
@@ -1118,6 +1125,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(lambda_var))?;
     m.add_wrapped(wrap_pyfunction!(array_transform))?;
     m.add_wrapped(wrap_pyfunction!(array_any_match))?;
+    m.add_wrapped(wrap_pyfunction!(array_filter))?;
 
     // Array Functions
     m.add_wrapped(wrap_pyfunction!(array_append))?;

docs/source/user-guide/common-operations/expressions.rst

@@ -150,8 +150,10 @@ Higher-order functions and lambdas
 
 Some array functions are *higher-order*: they take a lambda that runs once per
 element. :py:func:`~datafusion.functions.array_transform` maps a lambda over
-every element, and :py:func:`~datafusion.functions.array_any_match` returns
-whether any element satisfies a predicate lambda.
+every element, :py:func:`~datafusion.functions.array_filter` keeps the elements
+for which a predicate lambda is true, and
+:py:func:`~datafusion.functions.array_any_match` returns whether any element
+satisfies a predicate lambda.
 
 The simplest way to supply a lambda is a Python ``lambda``. Its parameter names
 become the lambda parameters, and its return value becomes the body.
@@ -164,6 +166,7 @@ become the lambda parameters, and its return value becomes the body.
     ctx = SessionContext()
     df = ctx.from_pydict({"a": [[1, 2, 3], [4, 5]]})
     df.select(f.array_transform(col("a"), lambda v: v * 2).alias("doubled"))
+    df.select(f.array_filter(col("a"), lambda v: v > 2).alias("big_only"))
     df.select(f.array_any_match(col("a"), lambda v: v > 3).alias("has_big"))
 
 If you need explicit control over parameter names, build the lambda with

python/datafusion/functions.py

@@ -85,6 +85,7 @@
     "array_empty",
     "array_except",
     "array_extract",
+    "array_filter",
     "array_has",
     "array_has_all",
     "array_has_any",
@@ -217,6 +218,7 @@
     "list_empty",
     "list_except",
     "list_extract",
+    "list_filter",
     "list_has",
     "list_has_all",
     "list_has_any",
@@ -632,6 +634,47 @@ def list_any_match(array: Expr, predicate: Expr | Callable[..., Any]) -> Expr:
     return array_any_match(array, predicate)
 
 
+def array_filter(array: Expr, predicate: Expr | Callable[..., Any]) -> Expr:
+    """Keep the elements of ``array`` for which ``predicate`` is ``True``.
+
+    ``predicate`` may be a Python callable, converted to a lambda
+    automatically, or an explicit lambda built with :py:func:`lambda_`. It must
+    return a boolean expression. The result is a new array containing only the
+    matching elements.
+
+    Examples:
+        Using a Python callable:
+
+        >>> ctx = dfn.SessionContext()
+        >>> df = ctx.from_pydict({"a": [[1, 2, 3, 4, 5]]})
+        >>> df.select(
+        ...     F.array_filter(col("a"), lambda v: v > 2).alias("f")
+        ... ).collect_column("f")[0].as_py()
+        [3, 4, 5]
+
+        Using an explicit lambda built with :py:func:`lambda_`:
+
+        >>> predicate = F.lambda_(["v"], F.lambda_var("v") > lit(2))
+        >>> df.select(
+        ...     F.array_filter(col("a"), predicate).alias("f")
+        ... ).collect_column("f")[0].as_py()
+        [3, 4, 5]
+
+    See Also:
+        :py:func:`array_transform`, :py:func:`array_any_match`, :py:func:`lambda_`.
+    """
+    return Expr(f.array_filter(array.expr, _to_lambda(predicate).expr))
+
+
+def list_filter(array: Expr, predicate: Expr | Callable[..., Any]) -> Expr:
+    """Keep the elements of a list for which a predicate is ``True``.
+
+    See Also:
+        This is an alias for :py:func:`array_filter`.
+    """
+    return array_filter(array, predicate)
+
+
 def in_list(arg: Expr, values: list[Expr], negated: bool = False) -> Expr:
     """Returns whether the argument is contained within the list ``values``.
 

python/tests/test_lambda.py

@@ -78,6 +78,23 @@ def _column(df, expr, name):
             [False, True],
             id="list_any_match_alias",
         ),
+        pytest.param(
+            lambda: f.array_filter(col("a"), lambda v: v > 2),
+            [[3], [4, 5]],
+            id="array_filter_callable",
+        ),
+        pytest.param(
+            lambda: f.array_filter(
+                col("a"), f.lambda_(["v"], f.lambda_var("v") > lit(2))
+            ),
+            [[3], [4, 5]],
+            id="array_filter_explicit_lambda",
+        ),
+        pytest.param(
+            lambda: f.list_filter(col("a"), lambda v: v > 2),
+            [[3], [4, 5]],
+            id="list_filter_alias",
+        ),
     ],
 )
 def test_higher_order_function_results(df, build_expr, expected):

skills/datafusion_python/SKILL.md

@@ -496,11 +496,12 @@ return value becomes the body:
 
 ```python
 F.array_transform(col("a"), lambda v: v * 2)    # map: [1,2,3] -> [2,4,6]
+F.array_filter(col("a"), lambda v: v > 2)        # filter: [1,2,3] -> [3]
 F.array_any_match(col("a"), lambda v: v > 3)     # predicate: any element > 3
 ```
 
-Aliases: `list_transform` for `array_transform`; `any_match` / `list_any_match`
-for `array_any_match`.
+Aliases: `list_transform` for `array_transform`; `list_filter` for
+`array_filter`; `any_match` / `list_any_match` for `array_any_match`.
 
 For explicit parameter names, build the lambda by hand: