Add lambda support and array_transform udf

datafusion/catalog-listing/src/helpers.rs

@@ -86,7 +86,9 @@ pub fn expr_applicable_for_cols(col_names: &[&str], expr: &Expr) -> bool {
         | Expr::InSubquery(_)
         | Expr::ScalarSubquery(_)
         | Expr::GroupingSet(_)
-        | Expr::Case(_) => Ok(TreeNodeRecursion::Continue),
+        | Expr::Case(_)
+        | Expr::Lambda(_)
+        | Expr::LambdaColumn(_) => Ok(TreeNodeRecursion::Continue),
 
         Expr::ScalarFunction(scalar_function) => {
             match scalar_function.func.signature().volatility {

datafusion/common/src/utils/mod.rs

@@ -22,20 +22,26 @@ pub mod memory;
 pub mod proxy;
 pub mod string_utils;
 
-use crate::error::{_exec_datafusion_err, _internal_datafusion_err, _internal_err};
+use crate::error::{
+    _exec_datafusion_err, _exec_err, _internal_datafusion_err, _internal_err,
+};
 use crate::{Result, ScalarValue};
 use arrow::array::{
     cast::AsArray, Array, ArrayRef, FixedSizeListArray, LargeListArray, ListArray,
     OffsetSizeTrait,
 };
+use arrow::array::{ArrowPrimitiveType, PrimitiveArray};
 use arrow::buffer::OffsetBuffer;
 use arrow::compute::{partition, SortColumn, SortOptions};
-use arrow::datatypes::{DataType, Field, SchemaRef};
+use arrow::datatypes::{
+    ArrowNativeType, DataType, Field, Int32Type, Int64Type, SchemaRef,
+};
 #[cfg(feature = "sql")]
 use sqlparser::{ast::Ident, dialect::GenericDialect, parser::Parser};
 use std::borrow::{Borrow, Cow};
 use std::cmp::{min, Ordering};
 use std::collections::HashSet;
+use std::iter::repeat_n;
 use std::num::NonZero;
 use std::ops::Range;
 use std::sync::Arc;
@@ -939,6 +945,121 @@ pub fn take_function_args<const N: usize, T>(
     })
 }
 
+/// [0, 2, 2, 5, 6] -> [0, 0, 2, 2, 2, 3]
+pub fn make_list_array_indices<T: ArrowPrimitiveType>(
+    offsets: &OffsetBuffer<T::Native>,
+) -> PrimitiveArray<T> {
+    let mut indices = Vec::with_capacity(
+        offsets.last().unwrap().as_usize() - offsets.first().unwrap().as_usize(),
+    );
+
+    for (i, (&start, &end)) in std::iter::zip(&offsets[..], &offsets[1..]).enumerate() {
+        indices.extend(repeat_n(
+            T::Native::usize_as(i),
+            end.as_usize() - start.as_usize(),
+        ));
+    }
+
+    PrimitiveArray::new(indices.into(), None)
+}
+
+/// [0, 2, 2, 5, 6] -> [0, 1, 0, 1, 2, 0]
+pub fn make_list_element_indices<T: ArrowPrimitiveType>(
+    offsets: &OffsetBuffer<T::Native>,
+) -> PrimitiveArray<T> {
+    let mut indices =
+        Vec::with_capacity(offsets.last().unwrap().as_usize() - offsets[0].as_usize());
+
+    for (&start, &end) in std::iter::zip(&offsets[..], &offsets[1..]) {
+        indices.extend(
+            (0..end.as_usize() - start.as_usize()).map(|i| T::Native::usize_as(i)),
+        );
+    }
+
+    PrimitiveArray::new(indices.into(), None)
+}
+
+/// (3, 2) -> [0, 0, 1, 1, 2, 2]
+pub fn make_fsl_array_indices(
+    list_size: i32,
+    array_len: usize,
+) -> PrimitiveArray<Int32Type> {
+    let mut indices = Vec::with_capacity(list_size as usize * array_len);
+
+    for i in 0..array_len {
+        indices.extend(repeat_n(i as i32, list_size as usize));
+    }
+
+    PrimitiveArray::new(indices.into(), None)
+}
+
+/// (3, 2) -> [0, 1, 0, 1, 0, 1]
+pub fn make_fsl_element_indices(
+    list_size: i32,
+    array_len: usize,
+) -> PrimitiveArray<Int32Type> {
+    let mut indices = Vec::with_capacity(list_size as usize * array_len);
+
+    if array_len > 0 {
+        indices.extend((0..list_size as usize).map(|j| j as i32));
+
+        for _ in 1..array_len {
+            indices.extend_from_within(0..list_size as usize);
+        }
+    }
+
+    PrimitiveArray::new(indices.into(), None)
+}
+
+pub fn list_values(array: &dyn Array) -> Result<&ArrayRef> {
+    match array.data_type() {
+        DataType::List(_) => Ok(array.as_list::<i32>().values()),
+        DataType::LargeList(_) => Ok(array.as_list::<i64>().values()),
+        DataType::FixedSizeList(_, _) => Ok(array.as_fixed_size_list().values()),
+        other => _exec_err!("expected list, got {other}"),
+    }
+}
+
+pub fn list_indices(array: &dyn Array) -> Result<ArrayRef> {
+    match array.data_type() {
+        DataType::List(_) => Ok(Arc::new(make_list_array_indices::<Int32Type>(
+            array.as_list().offsets(),
+        ))),
+        DataType::LargeList(_) => Ok(Arc::new(make_list_array_indices::<Int64Type>(
+            array.as_list().offsets(),
+        ))),
+        DataType::FixedSizeList(_, _) => {
+            let fixed_size_list = array.as_fixed_size_list();
+
+            Ok(Arc::new(make_fsl_array_indices(
+                fixed_size_list.value_length(),
+                fixed_size_list.len(),
+            )))
+        }
+        other => _exec_err!("expected list, got {other}"),
+    }
+}
+
+pub fn elements_indices(array: &dyn Array) -> Result<ArrayRef> {
+    match array.data_type() {
+        DataType::List(_) => Ok(Arc::new(make_list_element_indices::<Int32Type>(
+            array.as_list::<i32>().offsets(),
+        ))),
+        DataType::LargeList(_) => Ok(Arc::new(make_list_element_indices::<Int64Type>(
+            array.as_list::<i64>().offsets(),
+        ))),
+        DataType::FixedSizeList(_, _) => {
+            let fixed_size_list = array.as_fixed_size_list();
+
+            Ok(Arc::new(make_fsl_element_indices(
+                fixed_size_list.value_length(),
+                fixed_size_list.len(),
+            )))
+        }
+        other => _exec_err!("expected list, got {other}"),
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

datafusion/expr/src/expr.rs

@@ -398,6 +398,30 @@ pub enum Expr {
     OuterReferenceColumn(FieldRef, Column),
     /// Unnest expression
     Unnest(Unnest),
+    /// Lambda expression
+    Lambda(Lambda),
+    LambdaColumn(LambdaColumn),
+}
+
+#[derive(Clone, PartialEq, PartialOrd, Eq, Debug, Hash)]
+pub struct LambdaColumn {
+    pub name: String,
+    pub field: FieldRef,
+    pub spans: Spans,
+}
+
+impl LambdaColumn {
+    pub fn new(name: String, field: FieldRef) -> Self {
+        Self {
+            name,
+            field,
+            spans: Spans::new(),
+        }
+    }
+
+    pub fn spans_mut(&mut self) -> &mut Spans {
+        &mut self.spans
+    }
 }
 
 impl Default for Expr {
@@ -1211,6 +1235,23 @@ impl GroupingSet {
     }
 }
 
+/// Lambda expression.
+#[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)]
+pub struct Lambda {
+    pub params: Vec<String>,
+    pub body: Box<Expr>,
+}
+
+impl Lambda {
+    /// Create a new lambda expression
+    pub fn new(params: Vec<String>, body: Expr) -> Self {
+        Self {
+            params,
+            body: Box::new(body),
+        }
+    }
+}
+
 #[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)]
 #[cfg(not(feature = "sql"))]
 pub struct IlikeSelectItem {
@@ -1525,6 +1566,8 @@ impl Expr {
             #[expect(deprecated)]
             Expr::Wildcard { .. } => "Wildcard",
             Expr::Unnest { .. } => "Unnest",
+            Expr::Lambda { .. } => "Lambda",
+            Expr::LambdaColumn { .. } => "LambdaColumn",
         }
     }
 
@@ -2078,7 +2121,9 @@ impl Expr {
             | Expr::Wildcard { .. }
             | Expr::WindowFunction(..)
             | Expr::Literal(..)
-            | Expr::Placeholder(..) => false,
+            | Expr::Placeholder(..)
+            | Expr::Lambda(..)
+            | Expr::LambdaColumn(..) => false,
         }
     }
 
@@ -2674,6 +2719,17 @@ impl HashNode for Expr {
                 column.hash(state);
             }
             Expr::Unnest(Unnest { expr: _expr }) => {}
+            Expr::Lambda(Lambda { params, body: _ }) => {
+                params.hash(state);
+            }
+            Expr::LambdaColumn(LambdaColumn {
+                name,
+                field,
+                spans: _,
+            }) => {
+                name.hash(state);
+                field.hash(state);
+            }
         };
     }
 }
@@ -2987,6 +3043,12 @@ impl Display for SchemaDisplay<'_> {
                     }
                 }
             }
+            Expr::Lambda(Lambda { params, body }) => {
+                write!(f, "({}) -> {body}", display_comma_separated(params))
+            }
+            Expr::LambdaColumn(c) => {
+                write!(f, "{}", c.name)
+            }
         }
     }
 }
@@ -3167,6 +3229,9 @@ impl Display for SqlDisplay<'_> {
                     }
                 }
             }
+            Expr::Lambda(Lambda { params, body }) => {
+                write!(f, "({}) -> {}", params.join(", "), SchemaDisplay(body))
+            }
             _ => write!(f, "{}", self.0),
         }
     }
@@ -3474,6 +3539,12 @@ impl Display for Expr {
             Expr::Unnest(Unnest { expr }) => {
                 write!(f, "{UNNEST_COLUMN_PREFIX}({expr})")
             }
+            Expr::Lambda(Lambda { params, body }) => {
+                write!(f, "({}) -> {body}", params.join(", "))
+            }
+            Expr::LambdaColumn(c) => {
+                write!(f, "{}", c.name)
+            }
         }
     }
 }

datafusion/expr/src/expr_schema.rs

@@ -18,17 +18,22 @@
 use super::{Between, Expr, Like};
 use crate::expr::{
     AggregateFunction, AggregateFunctionParams, Alias, BinaryExpr, Cast, InList,
-    InSubquery, Placeholder, ScalarFunction, TryCast, Unnest, WindowFunction,
+    InSubquery, Lambda, Placeholder, ScalarFunction, TryCast, Unnest, WindowFunction,
     WindowFunctionParams,
 };
+use crate::expr::{FieldMetadata, LambdaColumn};
 use crate::type_coercion::functions::{
-    data_types_with_scalar_udf, fields_with_aggregate_udf, fields_with_window_udf,
+    fields_with_aggregate_udf, fields_with_window_udf,
+};
+use crate::{
+    type_coercion::functions::data_types_with_scalar_udf, udf::ReturnFieldArgs, utils,
+    LogicalPlan, Projection, Subquery, WindowFunctionDefinition,
+};
+use arrow::datatypes::FieldRef;
+use arrow::{
+    compute::can_cast_types,
+    datatypes::{DataType, Field},
 };
-use crate::udf::ReturnFieldArgs;
-use crate::{utils, LogicalPlan, Projection, Subquery, WindowFunctionDefinition};
-use arrow::compute::can_cast_types;
-use arrow::datatypes::{DataType, Field, FieldRef};
-use datafusion_common::metadata::FieldMetadata;
 use datafusion_common::{
     not_impl_err, plan_datafusion_err, plan_err, Column, DataFusionError, ExprSchema,
     Result, Spans, TableReference,
@@ -229,6 +234,10 @@ impl ExprSchemable for Expr {
                 // Grouping sets do not really have a type and do not appear in projections
                 Ok(DataType::Null)
             }
+            Expr::Lambda(Lambda { params: _, body }) => body.get_type(schema),
+            Expr::LambdaColumn(LambdaColumn { name: _, field, .. }) => {
+                Ok(field.data_type().clone())
+            }
         }
     }
 
@@ -347,6 +356,8 @@ impl ExprSchemable for Expr {
                 // in projections
                 Ok(true)
             }
+            Expr::Lambda(l) => l.body.nullable(input_schema),
+            Expr::LambdaColumn(c) => Ok(c.field.is_nullable()),
         }
     }
 
@@ -543,6 +554,7 @@ impl ExprSchemable for Expr {
                     .into_iter()
                     .map(|f| (f.data_type().clone(), f))
                     .unzip();
+
                 // Verify that function is invoked with correct number and type of arguments as defined in `TypeSignature`
                 let new_data_types = data_types_with_scalar_udf(&arg_types, func)
                     .map_err(|err| {
@@ -573,9 +585,16 @@ impl ExprSchemable for Expr {
                         _ => None,
                     })
                     .collect::<Vec<_>>();
+
+                let lambdas = args
+                    .iter()
+                    .map(|e| matches!(e, Expr::Lambda { .. }))
+                    .collect::<Vec<_>>();
+
                 let args = ReturnFieldArgs {
                     arg_fields: &new_fields,
                     scalar_arguments: &arguments,
+                    lambdas: &lambdas,
                 };
 
                 func.return_field_from_args(args)
@@ -600,11 +619,13 @@ impl ExprSchemable for Expr {
             | Expr::Wildcard { .. }
             | Expr::GroupingSet(_)
             | Expr::Placeholder(_)
-            | Expr::Unnest(_) => Ok(Arc::new(Field::new(
+            | Expr::Unnest(_)
+            | Expr::Lambda(_) => Ok(Arc::new(Field::new(
                 &schema_name,
                 self.get_type(schema)?,
                 self.nullable(schema)?,
             ))),
+            Expr::LambdaColumn(c) => Ok(Arc::clone(&c.field)),
         }?;
 
         Ok((