feat: run filter lambdas natively with DataFusion's array_filter

docs/source/user-guide/latest/expressions.md

@@ -388,7 +388,7 @@ to Spark ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enab
 | `aggregate` | ✅ | Codegen dispatch |  |
 | `array_sort` | ✅ | Native | The default comparator runs natively; custom comparators and arrays with non-default string collations use the JVM codegen dispatcher |
 | `exists` | ✅ | Codegen dispatch |  |
-| `filter` | ✅ | Native | General lambda routed through the JVM codegen dispatcher; the `array_compact` form runs natively |
+| `filter` | ✅ | Native | A one-parameter lambda whose body is natively supported runs natively, as does the `array_compact` form; other lambdas use the JVM codegen dispatcher. See [Expression Tuning](tuning/expressions.md#lambdas-and-higher-order-functions) |
 | `forall` | ✅ | Codegen dispatch |  |
 | `map_filter` | ✅ | Codegen dispatch |  |
 | `map_zip_with` | ✅ | Codegen dispatch |  |

docs/source/user-guide/latest/index.rst

@@ -96,6 +96,7 @@ the navigation menu to read more.
    Shuffle <tuning/shuffle>
    Celeborn Shuffle <tuning/celeborn>
    Operators <tuning/operators>
+   Expressions <tuning/expressions>
    Row/Columnar Transitions <tuning/transitions>
 
 .. toctree::

docs/source/user-guide/latest/tuning.md

@@ -35,6 +35,8 @@ guide is split into the following pages:
   `spark.comet.scan.icebergNative.dataFileConcurrencyLimit`.
 - [Operator Tuning](tuning/operators.md): joins, adaptive partial aggregation, and sorting on
   floating-point values.
+- [Expression Tuning](tuning/expressions.md): lambdas and higher-order functions such as `filter`, and the
+  built-in array functions that replace them.
 - [Reducing Row/Columnar Conversion Overhead](tuning/transitions.md): stages in which many
   operators fall back to Spark, and the experimental direct columnar-to-row converter.
 

docs/source/user-guide/latest/tuning/expressions.md

@@ -0,0 +1,80 @@
+<!---
+Licensed to the Apache Software Foundation (ASF) under one
+or more contributor license agreements.  See the NOTICE file
+distributed with this work for additional information
+regarding copyright ownership.  The ASF licenses this file
+to you under the Apache License, Version 2.0 (the
+"License"); you may not use this file except in compliance
+with the License.  You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing,
+software distributed under the License is distributed on an
+"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+KIND, either express or implied.  See the License for the
+specific language governing permissions and limitations
+under the License.
+-->
+
+# Expression Tuning
+
+## Lambdas and Higher-Order Functions
+
+A higher-order function takes a lambda: `filter(arr, x -> ...)`, `transform`, `exists`, `forall`, `aggregate`,
+`zip_with` and `map_filter`. How fast Comet runs one depends on whether it can run the lambda natively.
+
+| Lambda                                                                                                                                                                    | How Comet runs it                                                                                                                                                                |
+| ------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `filter(arr, x -> x IS NOT NULL)` (what `array_compact` expands to), and `filter(arr, x -> nvl(x, '') != '')` on an array of numbers, dates or timestamps                 | Natively, as `array_compact`.                                                                                                                                                    |
+| `filter` with a one-parameter lambda whose body uses only natively supported expressions, over an array of primitive, decimal, string, binary, date or timestamp elements | Natively. The lambda runs once over all the elements of a batch, vectorized, instead of once per element.                                                                        |
+| Any other lambda, and every other higher-order function                                                                                                                   | Through the [JVM codegen dispatcher](../scala_java_udfs.md). Spark's interpreter evaluates the lambda for each element inside a Comet kernel, so expect Spark's speed, not more. |
+
+The dispatcher has no advantage over Spark for a lambda, because Spark also evaluates a lambda one element at a time.
+It adds the cost of handing each batch to the JVM and back. The rest of the query stays native either way.
+
+To get the most out of Comet when a query uses lambdas:
+
+- **Prefer a built-in array function to a lambda.** Most of them run natively with no per-element evaluation, for
+  example `array_compact`, `array_remove`, `array_distinct` and `array_position`. See the
+  [Supported Expressions](../expressions.md) tables for the caveats of each.
+  `array_compact(arr)` returns what `filter(arr, x -> x IS NOT NULL)` does for any array, and
+  `array_remove(array_compact(arr), '')` returns what `filter(arr, x -> nvl(x, '') != '')` does for an array of
+  strings with the default collation.
+- **Do not test a numeric array for nulls with a string comparison.** With ANSI mode off, Spark casts every element
+  of `filter(int_arr, x -> nvl(x, '') != '')` to a string to compare it with `''`; with ANSI mode on the query does
+  not analyze. Comet recognizes exactly this form, with either operand first, and runs it as `array_compact`, which
+  never casts. A variation, such as another default value or an extra condition, runs as a general lambda and does
+  cast every element. Write `x IS NOT NULL`, or call `array_compact`.
+- **Keep the lambda body to natively supported expressions.** A body that calls a regular expression function, a
+  UDF or another lambda cannot be split from the dispatcher, so the whole lambda runs there. Use
+  [`spark.comet.explain.codegen.enabled`](../understanding-comet-plans.md#sparkcometexplaincodegenenabled) to see which
+  expressions of a query went through the dispatcher.
+- **Keep `AND`, `OR` and `IN` away from operations that can raise.** Spark evaluates them lazily, so
+  `x != 0 AND 100 / x > 1` never divides by zero, but the native lambda evaluates every operand for every element. A
+  body that combines them with an operation that can raise stays on the dispatcher. That covers a division, an
+  addition or a cast under ANSI mode, which raise only in that mode, and `element_at`, which raises for the index 0 in
+  any mode. `CASE WHEN`, `IF` and `COALESCE` are lazy natively too, so
+  `CASE WHEN x = 0 THEN false ELSE 100 / x > 1 END` runs natively.
+- **Capture only numeric, boolean, date or timestamp columns.** DataFusion copies each column the lambda reads once per
+  element of the array before it evaluates the lambda, so a captured string, binary, array, struct or map column is
+  copied as many times as the array has elements. A lambda that reads one stays on the dispatcher.
+- **Use a lambda of one parameter.** `filter(arr, (x, i) -> i > 0)`, which also receives the element's index, runs
+  through the dispatcher.
+
+`spark.comet.exec.nativeLambda.enabled=false` makes Comet run every lambda through the dispatcher.
+
+### Example
+
+`CometHigherOrderFunctionBenchmark` filters 1 million rows of ten-element arrays, in which every fifth element is
+`NULL`. These are the best times of the Spark baseline and of Comet on an Apple M-series laptop, with
+`spark.comet.exec.nativeLambda.enabled` off and at its default:
+
+| Query                                                       | Spark  | Comet, nativeLambda off | Comet, default |
+| ----------------------------------------------------------- | ------ | ----------------------- | -------------- |
+| `filter(strs, x -> nvl(x, '') != '')` (string array)        | 616 ms | 547 ms                  | 307 ms         |
+| `array_remove(array_compact(strs), '')` (same result)       | 587 ms | 252 ms                  | 252 ms         |
+| `filter(ints, x -> x IS NOT NULL)`                          | 286 ms | 124 ms                  | 125 ms         |
+| `filter(ints, x -> x > threshold)`, `threshold` is a column | 346 ms | 335 ms                  | 110 ms         |
+
+Rewriting a lambda to built-in functions is still the fastest form, but a native lambda gets most of the way there.

native/core/src/execution/planner.rs

@@ -67,8 +67,11 @@ use datafusion::{
     execution::FunctionRegistry,
     logical_expr::Operator as DataFusionOperator,
     physical_expr::{
-        expressions::{BinaryExpr, CastExpr, Column, IsNullExpr, Literal as DataFusionLiteral},
-        PhysicalExpr, PhysicalSortExpr, ScalarFunctionExpr,
+        expressions::{
+            BinaryExpr, CastExpr, Column, IsNullExpr, LambdaExpr, LambdaVariable,
+            Literal as DataFusionLiteral,
+        },
+        HigherOrderFunctionExpr, PhysicalExpr, PhysicalSortExpr, ScalarFunctionExpr,
     },
     physical_plan::{
         empty::EmptyExec, limit::LocalLimitExec, sorts::sort::SortExec, ExecutionPlan,
@@ -93,6 +96,7 @@ use crate::parquet::parquet_support::{
 use datafusion::common::ScalarValue;
 use datafusion::datasource::listing::PartitionedFile;
 use datafusion::datasource::object_store::ObjectStoreUrl;
+use datafusion::functions_nested::array_filter::array_filter_higher_order_function;
 use datafusion::logical_expr::type_coercion::functions::fields_with_udf;
 use datafusion::logical_expr::{ReturnFieldArgs, ScalarUDF, TypeSignature};
 use datafusion::physical_expr::expressions::Literal;
@@ -1142,6 +1146,61 @@ impl PhysicalPlanner {
                 ));
                 Ok(expr)
             }
+            ExprStruct::HigherOrderFunc(hof) => {
+                let array =
+                    self.create_expr(hof.array.as_ref().unwrap(), Arc::clone(&input_schema))?;
+                // The lambda parameter takes the element field of the array as planned, since
+                // DataFusion checks it against the field of the list it evaluates.
+                let array_field = array.return_field(input_schema.as_ref())?;
+                let element_field = match array_field.data_type() {
+                    DataType::List(f) | DataType::LargeList(f) => f,
+                    other => {
+                        return Err(GeneralError(format!(
+                            "{} expects an array argument, got {other}",
+                            hof.func
+                        )))
+                    }
+                };
+                // The lambda body sees the enclosing schema with the parameter appended, so
+                // `BoundReference`s keep their indexes and a `LambdaVariable` is found by name.
+                let mut fields: Vec<FieldRef> = input_schema.fields().to_vec();
+                fields.push(Arc::new(
+                    element_field.as_ref().clone().with_name(&hof.param_name),
+                ));
+                let lambda_schema = Arc::new(Schema::new_with_metadata(
+                    fields,
+                    input_schema.metadata().clone(),
+                ));
+                let body = self.create_expr(hof.body.as_ref().unwrap(), lambda_schema)?;
+                let lambda: Arc<dyn PhysicalExpr> =
+                    Arc::new(LambdaExpr::try_new(vec![hof.param_name.clone()], body)?);
+                let func = match hof.func.as_str() {
+                    "array_filter" => array_filter_higher_order_function(),
+                    other => {
+                        return Err(GeneralError(format!(
+                            "Unsupported higher-order function '{other}'"
+                        )))
+                    }
+                };
+                Ok(Arc::new(HigherOrderFunctionExpr::try_new_with_schema(
+                    func,
+                    vec![array, lambda],
+                    input_schema.as_ref(),
+                    Arc::new(ConfigOptions::default()),
+                )?))
+            }
+            ExprStruct::LambdaVariable(var) => {
+                let idx = input_schema.index_of(&var.name).map_err(|_| {
+                    GeneralError(format!(
+                        "Lambda variable '{}' is not in scope. Schema: {input_schema}",
+                        var.name
+                    ))
+                })?;
+                Ok(Arc::new(LambdaVariable::new(
+                    idx,
+                    Arc::clone(&input_schema.fields()[idx]),
+                )))
+            }
             expr => Err(GeneralError(format!("Not implemented: {expr:?}"))),
         }
     }
@@ -5285,6 +5344,109 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_higher_order_filter() -> Result<(), DataFusionError> {
+        use arrow::array::ListBuilder;
+        use arrow::array::StringBuilder;
+        use datafusion_comet_proto::spark_expression::{
+            BinaryExpr, BoundReference, HigherOrderFunc, LambdaVariable,
+        };
+
+        fn expr(expr_struct: ExprStruct) -> spark_expression::Expr {
+            spark_expression::Expr {
+                expr_struct: Some(expr_struct),
+                query_context: None,
+                expr_id: None,
+            }
+        }
+        fn bound(index: i32) -> spark_expression::Expr {
+            expr(Bound(BoundReference {
+                index,
+                datatype: None,
+            }))
+        }
+        fn variable(name: &str) -> spark_expression::Expr {
+            expr(ExprStruct::LambdaVariable(LambdaVariable {
+                name: name.to_string(),
+            }))
+        }
+        fn string_lists(rows: Vec<Option<Vec<Option<&str>>>>) -> ListArray {
+            let mut builder = ListBuilder::new(StringBuilder::new());
+            for row in rows {
+                match row {
+                    Some(values) => {
+                        values
+                            .into_iter()
+                            .for_each(|v| builder.values().append_option(v));
+                        builder.append(true);
+                    }
+                    None => builder.append(false),
+                }
+            }
+            builder.finish()
+        }
+
+        // array_filter(a, x#1 -> x#1 = c) over a: array<string> and c: string
+        fn filter(func: &str, array: spark_expression::Expr) -> spark_expression::Expr {
+            expr(ExprStruct::HigherOrderFunc(Box::new(HigherOrderFunc {
+                func: func.to_string(),
+                array: Some(Box::new(array)),
+                param_name: "x#1".to_string(),
+                body: Some(Box::new(expr(ExprStruct::Eq(Box::new(BinaryExpr {
+                    left: Some(Box::new(variable("x#1"))),
+                    right: Some(Box::new(bound(1))),
+                }))))),
+            })))
+        }
+
+        let list = string_lists(vec![
+            Some(vec![Some("a"), None, Some("b"), Some("a")]),
+            Some(vec![]),
+            None,
+            Some(vec![Some("b")]),
+        ]);
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("a", list.data_type().clone(), true),
+            Field::new("c", DataType::Utf8, true),
+        ]));
+        let batch = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![
+                Arc::new(list),
+                Arc::new(StringArray::from(vec!["a", "a", "a", "x"])),
+            ],
+        )?;
+
+        let session_ctx = SessionContext::new();
+        let planner = PhysicalPlanner::new(Arc::from(session_ctx), 0);
+
+        // The lambda keeps the elements equal to the captured column, and drops NULL elements.
+        let physical = planner
+            .create_expr(&filter("array_filter", bound(0)), Arc::clone(&schema))
+            .unwrap();
+        let result = physical.evaluate(&batch)?.into_array(batch.num_rows())?;
+        let expected = string_lists(vec![
+            Some(vec![Some("a"), Some("a")]),
+            Some(vec![]),
+            None,
+            Some(vec![]),
+        ]);
+        assert_eq!(&result, &(Arc::new(expected) as ArrayRef));
+
+        // The argument must be an array, and only array_filter is implemented.
+        let not_array = planner.create_expr(&filter("array_filter", bound(1)), Arc::clone(&schema));
+        assert!(not_array
+            .unwrap_err()
+            .to_string()
+            .contains("expects an array argument"));
+        let unknown = planner.create_expr(&filter("array_zip", bound(0)), Arc::clone(&schema));
+        assert!(unknown
+            .unwrap_err()
+            .to_string()
+            .contains("Unsupported higher-order function"));
+        Ok(())
+    }
+
     #[test]
     fn test_create_array() {
         let session_ctx = SessionContext::new();

native/proto/src/proto/expr.proto

@@ -96,6 +96,8 @@ message Expr {
     Uuid uuid = 74;
     NativeScalarUdf native_scalar_udf = 75;
     AtLeastNNonNulls at_least_n_non_nulls = 76;
+    HigherOrderFunc higher_order_func = 77;
+    LambdaVariable lambda_variable = 78;
   }
 
   reserved 20;
@@ -739,6 +741,27 @@ message ArraysZip {
   repeated string names = 2;
 }
 
+// A DataFusion higher-order function applied to one array argument and a
+// one-parameter lambda, e.g. `array_filter(array, x -> body)`. The lambda body
+// is an ordinary expression tree in which `LambdaVariable` nodes refer to the
+// lambda parameter and `BoundReference` nodes refer to columns of the
+// enclosing operator's input.
+message HigherOrderFunc {
+  // DataFusion higher-order function name, e.g. "array_filter".
+  string func = 1;
+  Expr array = 2;
+  // Unique name of the lambda parameter. `LambdaVariable.name` nodes in the
+  // body resolve against it.
+  string param_name = 3;
+  Expr body = 4;
+}
+
+// Reference to a lambda parameter from inside a `HigherOrderFunc` body. The
+// element type is taken from the array argument at planning time.
+message LambdaVariable {
+  string name = 1;
+}
+
 // Scalar UDF dispatched to the JVM via JNI. Native side exports input arrays
 // through Arrow C Data Interface, calls CometUdfBridge.evaluate, and imports
 // the result.

native/spark-expr/src/conditional_funcs/case_when.rs

@@ -34,7 +34,7 @@ use datafusion::common::{
 use datafusion::logical_expr::type_coercion::binary::type_union_coercion;
 use datafusion::logical_expr::{ColumnarValue, Operator};
 use datafusion::physical_expr::expressions::{
-    BinaryExpr, CaseExpr, Column, IsNotNullExpr, IsNullExpr, Literal, NotExpr,
+    BinaryExpr, CaseExpr, Column, IsNotNullExpr, IsNullExpr, LambdaVariable, Literal, NotExpr,
 };
 use datafusion::physical_expr::PhysicalExpr;
 use std::fmt::{Display, Formatter};
@@ -451,7 +451,7 @@ pub(super) fn cannot_fail(expr: &Arc<dyn PhysicalExpr>, input_schema: &Schema) -
 }
 
 fn infallible(expr: &Arc<dyn PhysicalExpr>, input_schema: &Schema, parsing: bool) -> bool {
-    if expr.is::<Column>() || expr.is::<Literal>() {
+    if expr.is::<Column>() || expr.is::<Literal>() || expr.is::<LambdaVariable>() {
         return true;
     }
     let node_is_infallible = if let Some(binary) = expr.downcast_ref::<BinaryExpr>() {