feat: Lambda function support from DataFusion, illustrated with array_filter

native/core/src/execution/planner.rs

@@ -72,9 +72,9 @@ use datafusion::{
     prelude::SessionContext,
 };
 use datafusion_comet_spark_expr::{
-    create_comet_physical_fun, create_comet_physical_fun_with_eval_mode, BinaryOutputStyle,
-    BloomFilterAgg, BloomFilterMightContain, CsvWriteOptions, EvalMode, SparkArraysZipFunc,
-    SparkBloomFilterVersion, SumInteger, ToCsv,
+    create_comet_hof_func, create_comet_physical_fun, create_comet_physical_fun_with_eval_mode,
+    BinaryOutputStyle, BloomFilterAgg, BloomFilterMightContain, CsvWriteOptions, EvalMode,
+    SparkArraysZipFunc, SparkBloomFilterVersion, SumInteger, ToCsv,
 };
 use datafusion_spark::function::aggregate::collect::SparkCollectSet;
 use iceberg::expr::Bind;
@@ -95,9 +95,9 @@ use datafusion::logical_expr::{
     AggregateUDF, ReturnFieldArgs, ScalarUDF, TypeSignature, WindowFrame, WindowFrameBound,
     WindowFrameUnits, WindowFunctionDefinition,
 };
-use datafusion::physical_expr::expressions::{Literal, StatsType};
+use datafusion::physical_expr::expressions::{LambdaExpr, LambdaVariable, Literal, StatsType};
 use datafusion::physical_expr::window::WindowExpr;
-use datafusion::physical_expr::LexOrdering;
+use datafusion::physical_expr::{HigherOrderFunctionExpr, LexOrdering};
 
 use crate::parquet::parquet_exec::init_datasource_exec;
 use arrow::array::{
@@ -113,7 +113,7 @@ use datafusion::physical_plan::filter::FilterExec;
 use datafusion::physical_plan::joins::NestedLoopJoinExec;
 use datafusion::physical_plan::limit::GlobalLimitExec;
 use datafusion::physical_plan::unnest::{ListUnnest, UnnestExec};
-use datafusion_comet_proto::spark_expression::ListLiteral;
+use datafusion_comet_proto::spark_expression::{HigherOrderFunc, LambdaFunction, ListLiteral};
 use datafusion_comet_proto::spark_operator::SparkFilePartition;
 use datafusion_comet_proto::{
     spark_expression::{
@@ -532,6 +532,19 @@ impl PhysicalPlanner {
                     _ => func,
                 }
             }
+            ExprStruct::HighOrderFunc(hof) => {
+                self.create_high_order_function_expr(hof, input_schema)
+            }
+            ExprStruct::NamedLambdaVariable(nlv) => {
+                let idx = input_schema.index_of(&nlv.name).map_err(|_| {
+                    GeneralError(format!(
+                        "NamedLambdaVariable '{}' not found in enclosing lambda schema",
+                        nlv.name
+                    ))
+                })?;
+                let field = Arc::clone(&input_schema.fields()[idx]);
+                Ok(Arc::new(LambdaVariable::new(idx, field)))
+            }
             ExprStruct::CaseWhen(case_when) => {
                 let when_then_pairs = case_when
                     .when
@@ -3113,6 +3126,89 @@ impl PhysicalPlanner {
         }
     }
 
+    fn create_high_order_function_expr(
+        &self,
+        expr: &HigherOrderFunc,
+        input_schema: SchemaRef,
+    ) -> Result<Arc<dyn PhysicalExpr>, ExecutionError> {
+        let comet_hof_func =
+            create_comet_hof_func(expr.func_name.as_str(), &self.session_ctx.state())?;
+
+        let value_args = expr
+            .value_args
+            .iter()
+            .map(|e| self.create_expr(e, Arc::clone(&input_schema)))
+            .collect::<Result<Vec<_>, _>>()?;
+
+        let lambdas = expr
+            .lambdas
+            .iter()
+            .map(|l| self.create_lambda_expr(l, &input_schema))
+            .collect::<Result<Vec<_>, _>>()?;
+
+        // NOTE: assumes all value arguments precede all lambda arguments.
+        // Holds for array_filter and the current single-lambda HOFs, but would
+        // NOT generalize to a future HOF with interleaved value/lambda args
+        // (e.g. f(value, lambda, value, lambda)). Revisit this split if such a
+        // function is added.
+        let mut args: Vec<Arc<dyn PhysicalExpr>> =
+            Vec::with_capacity(value_args.len() + lambdas.len());
+        args.extend(value_args);
+        args.extend(lambdas);
+
+        let higher_order_function_expr = HigherOrderFunctionExpr::try_new_with_schema(
+            comet_hof_func,
+            args,
+            &input_schema,
+            Arc::new(ConfigOptions::default()),
+        )?;
+
+        Ok(Arc::new(higher_order_function_expr))
+    }
+
+    fn create_lambda_expr(
+        &self,
+        lambda: &LambdaFunction,
+        input_schema: &SchemaRef,
+    ) -> Result<Arc<dyn PhysicalExpr>, ExecutionError> {
+        let mut body_fields: Vec<Arc<Field>> =
+            input_schema.fields().iter().map(Arc::clone).collect();
+
+        for arg in &lambda.args {
+            let data_type = arg.data_type.as_ref().ok_or_else(|| {
+                DataFusionError::Internal("lambda variable without data type".to_string())
+            })?;
+            let arrow_data_type = to_arrow_datatype(data_type);
+            body_fields.push(Arc::new(Field::new(
+                &arg.name,
+                arrow_data_type,
+                arg.nullable,
+            )));
+        }
+
+        let body_schema = Arc::new(Schema::new(
+            body_fields
+                .iter()
+                .map(|f| f.as_ref().clone())
+                .collect::<Vec<Field>>(),
+        ));
+
+        let lambda_body = lambda
+            .body
+            .as_ref()
+            .ok_or_else(|| DataFusionError::Internal("lambda has no body".to_string()))?;
+        let body_expr = self.create_expr(lambda_body, body_schema)?;
+
+        Ok(Arc::new(LambdaExpr::try_new(
+            lambda
+                .args
+                .iter()
+                .map(|a| a.name.clone())
+                .collect::<Vec<_>>(),
+            body_expr,
+        )?))
+    }
+
     fn create_scalar_function_expr(
         &self,
         expr: &ScalarFunc,

native/proto/src/proto/expr.proto

@@ -93,6 +93,8 @@ message Expr {
     JvmScalarUdf jvm_scalar_udf = 70;
     PreciseTimestampConversion precise_timestamp_conversion = 71;
     Shuffle shuffle = 72;
+    HigherOrderFunc high_order_func = 73;
+    NamedLambdaVariable named_lambda_variable = 74;
   }
 
   // Optional QueryContext for error reporting (contains SQL text and position)
@@ -557,3 +559,20 @@ message JvmScalarUdf {
   // Whether the result column may contain nulls.
   bool return_nullable = 4;
 }
+
+message HigherOrderFunc {
+  string func_name = 1;
+  repeated Expr value_args = 2;
+  repeated LambdaFunction lambdas = 3;
+}
+
+message NamedLambdaVariable {
+  string name = 1;
+  DataType data_type = 2;
+  bool nullable = 3;
+}
+
+message LambdaFunction {
+  Expr body = 1;
+  repeated NamedLambdaVariable args = 2;
+}

native/spark-expr/src/comet_high_order_funcs.rs

@@ -0,0 +1,30 @@
+// 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.
+
+use datafusion::common::DataFusionError;
+use datafusion::execution::FunctionRegistry;
+use datafusion::logical_expr::HigherOrderUDF;
+use std::sync::Arc;
+
+pub fn create_comet_hof_func(
+    func_name: &str,
+    registry: &dyn FunctionRegistry,
+) -> Result<Arc<HigherOrderUDF>, DataFusionError> {
+    registry.higher_order_function(func_name).map_err(|e| {
+        DataFusionError::Execution(format!("HOF {func_name} not found in the registry: {e}"))
+    })
+}

native/spark-expr/src/lib.rs

@@ -42,6 +42,7 @@ pub use predicate_funcs::{spark_isnan, RLike};
 
 mod agg_funcs;
 mod array_funcs;
+mod comet_high_order_funcs;
 mod comet_scalar_funcs;
 pub mod hash_funcs;
 
@@ -70,6 +71,7 @@ pub use conditional_funcs::*;
 pub use conversion_funcs::*;
 pub use nondetermenistic_funcs::*;
 
+pub use comet_high_order_funcs::create_comet_hof_func;
 pub use comet_scalar_funcs::{
     create_comet_physical_fun, create_comet_physical_fun_with_eval_mode,
     register_all_comet_functions,

spark/src/main/scala/org/apache/comet/CometConf.scala

@@ -381,6 +381,16 @@ object CometConf extends ShimCometConf {
       .booleanConf
       .createWithDefault(true)
 
+  val COMET_EXEC_HIGHER_ORDER_FUNCTION_NATIVE_ENABLED: ConfigEntry[Boolean] =
+    conf("spark.comet.exec.higherOrderFunction.native.enabled")
+      .category(CATEGORY_EXEC)
+      .doc(
+        "When enabled, supported higher-order functions (e.g. filter) are executed by the " +
+          "native DataFusion engine. Shapes the native path cannot handle fall back to the " +
+          "codegen dispatcher, and finally to Spark.")
+      .booleanConf
+      .createWithDefault(true)
+
   val COMET_EXEC_SHUFFLE_WITH_HASH_PARTITIONING_ENABLED: ConfigEntry[Boolean] =
     conf("spark.comet.native.shuffle.partitioning.hash.enabled")
       .category(CATEGORY_SHUFFLE)