[RFC] Add lambda support and array_transform udf

datafusion-examples/examples/sql_ops/frontend.rs

@@ -22,8 +22,8 @@ use datafusion::common::{TableReference, plan_err};
 use datafusion::config::ConfigOptions;
 use datafusion::error::Result;
 use datafusion::logical_expr::{
-    AggregateUDF, Expr, LogicalPlan, ScalarUDF, TableProviderFilterPushDown, TableSource,
-    WindowUDF,
+    AggregateUDF, Expr, LambdaUDF, LogicalPlan, ScalarUDF, TableProviderFilterPushDown,
+    TableSource, WindowUDF,
 };
 use datafusion::optimizer::{
     Analyzer, AnalyzerRule, Optimizer, OptimizerConfig, OptimizerContext, OptimizerRule,
@@ -155,6 +155,10 @@ impl ContextProvider for MyContextProvider {
         None
     }
 
+    fn get_lambda_meta(&self, _name: &str) -> Option<Arc<dyn LambdaUDF>> {
+        None
+    }
+
     fn get_aggregate_meta(&self, _name: &str) -> Option<Arc<AggregateUDF>> {
         None
     }
@@ -175,6 +179,10 @@ impl ContextProvider for MyContextProvider {
         Vec::new()
     }
 
+    fn udlf_names(&self) -> Vec<String> {
+        Vec::new()
+    }
+
     fn udaf_names(&self) -> Vec<String> {
         Vec::new()
     }

datafusion/catalog-listing/src/helpers.rs

@@ -85,7 +85,9 @@ pub fn expr_applicable_for_cols(col_names: &[&str], expr: &Expr) -> bool {
         | Expr::ScalarSubquery(_)
         | Expr::SetComparison(_)
         | Expr::GroupingSet(_)
-        | Expr::Case(_) => Ok(TreeNodeRecursion::Continue),
+        | Expr::Case(_)
+        | Expr::Lambda(_)
+        | Expr::LambdaVariable(_) => Ok(TreeNodeRecursion::Continue),
 
         Expr::ScalarFunction(scalar_function) => {
             match scalar_function.func.signature().volatility {
@@ -97,6 +99,16 @@ pub fn expr_applicable_for_cols(col_names: &[&str], expr: &Expr) -> bool {
                 }
             }
         }
+        Expr::LambdaFunction(lambda_function) => {
+            match lambda_function.func.signature().volatility {
+                Volatility::Immutable => Ok(TreeNodeRecursion::Continue),
+                // TODO: Stable functions could be `applicable`, but that would require access to the context
+                Volatility::Stable | Volatility::Volatile => {
+                    is_applicable = false;
+                    Ok(TreeNodeRecursion::Stop)
+                }
+            }
+        }
 
         // TODO other expressions are not handled yet:
         // - AGGREGATE and WINDOW should not end up in filter conditions, except maybe in some edge cases

datafusion/common/src/utils/mod.rs

@@ -24,7 +24,7 @@ pub mod proxy;
 pub mod string_utils;
 
 use crate::assert_or_internal_err;
-use crate::error::{_exec_datafusion_err, _internal_datafusion_err};
+use crate::error::{_exec_datafusion_err, _exec_err, _internal_datafusion_err};
 use crate::{Result, ScalarValue};
 use arrow::array::{
     Array, ArrayRef, FixedSizeListArray, LargeListArray, ListArray, OffsetSizeTrait,
@@ -971,11 +971,27 @@ pub fn take_function_args<const N: usize, T>(
     })
 }
 
+/// Returns the inner values of a list, or an error otherwise
+/// For [`ListArray`] and [`LargeListArray`], if it's sliced, it returns a
+/// sliced array too. Therefore, too reconstruct a list using it,
+/// you must adjust the offsets using [`adjust_offsets_for_slice`]
+pub fn list_values(array: &dyn Array) -> Result<ArrayRef> {
+    match array.data_type() {
+        DataType::List(_) => Ok(Arc::clone(array.as_list::<i32>().values())),
+        DataType::LargeList(_) => Ok(Arc::clone(array.as_list::<i64>().values())),
+        DataType::FixedSizeList(_, _) => {
+            Ok(Arc::clone(array.as_fixed_size_list().values()))
+        }
+        other => _exec_err!("expected list, got {other}"),
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use crate::ScalarValue::Null;
     use arrow::array::Float64Array;
+    use sqlparser::ast::Ident;
 
     #[test]
     fn test_bisect_linear_left_and_right() -> Result<()> {

datafusion/core/src/bin/print_functions_docs.rs

@@ -18,7 +18,7 @@
 use datafusion::execution::SessionStateDefaults;
 use datafusion_common::{HashSet, Result, not_impl_err};
 use datafusion_expr::{
-    AggregateUDF, DocSection, Documentation, ScalarUDF, WindowUDF,
+    AggregateUDF, DocSection, Documentation, LambdaUDF, ScalarUDF, WindowUDF,
     aggregate_doc_sections, scalar_doc_sections, window_doc_sections,
 };
 use itertools::Itertools;
@@ -282,6 +282,18 @@ impl DocProvider for WindowUDF {
     }
 }
 
+impl DocProvider for dyn LambdaUDF {
+    fn get_name(&self) -> String {
+        self.name().to_string()
+    }
+    fn get_aliases(&self) -> Vec<String> {
+        self.aliases().iter().map(|a| a.to_string()).collect()
+    }
+    fn get_documentation(&self) -> Option<&Documentation> {
+        self.documentation()
+    }
+}
+
 #[expect(clippy::borrowed_box)]
 fn get_names_and_aliases(functions: &Vec<&Box<dyn DocProvider>>) -> Vec<String> {
     functions

datafusion/core/src/datasource/listing_table_factory.rs

@@ -599,6 +599,11 @@ mod tests {
             ) -> &HashMap<String, Arc<datafusion_expr::ScalarUDF>> {
                 unimplemented!()
             }
+            fn lambda_functions(
+                &self,
+            ) -> &HashMap<String, Arc<dyn datafusion_expr::LambdaUDF>> {
+                unimplemented!()
+            }
             fn aggregate_functions(
                 &self,
             ) -> &HashMap<String, Arc<datafusion_expr::AggregateUDF>> {

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

@@ -83,6 +83,7 @@ use datafusion_execution::disk_manager::{
     DEFAULT_MAX_TEMP_DIRECTORY_SIZE, DiskManagerBuilder,
 };
 use datafusion_execution::registry::SerializerRegistry;
+use datafusion_expr::LambdaUDF;
 pub use datafusion_expr::execution_props::ExecutionProps;
 #[cfg(feature = "sql")]
 use datafusion_expr::planner::RelationPlanner;
@@ -1976,6 +1977,10 @@ impl FunctionRegistry for SessionContext {
         self.state.read().udf(name)
     }
 
+    fn udlf(&self, name: &str) -> Result<Arc<dyn LambdaUDF>> {
+        self.state.read().udlf(name)
+    }
+
     fn udaf(&self, name: &str) -> Result<Arc<AggregateUDF>> {
         self.state.read().udaf(name)
     }
@@ -1988,6 +1993,13 @@ impl FunctionRegistry for SessionContext {
         self.state.write().register_udf(udf)
     }
 
+    fn register_udlf(
+        &mut self,
+        udlf: Arc<dyn LambdaUDF>,
+    ) -> Result<Option<Arc<dyn LambdaUDF>>> {
+        self.state.write().register_udlf(udlf)
+    }
+
     fn register_udaf(
         &mut self,
         udaf: Arc<AggregateUDF>,
@@ -2017,6 +2029,10 @@ impl FunctionRegistry for SessionContext {
         self.state.write().register_expr_planner(expr_planner)
     }
 
+    fn udlfs(&self) -> HashSet<String> {
+        self.state.read().udlfs()
+    }
+
     fn udafs(&self) -> HashSet<String> {
         self.state.read().udafs()
     }