feat: Lambda function support from DataFusion, illustrated with array_filter

native/core/src/execution/lambda.rs

@@ -0,0 +1,69 @@
+// 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.
+
+//! Helpers for planning DataFusion higher-order functions (HOFs) coming
+//! from Spark.
+//!
+//! The planner needs three things that don't belong in `planner.rs`:
+//! 1. A stack of *lambda scopes* so nested `NamedLambdaVariable`s resolve
+//!    by Spark `exprId` (immune to name shadowing / column collisions).
+//! 2. A drop-guard that pops a scope on any exit path (`?`, panic-safe).
+//! 3. A tiny `PhysicalExpr` wrapper that keeps *unused* lambda parameters
+//!    visible in `children()` so `LambdaExpr::new`'s projection compaction
+//!    stays consistent with the runtime batch layout.
+
+use std::cell::RefCell;
+use std::collections::HashMap;
+
+use arrow::datatypes::FieldRef;
+use datafusion::common::Result;
+
+/// Maps Spark `exprId` -> (column index in the extended body schema, field).
+pub(crate) type LambdaScope = HashMap<i64, (usize, FieldRef)>;
+
+/// A stack of lambda variable scopes, innermost last.
+/// Planning is single-threaded per planner, so `RefCell` is sufficient to manage
+/// the stack of scopes during the recursive planning process.
+#[derive(Default)]
+pub(crate) struct LambdaScopes {
+    stack: RefCell<Vec<LambdaScope>>,
+}
+
+impl LambdaScopes {
+    /// Resolve a lambda variable by Spark `exprId`, searching innermost
+    /// scope first.
+    pub(crate) fn resolve_variable(&self, expr_id: i64) -> Option<(usize, FieldRef)> {
+        self.stack
+            .borrow()
+            .iter()
+            .rev()
+            .find_map(|s| s.get(&expr_id).cloned())
+    }
+
+    /// Push `scope`, run `f`, pop unconditionally. The pop happens on both
+    /// the `Ok` and `Err` paths — this replaces the earlier RAII guard.
+    pub(crate) fn with_scope<T, E>(
+        &self,
+        scope: LambdaScope,
+        f: impl FnOnce() -> Result<T, E>,
+    ) -> Result<T, E> {
+        self.stack.borrow_mut().push(scope);
+        let out = f();
+        self.stack.borrow_mut().pop();
+        out
+    }
+}

native/core/src/execution/mod.rs

@@ -25,6 +25,7 @@ pub mod operators;
 pub(crate) mod planner;
 pub mod serde;
 pub use datafusion_comet_shuffle as shuffle;
+pub mod lambda;
 mod memory_pools;
 pub(crate) mod sort;
 pub(crate) mod spark_config;

native/core/src/execution/planner.rs

@@ -91,9 +91,10 @@ use datafusion::{
     prelude::SessionContext,
 };
 use datafusion_comet_spark_expr::{
-    create_comet_physical_fun, create_comet_physical_fun_with_eval_mode, BinaryOutputStyle,
-    BloomFilterAgg, BloomFilterMightContain, CometCollectList, CometCollectSet, CsvWriteOptions,
-    EvalMode, SparkArraysZipFunc, SparkBloomFilterVersion, SparkPercentile, SumInteger, ToCsv,
+    create_comet_hof_func, create_comet_physical_fun, create_comet_physical_fun_with_eval_mode,
+    BinaryOutputStyle, BloomFilterAgg, BloomFilterMightContain, CometCollectList, CometCollectSet,
+    CsvWriteOptions, EvalMode, SparkArraysZipFunc, SparkBloomFilterVersion, SparkPercentile,
+    SumInteger, ToCsv,
 };
 use iceberg::expr::Bind;
 
@@ -111,13 +112,15 @@ use datafusion::datasource::listing::PartitionedFile;
 use datafusion::logical_expr::type_coercion::functions::fields_with_udf;
 use datafusion::logical_expr::type_coercion::other::get_coerce_type_for_case_expression;
 use datafusion::logical_expr::{
-    AggregateUDF, ReturnFieldArgs, ScalarUDF, TypeSignature, WindowFrame, WindowFrameBound,
-    WindowFrameUnits, WindowFunctionDefinition,
+    AggregateUDF, HigherOrderUDF, LambdaParametersProgress, ReturnFieldArgs, ScalarUDF,
+    TypeSignature, ValueOrLambda, 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::execution::lambda::{LambdaScope, LambdaScopes};
 use crate::parquet::parquet_exec::init_datasource_exec;
 use arrow::array::{
     new_empty_array, Array, ArrayRef, BinaryBuilder, BooleanArray, Date32Array, Decimal128Array,
@@ -132,7 +135,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;
-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::{
@@ -307,6 +310,9 @@ pub struct PhysicalPlanner {
     /// Task-owned destination for remote shuffle blocks, registered on the driving Spark task
     /// thread before native planning. Only explicit RSS destinations may use it.
     shuffle_partition_pusher: Option<Arc<dyn ShufflePartitionPusher>>,
+    /// Stack of lambda variable scopes, innermost last. Planning is
+    /// single-threaded per planner, so `RefCell` is sufficient.
+    lambda_scopes: LambdaScopes,
 }
 
 impl Default for PhysicalPlanner {
@@ -326,6 +332,7 @@ impl PhysicalPlanner {
             task_context: None,
             class_loader: None,
             shuffle_partition_pusher: None,
+            lambda_scopes: LambdaScopes::default(),
         }
     }
 
@@ -725,6 +732,18 @@ impl PhysicalPlanner {
                     _ => func,
                 }
             }
+            ExprStruct::HighOrderFunc(hof) => {
+                self.create_high_order_function_expr(hof, input_schema)
+            }
+            ExprStruct::NamedLambdaVariable(nlv) => {
+                let (idx, field) = self.lambda_scopes.resolve_variable(nlv.expr_id).ok_or_else(|| {
+                    GeneralError(format!(
+                        "Lambda variable '{}' (exprId={}) is not bound in any enclosing lambda scope",
+                        nlv.name, nlv.expr_id
+                    ))
+                })?;
+                Ok(Arc::new(LambdaVariable::new(idx, field)))
+            }
             ExprStruct::CaseWhen(case_when) => {
                 let when_then_pairs = case_when
                     .when
@@ -3670,6 +3689,134 @@ impl PhysicalPlanner {
         }
     }
 
+    fn create_high_order_function_expr(
+        &self,
+        expr: &HigherOrderFunc,
+        input_schema: SchemaRef,
+    ) -> Result<Arc<dyn PhysicalExpr>, ExecutionError> {
+        let udf = create_comet_hof_func(&expr.func_name, &self.session_ctx.state())?;
+
+        // 1. Plan value args.
+        let value_args: Vec<Arc<dyn PhysicalExpr>> = expr
+            .value_args
+            .iter()
+            .map(|e| self.create_expr(e, Arc::clone(&input_schema)))
+            .collect::<Result<_, _>>()?;
+
+        // 2. Resolve lambda param field types via the UDF (mirrors runtime).
+        let param_fields = Self::resolve_lambda_param_fields(
+            &udf,
+            &expr.func_name,
+            &value_args,
+            expr.lambdas.len(),
+            input_schema.as_ref(),
+        )?;
+
+        // 3. Plan lambdas with resolved param fields.
+        let lambdas: Vec<Arc<dyn PhysicalExpr>> = expr
+            .lambdas
+            .iter()
+            .zip(&param_fields)
+            .map(|(l, fields)| self.create_lambda_expr(l, &input_schema, fields))
+            .collect::<Result<_, _>>()?;
+
+        // 4. NOTE: assumes value args precede lambdas (holds for array_filter).
+        let mut args = value_args;
+        args.extend(lambdas);
+
+        Ok(Arc::new(HigherOrderFunctionExpr::try_new_with_schema(
+            udf,
+            args,
+            &input_schema,
+            Arc::new(ConfigOptions::default()),
+        )?))
+    }
+
+    fn resolve_lambda_param_fields(
+        udf: &HigherOrderUDF,
+        func_name: &str,
+        value_args: &[Arc<dyn PhysicalExpr>],
+        lambda_count: usize,
+        schema: &Schema,
+    ) -> Result<Vec<Vec<FieldRef>>, ExecutionError> {
+        let mut planning_fields: Vec<ValueOrLambda<FieldRef, Option<FieldRef>>> = value_args
+            .iter()
+            .map(|e| Ok(ValueOrLambda::Value(e.return_field(schema)?)))
+            .collect::<Result<_, DataFusionError>>()?;
+        planning_fields.extend(std::iter::repeat_n(
+            ValueOrLambda::Lambda(None),
+            lambda_count,
+        ));
+
+        match udf.lambda_parameters(0, &planning_fields)? {
+            LambdaParametersProgress::Complete(items) if items.len() >= lambda_count => Ok(items),
+            LambdaParametersProgress::Complete(items) => Err(GeneralError(format!(
+                "{func_name}: expected parameter fields for {lambda_count} lambdas, got {}",
+                items.len()
+            ))),
+            LambdaParametersProgress::Partial(_) => Err(GeneralError(format!(
+                "{func_name}: multi-step lambda resolution is not supported yet"
+            ))),
+        }
+    }
+
+    fn create_lambda_expr(
+        &self,
+        lambda: &LambdaFunction,
+        input_schema: &SchemaRef,
+        param_fields: &[FieldRef],
+    ) -> Result<Arc<dyn PhysicalExpr>, ExecutionError> {
+        if param_fields.len() < lambda.args.len() {
+            return Err(GeneralError(format!(
+                "lambda declares {} params but the function resolved only {}",
+                lambda.args.len(),
+                param_fields.len()
+            )));
+        }
+
+        // Build extended schema = input schema ++ lambda params, and the scope.
+        let mut body_fields: Vec<FieldRef> = input_schema.fields().iter().map(Arc::clone).collect();
+        let mut scope = LambdaScope::with_capacity(lambda.args.len());
+        let mut scope_entries: Vec<(usize, FieldRef)> = Vec::with_capacity(lambda.args.len());
+        let mut param_names: Vec<String> = Vec::with_capacity(lambda.args.len());
+
+        for (arg, resolved) in lambda.args.iter().zip(param_fields) {
+            // Runtime uses `param.renamed(name)` — do the same here.
+            let field: FieldRef = Arc::new(resolved.as_ref().clone().with_name(&arg.name));
+            let idx = body_fields.len();
+            if scope
+                .insert(arg.expr_id, (idx, Arc::clone(&field)))
+                .is_some()
+            {
+                return Err(GeneralError(format!(
+                    "duplicate lambda variable exprId {} ('{}')",
+                    arg.expr_id, arg.name
+                )));
+            }
+            scope_entries.push((idx, Arc::clone(&field)));
+            param_names.push(arg.name.clone());
+            body_fields.push(field);
+        }
+
+        let body_schema = Arc::new(Schema::new(
+            body_fields
+                .iter()
+                .map(|f| f.as_ref().clone())
+                .collect::<Vec<_>>(),
+        ));
+        let lambda_body = lambda
+            .body
+            .as_ref()
+            .ok_or_else(|| GeneralError("lambda has no body".to_string()))?;
+
+        // Plan the body under this scope; the guard pops on any `?` / drop.
+        let body_expr = self
+            .lambda_scopes
+            .with_scope(scope, || self.create_expr(lambda_body, body_schema))?;
+
+        Ok(Arc::new(LambdaExpr::try_new(param_names, body_expr)?))
+    }
+
     fn create_scalar_function_expr(
         &self,
         expr: &ScalarFunc,

native/proto/src/proto/expr.proto

@@ -94,6 +94,8 @@ message Expr {
     Shuffle shuffle = 72;
     RandStr rand_str = 73;
     Uuid uuid = 74;
+    HigherOrderFunc high_order_func = 75;
+    NamedLambdaVariable named_lambda_variable = 76;
   }
 
   reserved 20;
@@ -675,3 +677,21 @@ 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;
+  int64 expr_id = 4;
+}
+
+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

@@ -41,6 +41,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;
 
@@ -71,6 +72,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

@@ -409,6 +409,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_SHUFFLE_NATIVE_HASH_PARTITIONING_ENABLED: ConfigEntry[Boolean] =
     conf("spark.comet.shuffle.native.partitioning.hash.enabled")
       .withAlternative("spark.comet.native.shuffle.partitioning.hash.enabled")