finish lambda substrait support

Author: gstvg

datafusion/core/src/execution/context/mod.rs

@@ -1566,6 +1566,20 @@ impl SessionContext {
         state.register_udf(Arc::new(f)).ok();
     }
 
+    /// Registers a higher-order function within this context.
+    ///
+    /// Note in SQL queries, function names are looked up using
+    /// lowercase unless the query uses quotes. For example,
+    ///
+    /// - `SELECT MY_HIGHER_ORDER_FUNC(x)...` will look for a function named `"my_higher_order_func"`
+    /// - `SELECT "my_HIGHER_ORDER_FUNC"(x)` will look for a function named `"my_HIGHER_ORDER_FUNC"`
+    ///
+    /// Any functions registered with the function name or its aliases will be overwritten with this new function
+    pub fn register_higher_order_function(&self, f: Arc<dyn HigherOrderUDF>) {
+        let mut state = self.state.write();
+        state.register_higher_order_function(f).ok();
+    }
+
     /// Registers an aggregate UDF within this context.
     ///
     /// Note in SQL queries, aggregate names are looked up using

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::*;

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

@@ -44,13 +44,19 @@ 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 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) {
@@ -63,7 +69,7 @@ pub async fn from_scalar_function(
     // try to first match the requested function into registered udlfs, then udfs, built-in ops
     // and finally built-in expressions
     if let Ok(func) = udlf_func {
-        Ok(Expr::LambdaFunction(expr::LambdaFunction::new(
+        Ok(Expr::HigherOrderFunction(expr::HigherOrderFunction::new(
             func.to_owned(),
             args,
         )))