feat: add substrait lambda support

Author: gstvg

datafusion/substrait/src/extensions.rs

@@ -27,7 +27,7 @@ use substrait::proto::extensions::simple_extension_declaration::{
 /// types. This structs facilitates the use of these extensions in DataFusion.
 /// TODO: DF doesn't yet use extensions for type variations <https://github.com/apache/datafusion/issues/11544>
 /// TODO: DF doesn't yet provide valid extensionUris <https://github.com/apache/datafusion/issues/11545>
-#[derive(Default, Debug, PartialEq)]
+#[derive(Clone, Default, Debug, PartialEq)]
 pub struct Extensions {
     pub functions: HashMap<u32, String>, // anchor -> function name
     pub types: HashMap<u32, String>,     // anchor -> type name

datafusion/substrait/src/logical_plan/consumer/expr/field_reference.rs

@@ -21,7 +21,7 @@ use datafusion::logical_expr::Expr;
 use std::sync::Arc;
 use substrait::proto::expression::FieldReference;
 use substrait::proto::expression::field_reference::ReferenceType::DirectReference;
-use substrait::proto::expression::field_reference::RootType;
+use substrait::proto::expression::field_reference::{LambdaParameterReference, RootType};
 use substrait::proto::expression::reference_segment::ReferenceType::StructField;
 
 pub async fn from_field_reference(
@@ -56,9 +56,9 @@ pub(crate) fn from_substrait_field_reference(
                     Some(RootType::Expression(_)) => not_impl_err!(
                         "Expression root type in field reference is not supported"
                     ),
-                    Some(RootType::LambdaParameterReference(_)) => not_impl_err!(
-                        "Lambda parameter reference in field reference is not yet supported"
-                    ),
+                    Some(RootType::LambdaParameterReference(
+                        LambdaParameterReference { steps_out },
+                    )) => consumer.lambda_variable(*steps_out as usize, field_idx),
                 }
             }
             _ => not_impl_err!(

datafusion/substrait/src/logical_plan/consumer/expr/lambda.rs

@@ -0,0 +1,47 @@
+// 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::{DFSchema, substrait_err},
+    prelude::{Expr, lambda},
+};
+use substrait::proto;
+
+use crate::logical_plan::consumer::SubstraitConsumer;
+
+pub async fn from_lambda(
+    consumer: &impl SubstraitConsumer,
+    expr: &proto::expression::Lambda,
+    input_schema: &DFSchema,
+) -> datafusion::common::Result<Expr> {
+    let Some(parameters) = expr.parameters.as_ref() else {
+        return substrait_err!("Lambda expression without parameters is not allowed");
+    };
+
+    let (names, consumer_with_parameters) =
+        consumer.with_lambda_parameters(&parameters.types, input_schema)?;
+
+    let Some(body) = expr.body.as_ref() else {
+        return substrait_err!("Lambda expression without body is not allowed");
+    };
+
+    let body = consumer_with_parameters
+        .consume_expression(body, input_schema)
+        .await?;
+
+    Ok(lambda(names, body))
+}

datafusion/substrait/src/logical_plan/consumer/expr/mod.rs

@@ -20,6 +20,7 @@ mod cast;
 mod field_reference;
 mod function_arguments;
 mod if_then;
+mod lambda;
 mod literal;
 mod nested;
 mod scalar_function;
@@ -32,6 +33,7 @@ pub use cast::*;
 pub use field_reference::*;
 pub use function_arguments::*;
 pub use if_then::*;
+pub use lambda::*;
 pub use literal::*;
 pub use nested::*;
 pub use scalar_function::*;
@@ -95,8 +97,11 @@ pub async fn from_substrait_rex(
             RexType::DynamicParameter(expr) => {
                 consumer.consume_dynamic_parameter(expr, input_schema).await
             }
-            RexType::Lambda(_) | RexType::LambdaInvocation(_) => {
-                not_impl_err!("Lambda expressions are not yet supported")
+            RexType::Lambda(lambda) => {
+                consumer.consume_lambda(lambda.as_ref(), input_schema).await
+            }
+            RexType::LambdaInvocation(_) => {
+                not_impl_err!("Lambda invocations are not supported")
             }
         },
         None => substrait_err!("Expression must set rex_type: {expression:?}"),

datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs

@@ -45,6 +45,20 @@ pub async fn from_scalar_function(
     let fn_name = substrait_fun_name(fn_signature);
     let args = from_substrait_func_args(consumer, &f.arguments, input_schema).await?;
 
+    let udlf_func = consumer
+        .get_function_registry()
+        .higher_order_function(fn_name)
+        .or_else(|e| {
+            if let Some(alt_name) = substrait_to_df_name(fn_name) {
+                consumer
+                    .get_function_registry()
+                    .higher_order_function(alt_name)
+                    .or(Err(e))
+            } else {
+                Err(e)
+            }
+        });
+
     let udf_func = consumer.get_function_registry().udf(fn_name).or_else(|e| {
         if let Some(alt_name) = substrait_to_df_name(fn_name) {
             consumer.get_function_registry().udf(alt_name).or(Err(e))
@@ -53,9 +67,14 @@ pub async fn from_scalar_function(
         }
     });
 
-    // try to first match the requested function into registered udfs, then built-in ops
+    // try to first match the requested function into registered udlfs, then udfs, built-in ops
     // and finally built-in expressions
-    if let Ok(func) = udf_func {
+    if let Ok(func) = udlf_func {
+        Ok(Expr::HigherOrderFunction(expr::HigherOrderFunction::new(
+            func.to_owned(),
+            args,
+        )))
+    } else if let Ok(func) = udf_func {
         Ok(Expr::ScalarFunction(expr::ScalarFunction::new_udf(
             func.to_owned(),
             args,