add lambda substrait 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/mod.rs

@@ -93,8 +93,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,14 @@ 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().udlf(fn_name).or_else(|e| {
+        if let Some(alt_name) = substrait_to_df_name(fn_name) {
+            consumer.get_function_registry().udlf(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 +61,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::LambdaFunction(expr::LambdaFunction::new(
+            func.to_owned(),
+            args,
+        )))
+    } else if let Ok(func) = udf_func {
         Ok(Expr::ScalarFunction(expr::ScalarFunction::new_udf(
             func.to_owned(),
             args,

datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs

@@ -23,21 +23,24 @@ use super::{
     from_substrait_rex, from_window_function,
 };
 use crate::extensions::Extensions;
+use crate::logical_plan::consumer::from_substrait_type_without_names;
 use async_trait::async_trait;
-use datafusion::arrow::datatypes::DataType;
+use datafusion::arrow::datatypes::{DataType, Field, FieldRef};
 use datafusion::catalog::TableProvider;
 use datafusion::common::{
     DFSchema, ScalarValue, TableReference, not_impl_err, substrait_err,
 };
 use datafusion::execution::{FunctionRegistry, SessionState};
 use datafusion::logical_expr::{Expr, Extension, LogicalPlan};
+use datafusion::prelude::{lambda, lambda_var};
+use std::collections::VecDeque;
 use std::sync::{Arc, RwLock};
-use substrait::proto;
 use substrait::proto::expression as substrait_expression;
 use substrait::proto::expression::{
     Enum, FieldReference, IfThen, Literal, MultiOrList, Nested, ScalarFunction,
     SingularOrList, SwitchExpression, WindowFunction,
 };
+use substrait::proto::{self, Type};
 use substrait::proto::{
     AggregateRel, ConsistentPartitionWindowRel, CrossRel, DynamicParameter, ExchangeRel,
     Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel,
@@ -372,6 +375,14 @@ pub trait SubstraitConsumer: Send + Sync + Sized {
         not_impl_err!("Dynamic Parameter expression not supported")
     }
 
+    async fn consume_lambda(
+        &self,
+        expr: &proto::expression::Lambda,
+        input_schema: &DFSchema,
+    ) -> datafusion::common::Result<Expr> {
+        from_lambda(self, expr, input_schema).await
+    }
+
     // Outer Schema Stack
     // These methods manage a stack of outer schemas for correlated subquery support.
     // When entering a subquery, the enclosing query's schema is pushed onto the stack.
@@ -469,6 +480,32 @@ pub trait SubstraitConsumer: Send + Sync + Sized {
         };
         substrait_err!("Missing handler for user-defined literals {}", type_ref)
     }
+
+    fn with_lambda_parameters(
+        &self,
+        lambda_parameters: &[Type],
+    ) -> datafusion::common::Result<(Vec<String>, Self)>;
+
+    fn lambda_variable(
+        &self,
+        steps_out: usize,
+        field_idx: usize,
+    ) -> datafusion::common::Result<Expr>;
+}
+
+async fn from_lambda(
+    consumer: &impl SubstraitConsumer,
+    expr: &proto::expression::Lambda,
+    input_schema: &DFSchema,
+) -> datafusion::common::Result<Expr> {
+    let parameters = expr.parameters.as_ref().unwrap();
+
+    let (names, consumer) = consumer.with_lambda_parameters(&parameters.types)?;
+
+    let body = expr.body.as_ref().unwrap();
+    let body = consumer.consume_expression(body, input_schema).await?;
+
+    Ok(lambda(names, body))
 }
 
 /// Default SubstraitConsumer for converting standard Substrait without user-defined extensions.
@@ -478,6 +515,8 @@ pub struct DefaultSubstraitConsumer<'a> {
     pub(super) extensions: &'a Extensions,
     pub(super) state: &'a SessionState,
     outer_schemas: RwLock<Vec<Arc<DFSchema>>>,
+    lambdas_parameters: VecDeque<Vec<FieldRef>>,
+    num_lambda_parameters: usize,
 }
 
 impl<'a> DefaultSubstraitConsumer<'a> {
@@ -486,12 +525,61 @@ impl<'a> DefaultSubstraitConsumer<'a> {
             extensions,
             state,
             outer_schemas: RwLock::new(Vec::new()),
+            lambdas_parameters: VecDeque::new(),
+            num_lambda_parameters: 0,
         }
     }
 }
 
 #[async_trait]
 impl SubstraitConsumer for DefaultSubstraitConsumer<'_> {
+    fn with_lambda_parameters(
+        &self,
+        lambda_parameters: &[Type],
+    ) -> datafusion::common::Result<(Vec<String>, Self)> {
+        let mut lambdas_parameters = self.lambdas_parameters.clone();
+
+        let lambda_parameters = lambda_parameters
+            .iter()
+            .enumerate()
+            .map(|(i, ty)| {
+                let dt = from_substrait_type_without_names(self, ty)?;
+
+                Ok(Arc::new(Field::new(
+                    format!("p{}", i + self.num_lambda_parameters),
+                    dt,
+                    true,
+                )))
+            })
+            .collect::<datafusion::common::Result<Vec<_>>>()?;
+
+        let names = lambda_parameters.iter().map(|f| f.name().clone()).collect();
+        let num_lambda_parameters = self.num_lambda_parameters + lambda_parameters.len();
+
+        lambdas_parameters.push_front(lambda_parameters);
+
+        Ok((
+            names,
+            Self {
+                extensions: self.extensions,
+                state: self.state,
+                outer_schemas: RwLock::new(self.outer_schemas.read().unwrap().clone()),
+                lambdas_parameters,
+                num_lambda_parameters,
+            },
+        ))
+    }
+
+    fn lambda_variable(
+        &self,
+        steps_out: usize,
+        field_idx: usize,
+    ) -> datafusion::common::Result<Expr> {
+        let var = &self.lambdas_parameters[steps_out][field_idx];
+
+        Ok(lambda_var(var.name(), Arc::clone(var)))
+    }
+
     async fn resolve_table_ref(
         &self,
         table_ref: &TableReference,

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

@@ -150,10 +150,8 @@ pub fn to_substrait_rex(
         }
         Expr::Unnest(expr) => not_impl_err!("Cannot convert {expr:?} to Substrait"),
         Expr::LambdaFunction(expr) => producer.handle_lambda_function(expr, schema),
-        Expr::Lambda(expr) => not_impl_err!("Cannot convert {expr:?} to Substrait"),
-        Expr::LambdaVariable(expr) => {
-            not_impl_err!("Cannot convert {expr:?} to Substrait")
-        }
+        Expr::Lambda(expr) => producer.handle_lambda(expr, schema),
+        Expr::LambdaVariable(expr) => producer.handle_lambda_variable(expr, schema),
     }
 }
 

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

@@ -35,7 +35,37 @@ pub fn from_lambda_function(
     fun: &expr::LambdaFunction,
     schema: &DFSchemaRef,
 ) -> datafusion::common::Result<Expression> {
-    from_function(producer, fun.name(), &fun.args, schema)
+    let lambdas_parameters = fun.lambdas_parameters(schema)?;
+
+    let arguments = std::iter::zip(&fun.args, lambdas_parameters)
+        .map(|(arg, lambda_parameters)| {
+            let arg = match lambda_parameters {
+                Some(lambda_parameters) => {
+                    let mut producer =
+                        producer.with_lambda_parameters(lambda_parameters)?;
+
+                    producer.handle_expr(arg, schema)?
+                }
+                None => producer.handle_expr(arg, schema)?,
+            };
+
+            Ok(FunctionArgument {
+                arg_type: Some(ArgType::Value(arg)),
+            })
+        })
+        .collect::<datafusion::common::Result<_>>()?;
+
+    let function_anchor = producer.register_function(fun.name().to_string());
+    #[expect(deprecated)]
+    Ok(Expression {
+        rex_type: Some(RexType::ScalarFunction(ScalarFunction {
+            function_reference: function_anchor,
+            arguments,
+            output_type: None,
+            options: vec![],
+            args: vec![],
+        })),
+    })
 }
 
 fn from_function(