You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/04/28 09:03:50 UTC
(datafusion) branch gh-readonly-queue/main/pr-21679-54a551583bb13356c42e3c1d9e11b63079d2e40b created (now bbf67d999c)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-21679-54a551583bb13356c42e3c1d9e11b63079d2e40b
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at bbf67d999c Add lambda support and array_transform udf (#21679)
This branch includes the following new commits:
new bbf67d999c Add lambda support and array_transform udf (#21679)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: Add lambda support and array_transform udf (#21679)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-21679-54a551583bb13356c42e3c1d9e11b63079d2e40b
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit bbf67d999c0dcbaf47956233a5d8e78458c13ffa
Author: gstvg <28...@users.noreply.github.com>
AuthorDate: Tue Apr 28 06:03:28 2026 -0300
Add lambda support and array_transform udf (#21679)
This a clean version of #18921 to make it easier to review
**this is a breaking change due to adding variant to `Expr` enum, new
methods on traits `Session`, `FunctionRegistry` and `ContextProvider`
and a new arg on `TaskContext::new`**
This PR adds support for lambdas and the `array_transform` function used
to test the lambda implementation.
Example usage:
```sql
SELECT array_transform([2, 3], v -> v != 2);
[false, true]
-- arbitrally nested lambdas are also supported
SELECT array_transform([[[2, 3]]], m -> array_transform(m, l -> array_transform(l, v -> v*2)));
[[[4, 6]]]
```
Note: column capture has been removed for now and will be added on a
follow on PR, see #21172
Some comments on code snippets of this doc show what value each struct,
variant or field would hold after planning the first example above. Some
literals are simplified pseudo code
3 new `Expr` variants are added, `HigherOrderFunction`, owing a new
trait `HigherOrderUDF`, which is like a `ScalarFunction`/`ScalarUDFImpl`
with support for lambdas, `Lambda`, for the lambda body and it's
parameters names, and `LambdaVariable`, which is like `Column` but for
lambdas parameters.
Their logical representations:
```rust
enum Expr {
// array_transform([2, 3], v -> v != 2)
HigherOrderFunction(HigherOrderFunction),
// v -> v != 2
Lambda(Lambda),
// v, of the lambda body: v != 2
LambdaVariable(LambdaVariable),
...
}
// array_transform([2, 3], v -> v != 2)
struct HigherOrderFunction {
// global instance of array_transform
pub func: Arc<dyn HigherOrderUDF>,
// [Expr::ScalarValue([2, 3]), Expr::Lambda(v -> v != 2)]
pub args: Vec<Expr>,
}
// v -> v != 2
struct Lambda {
// ["v"]
pub params: Vec<String>,
// v != 2
pub body: Box<Expr>,
}
// v, of the lambda body: v != 2
struct LambdaVariable {
// "v"
pub name: String,
// Field::new("", DataType::Int32, false)
// Note: a follow on PR will make this field optional
// to free expr_api from specifying it beforehand,
// and add resolve_lambda_variables method to Expr,
// similar to Expr::Placeholder, see #21172
pub field: FieldRef,
pub spans: Spans,
}
```
The example would be planned into a tree like this:
```
HigherOrderFunctionExpression
name: array_transform
children:
1. ListExpression [2,3]
2. LambdaExpression
parameters: ["v"]
body:
BinaryExpression (!=)
left:
LambdaVariableExpression("v", Field::new("", Int32, false))
right:
LiteralExpression("2")
```
The physical counterparts definition:
```rust
struct HigherOrderFunctionExpr {
// global instance of array_transform
fun: Arc<dyn HigherOrderUDF>,
// "array_transform"
name: String,
// [LiteralExpr([2, 3], LambdaExpr("v -> v != 2"))]
args: Vec<Arc<dyn PhysicalExpr>>,
// [1], the positions at args that contains lambdas
lambda_positions: Vec<usize>,
// Field::new("", DataType::new_list(DataType::Boolean, false), false)
return_field: FieldRef,
config_options: Arc<ConfigOptions>,
}
struct LambdaExpr {
// ["v"]
params: Vec<String>,
// v -> v != 2
body: Arc<dyn PhysicalExpr>,
}
struct LambdaVariable {
// Field::new("v", DataType::Int32, false)
field: FieldRef,
// 0, the first and only parameter, "v"
index: usize,
}
```
Note: For those who primarly wants to check if this lambda
implementation supports their usecase and don't want to spend much time
here, it's okay to skip most collapsed blocks, as those serve mostly to
help code reviewers, with the exception of `HigherOrderUDF` and the
`array_transform` implementation of `HigherOrderUDF` relevant methods,
collapsed due to their size
The added `HigherOrderUDF` trait is almost a clone of `ScalarUDFImpl`,
with the exception of:
1. `return_field_from_args` and `invoke_with_args`, where now
`args.args` is a list of enums with two variants: `Value` or `Lambda`
instead of a list of values
2. the addition of `lambda_parameters`, which return a `Field` for each
parameter supported for every lambda argument based on the `Field` of
the non lambda arguments
3. the removal of `return_field` and the deprecated ones `is_nullable`
and `display_name`.
4. Not yet includes analogues to the methods preimage, placement,
evaluate_bounds, propagate_constraints, output_ordering and
preserves_lex_ordering
<details><summary>HigherOrderUDF</summary>
```rust
trait HigherOrderUDF {
/// Return the field of all the parameters supported by all the supported lambdas of this function
/// based on the field of the value arguments. If a lambda support multiple parameters, or if multiple
/// lambdas are supported and some are optional, all should be returned,
/// regardless of whether they are used on a particular invocation
///
/// Tip: If you have a [`HigherOrderFunction`] invocation, you can call the helper
/// [`HigherOrderFunction::lambda_parameters`] instead of this method directly
///
/// [`HigherOrderFunction`]: crate::expr::HigherOrderFunction
/// [`HigherOrderFunction::lambda_parameters`]: crate::expr::HigherOrderFunction::lambda_parameters
///
/// Example for array_transform:
///
/// `array_transform([2.0, 8.0], v -> v > 4.0)`
///
/// ```ignore
/// let lambda_parameters = array_transform.lambda_parameters(&[
/// Arc::new(Field::new("", DataType::new_list(DataType::Float32, false))), // the Field of the literal `[2, 8]`
/// ])?;
///
/// assert_eq!(
/// lambda_parameters,
/// vec![
/// // the lambda supported parameters, regardless of how many are actually used
/// vec![
/// // the value being transformed
/// Field::new("", DataType::Float32, false),
/// ]
/// ]
/// )
/// ```
///
/// The implementation can assume that some other part of the code has coerced
/// the actual argument types to match [`Self::signature`].
fn lambda_parameters(&self, value_fields: &[FieldRef]) -> Result<Vec<Vec<Field>>>;
fn return_field_from_args(&self, args: LambdaReturnFieldArgs) -> Result<FieldRef>;
fn invoke_with_args(&self, args: HigherOrderFunctionArgs) -> Result<ColumnarValue>;
// ... omitted methods that are similar in ScalarUDFImpl
}
/// An argument to a lambda function
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ValueOrLambda<V, L> {
/// A value with associated data
Value(V),
/// A lambda with associated data
Lambda(L),
}
/// Information about arguments passed to the function
///
/// This structure contains metadata about how the function was called
/// such as the type of the arguments, any scalar arguments and if the
/// arguments can (ever) be null
///
/// See [`HigherOrderUDF::return_field_from_args`] for more information
#[derive(Clone, Debug)]
pub struct LambdaReturnFieldArgs<'a> {
/// The data types of the arguments to the function
///
/// If argument `i` to the function is a lambda, it will be the field of the result of the
/// lambda if evaluated with the parameters returned from [`HigherOrderUDF::lambda_parameters`]
///
/// For example, with `array_transform([1], v -> v == 5)`
/// this field will be `[
/// ValueOrLambda::Value(Field::new("", DataType::List(DataType::Int32), false)),
/// ValueOrLambda::Lambda(Field::new("", DataType::Boolean, false))
/// ]`
pub arg_fields: &'a [ValueOrLambda<FieldRef, FieldRef>],
/// Is argument `i` to the function a scalar (constant)?
///
/// If the argument `i` is not a scalar, it will be None
///
/// For example, if a function is called like `array_transform([1], v -> v == 5)`
/// this field will be `[Some(ScalarValue::List(...), None]`
pub scalar_arguments: &'a [Option<&'a ScalarValue>],
}
/// Arguments passed to [`HigherOrderUDF::invoke_with_args`] when invoking a
/// lambda function.
#[derive(Debug, Clone)]
pub struct HigherOrderFunctionArgs {
/// The evaluated arguments and lambdas to the function
pub args: Vec<ValueOrLambda<ColumnarValue, LambdaArgument>>,
/// Field associated with each arg, if it exists
/// For lambdas, it will be the field of the result of
/// the lambda if evaluated with the parameters
/// returned from [`HigherOrderUDF::lambda_parameters`]
pub arg_fields: Vec<ValueOrLambda<FieldRef, FieldRef>>,
/// The number of rows in record batch being evaluated
pub number_rows: usize,
/// The return field of the lambda function returned
/// (from `return_field_from_args`) when creating the
/// physical expression from the logical expression
pub return_field: FieldRef,
/// The config options at execution time
pub config_options: Arc<ConfigOptions>,
}
/// A lambda argument to a HigherOrderFunction
#[derive(Clone, Debug)]
pub struct LambdaArgument {
/// The parameters defined in this lambda
///
/// For example, for `array_transform([2], v -> -v)`,
/// this will be `vec![Field::new("v", DataType::Int32, true)]`
params: Vec<FieldRef>,
/// The body of the lambda
///
/// For example, for `array_transform([2], v -> -v)`,
/// this will be the physical expression of `-v`
body: Arc<dyn PhysicalExpr>,
}
impl LambdaArgument {
/// Evaluate this lambda
/// `args` should evalute to the value of each parameter
/// of the correspondent lambda returned in [HigherOrderUDF::lambda_parameters].
pub fn evaluate(
&self,
args: &[&dyn Fn() -> Result<ArrayRef>],
) -> Result<ColumnarValue> {
let columns = args
.iter()
.take(self.params.len())
.map(|arg| arg())
.collect::<Result<_>>()?;
let schema = Arc::new(Schema::new(self.params.clone()));
let batch = RecordBatch::try_new(schema, columns)?;
self.body.evaluate(&batch)
}
}
```
</details>
<details><summary>array_transform lambda_parameters
implementation</summary>
```rust
impl HigherOrderUDF for ArrayTransform {
fn lambda_parameters(&self, value_fields: &[FieldRef]) -> Result<Vec<Vec<Field>>> {
let list = if value_fields.len() == 1 {
&value_fields[0]
} else {
return plan_err!(
"{} function requires 1 value arguments, got {}",
self.name(),
value_fields.len()
);
};
let field = match list.data_type() {
DataType::List(field) => field,
DataType::LargeList(field) => field,
DataType::FixedSizeList(field, _) => field,
_ => return plan_err!("expected list, got {list}"),
};
// we don't need to check whether the lambda contains more than two parameters,
// e.g. array_transform([], (v, i, j) -> v+i+j), as datafusion will do that for us
let value = Field::new("", field.data_type().clone(), field.is_nullable())
.with_metadata(field.metadata().clone());
Ok(vec![vec![value]])
}
}
```
</details>
<details><summary>array_transform return_field_from_args
implementation</summary>
```rust
fn value_lambda_pair<'a, V: Debug, L: Debug>(
name: &str,
args: &'a [ValueOrLambda<V, L>],
) -> Result<(&'a V, &'a L)> {
let [value, lambda] = take_function_args(name, args)?;
let (ValueOrLambda::Value(value), ValueOrLambda::Lambda(lambda)) = (value, lambda)
else {
return plan_err!(
"{name} expects a value followed by a lambda, got {value:?} and {lambda:?}"
);
};
Ok((value, lambda))
}
impl HigherOrderUDF for ArrayTransform {
fn return_field_from_args(
&self,
args: HigherOrderReturnFieldArgs,
) -> Result<Arc<Field>> {
let (list, lambda) = value_lambda_pair(self.name(), args.arg_fields)?;
// lambda is the resulting field of executing the lambda body
// with the parameters returned in lambda_parameters
let field = Arc::new(Field::new(
Field::LIST_FIELD_DEFAULT_NAME,
lambda.data_type().clone(),
lambda.is_nullable(),
));
let return_type = match list.data_type() {
DataType::List(_) => DataType::List(field),
DataType::LargeList(_) => DataType::LargeList(field),
DataType::FixedSizeList(_, size) => DataType::FixedSizeList(field, *size),
other => plan_err!("expected list, got {other}")?,
};
Ok(Arc::new(Field::new("", return_type, list.is_nullable())))
}
}
```
</details>
<details><summary>array_transform invoke_with_args
implementation</summary>
```rust
impl HigherOrderUDF for ArrayTransform {
fn invoke_with_args(&self, args: HigherOrderFunctionArgs) -> Result<ColumnarValue> {
let (list, lambda) = value_lambda_pair(self.name(), &args.args)?;
let list_array = list.to_array(args.number_rows)?;
// Fast path for fully null input array and also the only way to safely work with
// a fully null fixed size list array as it can't be handled by remove_list_null_values below
if list_array.null_count() == list_array.len() {
return Ok(ColumnarValue::Array(new_null_array(
args.return_type(),
list_array.len(),
)));
}
// as per list_values docs, if list_array is sliced, list_values will be sliced too,
// so before constructing the transformed array below, we must adjust the list offsets with
// adjust_offsets_for_slice
let list_values = list_values(&list_array)?;
// by passing closures, lambda.evaluate can evaluate only those actually needed
let values_param = || Ok(Arc::clone(&list_values));
// call the transforming lambda
let transformed_values = lambda
.evaluate(&[&values_param])?
.into_array(list_values.len())?;
let field = match args.return_field.data_type() {
DataType::List(field)
| DataType::LargeList(field)
| DataType::FixedSizeList(field, _) => Arc::clone(field),
_ => {
return exec_err!(
"{} expected ScalarFunctionArgs.return_field to be a list, got {}",
self.name(),
args.return_field
);
}
};
let transformed_list = match list_array.data_type() {
DataType::List(_) => {
let list = list_array.as_list();
// since we called list_values above which would return sliced values for
// a sliced list, we must adjust the offsets here as otherwise they would be invalid
let adjusted_offsets = adjust_offsets_for_slice(list);
Arc::new(ListArray::new(
field,
adjusted_offsets,
transformed_values,
list.nulls().cloned(),
)) as ArrayRef
}
DataType::LargeList(_) => {
let large_list = list_array.as_list();
// since we called list_values above which would return sliced values for
// a sliced list, we must adjust the offsets here as otherwise they would be invalid
let adjusted_offsets = adjust_offsets_for_slice(large_list);
Arc::new(LargeListArray::new(
field,
adjusted_offsets,
transformed_values,
large_list.nulls().cloned(),
))
}
DataType::FixedSizeList(_, value_length) => {
Arc::new(FixedSizeListArray::new(
field,
*value_length,
transformed_values,
list_array.as_fixed_size_list().nulls().cloned(),
))
}
other => exec_err!("expected list, got {other}")?,
};
Ok(ColumnarValue::Array(transformed_list))
}
}
```
</details>
<details><summary>How relevant HigherOrderUDF methods would be called
and what they would return during planning and evaluation of the
example</summary>
```rust
// this is called at sql planning
let lambda_parameters = lambda_udf.lambda_parameters(&[
Field::new("", DataType::new_list(DataType::Int32, false), false), // the Field of the [2, 3] literal
])?;
assert_eq!(
lambda_parameters,
vec![
// the parameters that *can* be declared on the lambda, and not only
// those actually declared: the implementation doesn't need to care
// about it
vec![
Field::new("", DataType::Int32, false), // the list inner value
]]
);
// this is called every time ExprSchemable is called on a HigherOrderFunction
let return_field = array_transform.return_field_from_args(&LambdaReturnFieldArgs {
arg_fields: &[
ValueOrLambda::Value(Field::new("", DataType::new_list(DataType::Int32, false), false)),
ValueOrLambda::Lambda(Field::new("", DataType::Boolean, false)), // the return_field of the expression "v != 2" when "v" is of the type returned in lambda_parameters
],
scalar_arguments // irrelevant
})?;
assert_eq!(return_field, Field::new("", DataType::new_list(DataType::Boolean, false), false));
let value = array_transform.evaluate(&HigherOrderFunctionArgs {
args: vec![
ValueOrLambda::Value(List([2, 3])),
ValueOrLambda::Lambda(LambdaArgument of `v -> v != 2`),
],
arg_fields, // same as above
number_rows: 1,
return_field, // same as above
config_options, // irrelevant
})?;
assert_eq!(value, BooleanArray::from([false, true]))
```
</details>
<br>
<br>
A pair HigherOrderUDF/HigherOrderUDFImpl like ScalarFunction was not
used because those exist only [to maintain backwards compatibility with
the older
API](https://docs.rs/datafusion/latest/datafusion/logical_expr/struct.ScalarUDF.html#api-note)
#8045
</details>
<br>
Why `LambdaVariable` and not `Column`:
Existing tree traversals that operate on columns would break if some
column nodes referenced to a lambda parameter and not a real column. In
the example query, projection pushdown would try to push the lambda
parameter "v", which won't exist in table "t".
Example of code of another traversal that would break:
```rust
fn minimize_join_filter(expr: Arc<dyn PhysicalExpr>, ...) -> JoinFilter {
let mut used_columns = HashSet::new();
expr.apply(|expr| {
if let Some(col) = expr.as_any().downcast_ref::<Column>() {
// if this is a lambda column, this function will break
used_columns.insert(col.index());
}
Ok(TreeNodeRecursion::Continue)
});
...
}
```
Furthermore, the implemention of `ExprSchemable` and
`PhysicalExpr::return_field` for `Column` expects that the schema it
receives as a argument contains an entry for its name, which is not the
case for lambda parameters.
By including a `FieldRef` on `LambdaVariable` that should be resolved
during construction time in the sql planner, `ExprSchemable` and
`PhysicalExpr::return_field` simply return it's own Field:
<details><summary>LambdaVariable ExprSchemable and
PhysicalExpr::return_field implementation </summary>
```rust
impl ExprSchemable for Expr {
fn to_field(
&self,
schema: &dyn ExprSchema,
) -> Result<(Option<TableReference>, Arc<Field>)> {
let (relation, schema_name) = self.qualified_name();
let field = match self {
Expr::LambdaVariable(l) => Ok(Arc::clone(&l.field)),
...
}?;
Ok((
relation,
Arc::new(field.as_ref().clone().with_name(schema_name)),
))
}
...
}
impl PhysicalExpr for LambdaVariable {
fn return_field(&self, _input_schema: &Schema) -> Result<FieldRef> {
Ok(Arc::clone(&self.field))
}
...
}
```
</details>
<br>
<details><summary>Possible alternatives discarded due to complexity,
requiring downstream changes and implementation size:</summary>
1. Add a new set of TreeNode methods that provides the set of lambdas
parameters names seen during the traversal, so column nodes can be
tested if they refer to a regular column or to a lambda parameter. Any
downstream user that wants to support lambdas would need use those
methods instead of the existing ones. This also would add 1k+ lines to
the PR.
```rust
impl Expr {
pub fn transform_with_lambdas_params<
F: FnMut(Self, &HashSet<String>) -> Result<Transformed<Self>>,
>(
self,
mut f: F,
) -> Result<Transformed<Self>> {}
}
```
How minimize_join_filter would looks like:
```rust
fn minimize_join_filter(expr: Arc<dyn PhysicalExpr>, ...) -> JoinFilter {
let mut used_columns = HashSet::new();
expr.apply_with_lambdas_params(|expr, lambdas_params| {
if let Some(col) = expr.as_any().downcast_ref::<Column>() {
// dont include lambdas parameters
if !lambdas_params.contains(col.name()) {
used_columns.insert(col.index());
}
}
Ok(TreeNodeRecursion::Continue)
})
...
}
```
2. Add a flag to the Column node indicating if it refers to a lambda
parameter. Still requires checking for it on existing tree traversals
that works on Columns (30+) and also downstream.
```rust
//logical
struct Column {
pub relation: Option<TableReference>,
pub name: String,
pub spans: Spans,
pub is_lambda_parameter: bool,
}
//physical
struct Column {
name: String,
index: usize,
is_lambda_parameter: bool,
}
```
How minimize_join_filter would look like:
```rust
fn minimize_join_filter(expr: Arc<dyn PhysicalExpr>, ...) -> JoinFilter {
let mut used_columns = HashSet::new();
expr.apply(|expr| {
if let Some(col) = expr.as_any().downcast_ref::<Column>() {
// dont include lambdas parameters
if !col.is_lambda_parameter {
used_columns.insert(col.index());
}
}
Ok(TreeNodeRecursion::Continue)
})
...
}
```
1. Add a new set of TreeNode methods that provides a schema that
includes the lambdas parameters for the scope of the node being
visited/transformed:
```rust
impl Expr {
pub fn transform_with_schema<
F: FnMut(Self, &DFSchema) -> Result<Transformed<Self>>,
>(
self,
schema: &DFSchema,
f: F,
) -> Result<Transformed<Self>> { ... }
... other methods
}
```
For any given HigherOrderFunction found during the traversal, a new
schema is created for each lambda argument that contains it's parameter,
returned from HigherOrderUDF::lambda_parameters
How it would look like:
```rust
pub fn infer_placeholder_types(self, schema: &DFSchema) -> Result<(Expr, bool)> {
let mut has_placeholder = false;
// Provide the schema as the first argument.
// Transforming closure receive an adjusted_schema as argument
self.transform_with_schema(schema, |mut expr, adjusted_schema| {
match &mut expr {
// Default to assuming the arguments are the same type
Expr::BinaryExpr(BinaryExpr { left, op: _, right }) => {
// use adjusted_schema and not schema. Those expressions may contain
// columns referring to a lambda parameter, which Field would only be
// available in adjusted_schema and not in schema
rewrite_placeholder(left.as_mut(), right.as_ref(), adjusted_schema)?;
rewrite_placeholder(right.as_mut(), left.as_ref(), adjusted_schema)?;
}
....
```
2. Make available trought LogicalPlan and ExecutionPlan nodes a schema
that includes all lambdas parameters from all expressions owned by the
node, and use this schema for tree traversals. For nodes which won't own
any expression, the regular schema can be returned
```rust
impl LogicalPlan {
fn lambda_extended_schema(&self) -> &DFSchema;
}
trait ExecutionPlan {
fn lambda_extended_schema(&self) -> &DFSchema;
}
//usage
impl LogicalPlan {
pub fn replace_params_with_values(
self,
param_values: &ParamValues,
) -> Result<LogicalPlan> {
self.transform_up_with_subqueries(|plan| {
// use plan.lambda_extended_schema() containing lambdas parameters
// instead of plan.schema() which wont
let lambda_extended_schema = Arc::clone(plan.lambda_extended_schema());
let name_preserver = NamePreserver::new(&plan);
plan.map_expressions(|e| {
// if this expression is child of lambda and contain columns referring it's parameters
// the lambda_extended_schema already contain them
let (e, has_placeholder) = e.infer_placeholder_types(&lambda_extended_schema)?;
....
```
</details>
<br>
---------
Co-authored-by: Raz Luvaton <16...@users.noreply.github.com>
Co-authored-by: Martin Grigorov <ma...@users.noreply.github.com>
Co-authored-by: Lía Adriana <li...@datadoghq.com>
---
Cargo.lock | 1 +
datafusion-examples/examples/sql_ops/frontend.rs | 12 +-
datafusion/catalog-listing/src/helpers.rs | 17 +-
datafusion/common/src/utils/mod.rs | 260 ++++++-
datafusion/core/src/bin/print_functions_docs.rs | 19 +-
.../core/src/datasource/listing_table_factory.rs | 5 +
datafusion/core/src/execution/context/mod.rs | 16 +
datafusion/core/src/execution/session_state.rs | 133 +++-
.../core/src/execution/session_state_defaults.rs | 11 +-
datafusion/core/tests/optimizer/mod.rs | 12 +-
datafusion/datasource-arrow/src/file_format.rs | 8 +-
datafusion/datasource/src/url.rs | 8 +-
datafusion/execution/src/task.rs | 41 +-
datafusion/expr/src/expr.rs | 185 ++++-
datafusion/expr/src/expr_fn.rs | 20 +-
datafusion/expr/src/expr_schema.rs | 75 +-
datafusion/expr/src/higher_order_function.rs | 772 +++++++++++++++++++++
datafusion/expr/src/lib.rs | 6 +
datafusion/expr/src/planner.rs | 10 +-
datafusion/expr/src/registry.rs | 52 ++
datafusion/expr/src/tree_node.rs | 38 +-
datafusion/expr/src/type_coercion/functions.rs | 364 +++++++++-
datafusion/expr/src/udf_eq.rs | 14 +-
datafusion/expr/src/utils.rs | 5 +-
datafusion/ffi/src/execution/task_ctx.rs | 2 +
datafusion/ffi/src/session/mod.rs | 10 +-
datafusion/functions-nested/Cargo.toml | 2 +
datafusion/functions-nested/src/array_transform.rs | 392 +++++++++++
datafusion/functions-nested/src/lib.rs | 41 +-
datafusion/functions-nested/src/macros_lambda.rs | 107 +++
datafusion/optimizer/src/analyzer/type_coercion.rs | 46 +-
.../optimizer/src/common_subexpr_eliminate.rs | 10 +-
datafusion/optimizer/src/push_down_filter.rs | 5 +-
.../src/simplify_expressions/expr_simplifier.rs | 8 +-
.../optimizer/tests/optimizer_integration.rs | 11 +
datafusion/physical-expr/src/expressions/case.rs | 23 +-
datafusion/physical-expr/src/expressions/lambda.rs | 160 +++++
.../src/expressions/lambda_variable.rs | 145 ++++
datafusion/physical-expr/src/expressions/mod.rs | 4 +
.../physical-expr/src/higher_order_function.rs | 686 ++++++++++++++++++
datafusion/physical-expr/src/lib.rs | 2 +
datafusion/physical-expr/src/planner.rs | 113 ++-
datafusion/proto/src/logical_plan/mod.rs | 21 +-
datafusion/proto/src/logical_plan/to_proto.rs | 5 +
datafusion/session/src/session.rs | 8 +-
datafusion/spark/src/lib.rs | 4 +-
datafusion/sql/examples/sql.rs | 10 +-
datafusion/sql/src/expr/function.rs | 256 ++++++-
datafusion/sql/src/expr/identifier.rs | 14 +
datafusion/sql/src/expr/mod.rs | 12 +-
datafusion/sql/src/planner.rs | 17 +
datafusion/sql/src/unparser/dialect.rs | 12 +
datafusion/sql/src/unparser/expr.rs | 104 ++-
datafusion/sql/tests/common/mod.rs | 24 +-
datafusion/sql/tests/sql_integration.rs | 175 ++++-
.../sqllogictest/test_files/higher_order.slt | 286 ++++++++
.../logical_plan/consumer/expr/scalar_function.rs | 1 +
.../src/logical_plan/producer/expr/mod.rs | 7 +
.../logical_plan/producer/expr/scalar_function.rs | 23 +-
.../logical_plan/producer/substrait_producer.rs | 19 +-
docs/source/user-guide/sql/scalar_functions.md | 34 +
61 files changed, 4784 insertions(+), 99 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index 899b78efcb..033a430c57 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2248,6 +2248,7 @@ dependencies = [
"datafusion-functions-aggregate",
"datafusion-functions-aggregate-common",
"datafusion-macros",
+ "datafusion-physical-expr",
"datafusion-physical-expr-common",
"hashbrown 0.17.0",
"itertools 0.14.0",
diff --git a/datafusion-examples/examples/sql_ops/frontend.rs b/datafusion-examples/examples/sql_ops/frontend.rs
index d7704f74db..b34c720a78 100644
--- a/datafusion-examples/examples/sql_ops/frontend.rs
+++ b/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, HigherOrderUDF, LogicalPlan, ScalarUDF,
+ TableProviderFilterPushDown, TableSource, WindowUDF,
};
use datafusion::optimizer::{
Analyzer, AnalyzerRule, Optimizer, OptimizerConfig, OptimizerContext, OptimizerRule,
@@ -154,6 +154,10 @@ impl ContextProvider for MyContextProvider {
None
}
+ fn get_higher_order_meta(&self, _name: &str) -> Option<Arc<dyn HigherOrderUDF>> {
+ None
+ }
+
fn get_aggregate_meta(&self, _name: &str) -> Option<Arc<AggregateUDF>> {
None
}
@@ -174,6 +178,10 @@ impl ContextProvider for MyContextProvider {
Vec::new()
}
+ fn higher_order_function_names(&self) -> Vec<String> {
+ Vec::new()
+ }
+
fn udaf_names(&self) -> Vec<String> {
Vec::new()
}
diff --git a/datafusion/catalog-listing/src/helpers.rs b/datafusion/catalog-listing/src/helpers.rs
index 62257ec027..c6305c3000 100644
--- a/datafusion/catalog-listing/src/helpers.rs
+++ b/datafusion/catalog-listing/src/helpers.rs
@@ -85,12 +85,26 @@ 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 {
Volatility::Immutable => Ok(TreeNodeRecursion::Continue),
// TODO: Stable functions could be `applicable`, but that would require access to the context
+ // https://github.com/apache/datafusion/issues/21690
+ Volatility::Stable | Volatility::Volatile => {
+ is_applicable = false;
+ Ok(TreeNodeRecursion::Stop)
+ }
+ }
+ }
+ Expr::HigherOrderFunction(hof) => {
+ match hof.func.signature().volatility {
+ Volatility::Immutable => Ok(TreeNodeRecursion::Continue),
+ // TODO: Stable functions could be `applicable`, but that would require access to the context
+ // https://github.com/apache/datafusion/issues/21690
Volatility::Stable | Volatility::Volatile => {
is_applicable = false;
Ok(TreeNodeRecursion::Stop)
@@ -102,6 +116,7 @@ pub fn expr_applicable_for_cols(col_names: &[&str], expr: &Expr) -> bool {
// - AGGREGATE and WINDOW should not end up in filter conditions, except maybe in some edge cases
// - Can `Wildcard` be considered as a `Literal`?
// - ScalarVariable could be `applicable`, but that would require access to the context
+ // https://github.com/apache/datafusion/issues/21690
// TODO: remove the next line after `Expr::Wildcard` is removed
#[expect(deprecated)]
Expr::AggregateFunction { .. }
diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs
index 2f01859f45..de00dac825 100644
--- a/datafusion/common/src/utils/mod.rs
+++ b/datafusion/common/src/utils/mod.rs
@@ -24,14 +24,19 @@ 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,
cast::AsArray,
};
+use arrow::array::{
+ Datum, GenericListArray, Int32Array, Int64Array, MutableArrayData, make_array,
+};
use arrow::array::{LargeListViewArray, ListViewArray};
use arrow::buffer::{OffsetBuffer, ScalarBuffer};
+use arrow::compute::kernels::cmp::neq;
+use arrow::compute::kernels::length::length;
use arrow::compute::{SortColumn, SortOptions, partition};
use arrow::datatypes::{DataType, Field, SchemaRef};
#[cfg(feature = "sql")]
@@ -1005,11 +1010,137 @@ 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(sliced_list_values(array.as_list::<i32>())),
+ DataType::LargeList(_) => Ok(sliced_list_values(array.as_list::<i64>())),
+ DataType::FixedSizeList(_, _) => {
+ Ok(Arc::clone(array.as_fixed_size_list().values()))
+ }
+ other => _exec_err!("expected list, got {other}"),
+ }
+}
+
+fn sliced_list_values<O: OffsetSizeTrait>(list: &GenericListArray<O>) -> ArrayRef {
+ let values = list.values();
+ let offsets = list.offsets();
+
+ if let (Some(first), Some(last)) = (offsets.first(), offsets.last()) {
+ let first = first.as_usize();
+ let last = last.as_usize();
+
+ if first != 0 || last != values.len() {
+ return values.slice(first, last - first);
+ }
+ }
+
+ Arc::clone(values)
+}
+
+/// If `list` is sliced, returns an adjusted offset buffer so that
+/// it points to the sliced portion of the list values, and not the whole list values
+pub fn adjust_offsets_for_slice<O: OffsetSizeTrait>(
+ list: &GenericListArray<O>,
+) -> OffsetBuffer<O> {
+ let offsets = list.offsets();
+
+ if let (Some(first), Some(last)) = (offsets.first(), offsets.last())
+ && (!first.is_zero() || last.as_usize() != list.values().len())
+ {
+ let offsets = offsets.iter().map(|offset| *offset - *first).collect();
+
+ //todo: use unsafe Offset::new_unchecked?
+ return OffsetBuffer::new(offsets);
+ }
+
+ offsets.clone()
+}
+
+/// For lists and large lists, truncates the sublist of null values
+/// Otherwise returns an error
+pub fn remove_list_null_values(array: &ArrayRef) -> Result<ArrayRef> {
+ // todo: handle list view and map
+ match array.data_type() {
+ DataType::List(_) => Ok(Arc::new(truncate_list_nulls(array.as_list::<i32>())?)),
+ DataType::LargeList(_) => {
+ Ok(Arc::new(truncate_list_nulls(array.as_list::<i64>())?))
+ }
+ dt => _exec_err!("expected List or LargeList, got {dt}"),
+ }
+}
+
+fn truncate_list_nulls<O: OffsetSizeTrait>(
+ list: &GenericListArray<O>,
+) -> Result<GenericListArray<O>> {
+ if let Some(nulls) = list.nulls()
+ && nulls.null_count() > 0
+ {
+ let lengths = length(list)?;
+ let zero: &dyn Datum = if lengths.data_type() == &DataType::Int32 {
+ &Int32Array::new_scalar(0)
+ } else {
+ &Int64Array::new_scalar(0)
+ };
+
+ let not_empty = neq(&lengths, zero)?;
+ let null_and_non_empty = &!nulls.inner() & not_empty.values();
+
+ if null_and_non_empty.count_set_bits() > 0 {
+ let array_data = list.values().to_data();
+ let offsets = list.offsets();
+ let capacity = offsets[offsets.len() - 1] - offsets[0];
+ let mut mutable_array_data =
+ MutableArrayData::new(vec![&array_data], false, capacity.as_usize());
+
+ let valid_or_empty = nulls.inner() | &!not_empty.values();
+
+ for (start, end) in valid_or_empty.set_slices() {
+ mutable_array_data.extend(
+ 0,
+ offsets[start].as_usize(),
+ offsets[end].as_usize(),
+ );
+ }
+
+ let lengths = std::iter::zip(offsets.lengths(), nulls)
+ .map(|(length, is_valid)| if is_valid { length } else { 0 });
+
+ let offsets = OffsetBuffer::from_lengths(lengths);
+ let values = make_array(mutable_array_data.freeze());
+
+ let field = match list.data_type() {
+ DataType::List(field) => field,
+ DataType::LargeList(field) => field,
+ _ => unreachable!(),
+ };
+
+ return Ok(GenericListArray::try_new(
+ Arc::clone(field),
+ offsets,
+ values,
+ list.nulls().cloned(),
+ )?);
+ }
+ }
+ Ok(list.clone())
+}
+
#[cfg(test)]
mod tests {
+ use std::sync::Arc;
+
use super::*;
use crate::ScalarValue::Null;
- use arrow::array::Float64Array;
+ use arrow::{
+ array::{Float64Array, Int32Array},
+ buffer::NullBuffer,
+ datatypes::Int32Type,
+ };
+ use sqlparser::ast::Ident;
#[test]
fn test_bisect_linear_left_and_right() -> Result<()> {
@@ -1309,4 +1440,129 @@ mod tests {
assert_eq!(expected, transposed);
Ok(())
}
+
+ #[test]
+ fn test_sliced_list_values() {
+ let data = vec![
+ Some(vec![Some(0), Some(1), Some(2)]),
+ None,
+ Some(vec![Some(3), None, Some(5)]),
+ Some(vec![Some(6), Some(7)]),
+ ];
+
+ let list = ListArray::from_iter_primitive::<Int32Type, _, _>(data);
+
+ assert_eq!(
+ sliced_list_values(&list).as_primitive(),
+ &Int32Array::from(vec![
+ Some(0),
+ Some(1),
+ Some(2),
+ Some(3),
+ None,
+ Some(5),
+ Some(6),
+ Some(7)
+ ])
+ );
+
+ assert_eq!(
+ sliced_list_values(&list.slice(0, 1)).as_primitive(),
+ &Int32Array::from(vec![Some(0), Some(1), Some(2)])
+ );
+
+ assert_eq!(
+ sliced_list_values(&list.slice(2, 1)).as_primitive(),
+ &Int32Array::from(vec![Some(3), None, Some(5)])
+ );
+
+ assert_eq!(
+ sliced_list_values(&list.slice(3, 1)).as_primitive(),
+ &Int32Array::from(vec![Some(6), Some(7)])
+ );
+
+ assert!(sliced_list_values(&list.slice(0, 0)).is_empty());
+ assert!(sliced_list_values(&list.slice(1, 0)).is_empty());
+ assert!(sliced_list_values(&list.slice(3, 0)).is_empty());
+ }
+
+ #[test]
+ fn test_adjust_offsets() {
+ let data = vec![
+ Some(vec![Some(0), Some(1), Some(2)]),
+ None,
+ Some(vec![Some(3), None, Some(5)]),
+ Some(vec![Some(6), Some(7)]),
+ ];
+ let list = ListArray::from_iter_primitive::<Int32Type, _, _>(data);
+
+ assert_eq!(
+ adjust_offsets_for_slice(&list),
+ OffsetBuffer::from_lengths([3, 0, 3, 2])
+ );
+
+ assert_eq!(
+ adjust_offsets_for_slice(&list.slice(0, 1)),
+ OffsetBuffer::from_lengths([3])
+ );
+
+ assert_eq!(
+ adjust_offsets_for_slice(&list.slice(1, 2)),
+ OffsetBuffer::from_lengths([0, 3])
+ );
+
+ assert_eq!(
+ adjust_offsets_for_slice(&list.slice(1, 3)),
+ OffsetBuffer::from_lengths([0, 3, 2])
+ );
+
+ assert_eq!(
+ adjust_offsets_for_slice(&list.slice(0, 0)),
+ OffsetBuffer::from_lengths([])
+ );
+
+ assert_eq!(
+ adjust_offsets_for_slice(&list.slice(1, 0)),
+ OffsetBuffer::from_lengths([])
+ );
+
+ assert_eq!(
+ adjust_offsets_for_slice(&list.slice(3, 0)),
+ OffsetBuffer::from_lengths([])
+ );
+ }
+
+ fn create_i32_list(
+ values: impl Into<Int32Array>,
+ offsets: OffsetBuffer<i32>,
+ nulls: Option<NullBuffer>,
+ ) -> ListArray {
+ let list_field = Arc::new(Field::new_list_field(DataType::Int32, true));
+
+ ListArray::new(list_field, offsets, Arc::new(values.into()), nulls)
+ }
+
+ #[test]
+ fn test_remove_list_null_values_list() {
+ let list = Arc::new(create_i32_list(
+ vec![100, 20, 10, 0, 0, 0, 0, 1, 50],
+ OffsetBuffer::<i32>::from_lengths(vec![3, 4, 0, 2, 0]),
+ Some(NullBuffer::from(vec![true, false, false, true, false])),
+ )) as ArrayRef;
+
+ let res = remove_list_null_values(&list).unwrap();
+ let res = res.as_list::<i32>();
+
+ let expected = Arc::new(create_i32_list(
+ vec![100, 20, 10, 1, 50],
+ OffsetBuffer::<i32>::from_lengths(vec![3, 0, 0, 2, 0]),
+ Some(NullBuffer::from(vec![true, false, false, true, false])),
+ )) as ArrayRef;
+ let expected = expected.as_list::<i32>();
+
+ assert_eq!(res, expected);
+ // check above skips inner value of nulls
+ assert_eq!(res.values(), expected.values());
+ assert_eq!(res.offsets(), expected.offsets());
+ }
}
diff --git a/datafusion/core/src/bin/print_functions_docs.rs b/datafusion/core/src/bin/print_functions_docs.rs
index 2466d42692..c34865a32d 100644
--- a/datafusion/core/src/bin/print_functions_docs.rs
+++ b/datafusion/core/src/bin/print_functions_docs.rs
@@ -18,12 +18,13 @@
use datafusion::execution::SessionStateDefaults;
use datafusion_common::{HashSet, Result, not_impl_err};
use datafusion_expr::{
- AggregateUDF, DocSection, Documentation, ScalarUDF, WindowUDF,
+ AggregateUDF, DocSection, Documentation, HigherOrderUDF, ScalarUDF, WindowUDF,
aggregate_doc_sections, scalar_doc_sections, window_doc_sections,
};
use itertools::Itertools;
use std::env::args;
use std::fmt::Write as _;
+use std::sync::Arc;
/// Print documentation for all functions of a given type to stdout
///
@@ -71,6 +72,10 @@ fn print_scalar_docs() -> Result<String> {
providers.push(Box::new(f.as_ref().clone()));
}
+ for f in SessionStateDefaults::default_higher_order_functions() {
+ providers.push(Box::new(f));
+ }
+
print_docs(providers, scalar_doc_sections::doc_sections())
}
@@ -282,6 +287,18 @@ impl DocProvider for WindowUDF {
}
}
+impl DocProvider for Arc<dyn HigherOrderUDF> {
+ 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
diff --git a/datafusion/core/src/datasource/listing_table_factory.rs b/datafusion/core/src/datasource/listing_table_factory.rs
index 31d6e36657..a513934675 100644
--- a/datafusion/core/src/datasource/listing_table_factory.rs
+++ b/datafusion/core/src/datasource/listing_table_factory.rs
@@ -575,6 +575,11 @@ mod tests {
) -> &HashMap<String, Arc<datafusion_expr::ScalarUDF>> {
unimplemented!()
}
+ fn higher_order_functions(
+ &self,
+ ) -> &HashMap<String, Arc<dyn datafusion_expr::HigherOrderUDF>> {
+ unimplemented!()
+ }
fn aggregate_functions(
&self,
) -> &HashMap<String, Arc<datafusion_expr::AggregateUDF>> {
diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs
index 95403dddac..87170f595f 100644
--- a/datafusion/core/src/execution/context/mod.rs
+++ b/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::HigherOrderUDF;
pub use datafusion_expr::execution_props::ExecutionProps;
#[cfg(feature = "sql")]
use datafusion_expr::planner::RelationPlanner;
@@ -1977,6 +1978,10 @@ impl FunctionRegistry for SessionContext {
self.state.read().udf(name)
}
+ fn higher_order_function(&self, name: &str) -> Result<Arc<dyn HigherOrderUDF>> {
+ self.state.read().higher_order_function(name)
+ }
+
fn udaf(&self, name: &str) -> Result<Arc<AggregateUDF>> {
self.state.read().udaf(name)
}
@@ -1989,6 +1994,13 @@ impl FunctionRegistry for SessionContext {
self.state.write().register_udf(udf)
}
+ fn register_higher_order_function(
+ &mut self,
+ function: Arc<dyn HigherOrderUDF>,
+ ) -> Result<Option<Arc<dyn HigherOrderUDF>>> {
+ self.state.write().register_higher_order_function(function)
+ }
+
fn register_udaf(
&mut self,
udaf: Arc<AggregateUDF>,
@@ -2018,6 +2030,10 @@ impl FunctionRegistry for SessionContext {
self.state.write().register_expr_planner(expr_planner)
}
+ fn higher_order_function_names(&self) -> HashSet<String> {
+ self.state.read().higher_order_function_names()
+ }
+
fn udafs(&self) -> HashSet<String> {
self.state.read().udafs()
}
diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs
index a5749e70ce..de5e6b97c1 100644
--- a/datafusion/core/src/execution/session_state.rs
+++ b/datafusion/core/src/execution/session_state.rs
@@ -61,7 +61,9 @@ use datafusion_expr::registry::{
SerializerRegistry,
};
use datafusion_expr::simplify::SimplifyContext;
-use datafusion_expr::{AggregateUDF, Explain, Expr, LogicalPlan, ScalarUDF, WindowUDF};
+use datafusion_expr::{
+ AggregateUDF, Explain, Expr, HigherOrderUDF, LogicalPlan, ScalarUDF, WindowUDF,
+};
use datafusion_optimizer::simplify_expressions::ExprSimplifier;
use datafusion_optimizer::{
Analyzer, AnalyzerRule, Optimizer, OptimizerConfig, OptimizerRule,
@@ -159,6 +161,8 @@ pub struct SessionState {
table_functions: HashMap<String, Arc<TableFunction>>,
/// Scalar functions that are registered with the context
scalar_functions: HashMap<String, Arc<ScalarUDF>>,
+ /// Higher order functions that are registered with the context
+ higher_order_functions: HashMap<String, Arc<dyn HigherOrderUDF>>,
/// Aggregate functions registered in the context
aggregate_functions: HashMap<String, Arc<AggregateUDF>>,
/// Window functions registered in the context
@@ -245,6 +249,7 @@ impl Debug for SessionState {
.field("physical_optimizers", &self.physical_optimizers)
.field("table_functions", &self.table_functions)
.field("scalar_functions", &self.scalar_functions)
+ .field("higher_order_functions", &self.higher_order_functions)
.field("aggregate_functions", &self.aggregate_functions)
.field("window_functions", &self.window_functions)
.field("prepared_plans", &self.prepared_plans)
@@ -281,6 +286,10 @@ impl Session for SessionState {
&self.scalar_functions
}
+ fn higher_order_functions(&self) -> &HashMap<String, Arc<dyn HigherOrderUDF>> {
+ &self.higher_order_functions
+ }
+
fn aggregate_functions(&self) -> &HashMap<String, Arc<AggregateUDF>> {
&self.aggregate_functions
}
@@ -924,6 +933,11 @@ impl SessionState {
&self.scalar_functions
}
+ /// Return reference to higher_order_functions
+ pub fn higher_order_functions(&self) -> &HashMap<String, Arc<dyn HigherOrderUDF>> {
+ &self.higher_order_functions
+ }
+
/// Return reference to aggregate_functions
pub fn aggregate_functions(&self) -> &HashMap<String, Arc<AggregateUDF>> {
&self.aggregate_functions
@@ -1020,6 +1034,7 @@ pub struct SessionStateBuilder {
catalog_list: Option<Arc<dyn CatalogProviderList>>,
table_functions: Option<HashMap<String, Arc<TableFunction>>>,
scalar_functions: Option<Vec<Arc<ScalarUDF>>>,
+ higher_order_functions: Option<Vec<Arc<dyn HigherOrderUDF>>>,
aggregate_functions: Option<Vec<Arc<AggregateUDF>>>,
window_functions: Option<Vec<Arc<WindowUDF>>>,
extension_types: Option<ExtensionTypeRegistryRef>,
@@ -1062,6 +1077,7 @@ impl SessionStateBuilder {
catalog_list: None,
table_functions: None,
scalar_functions: None,
+ higher_order_functions: None,
aggregate_functions: None,
window_functions: None,
extension_types: None,
@@ -1117,6 +1133,9 @@ impl SessionStateBuilder {
catalog_list: Some(existing.catalog_list),
table_functions: Some(existing.table_functions),
scalar_functions: Some(existing.scalar_functions.into_values().collect_vec()),
+ higher_order_functions: Some(
+ existing.higher_order_functions.into_values().collect_vec(),
+ ),
aggregate_functions: Some(
existing.aggregate_functions.into_values().collect_vec(),
),
@@ -1160,6 +1179,10 @@ impl SessionStateBuilder {
.get_or_insert_with(Vec::new)
.extend(SessionStateDefaults::default_scalar_functions());
+ self.higher_order_functions
+ .get_or_insert_with(Vec::new)
+ .extend(SessionStateDefaults::default_higher_order_functions());
+
self.aggregate_functions
.get_or_insert_with(Vec::new)
.extend(SessionStateDefaults::default_aggregate_functions());
@@ -1345,6 +1368,15 @@ impl SessionStateBuilder {
self
}
+ /// Set the map of [`HigherOrderUDF`]s
+ pub fn with_higher_order_functions(
+ mut self,
+ higher_order_functions: Vec<Arc<dyn HigherOrderUDF>>,
+ ) -> Self {
+ self.higher_order_functions = Some(higher_order_functions);
+ self
+ }
+
/// Set the map of [`AggregateUDF`]s
pub fn with_aggregate_functions(
mut self,
@@ -1518,6 +1550,7 @@ impl SessionStateBuilder {
catalog_list,
table_functions,
scalar_functions,
+ higher_order_functions,
aggregate_functions,
window_functions,
extension_types,
@@ -1556,6 +1589,7 @@ impl SessionStateBuilder {
}),
table_functions: table_functions.unwrap_or_default(),
scalar_functions: HashMap::new(),
+ higher_order_functions: HashMap::new(),
aggregate_functions: HashMap::new(),
window_functions: HashMap::new(),
extension_types: Arc::new(MemoryExtensionTypeRegistry::default()),
@@ -1611,6 +1645,29 @@ impl SessionStateBuilder {
}
}
+ if let Some(higher_order_functions) = higher_order_functions {
+ for function in higher_order_functions {
+ match state.register_higher_order_function(Arc::clone(&function)) {
+ Ok(Some(existing)) => {
+ debug!(
+ "Overwrote existing higher-order function '{}'",
+ existing.name()
+ );
+ }
+ Ok(None) => {
+ debug!("Registered higher-order function '{}'", function.name());
+ }
+ Err(err) => {
+ debug!(
+ "Failed to register higher-order function '{}': {}",
+ function.name(),
+ err
+ );
+ }
+ }
+ }
+ }
+
if let Some(aggregate_functions) = aggregate_functions {
aggregate_functions.into_iter().for_each(|udaf| {
let existing_udf = state.register_udaf(udaf);
@@ -1733,6 +1790,13 @@ impl SessionStateBuilder {
&mut self.scalar_functions
}
+ /// Returns the current scalar_functions value
+ pub fn higher_order_functions(
+ &mut self,
+ ) -> &mut Option<Vec<Arc<dyn HigherOrderUDF>>> {
+ &mut self.higher_order_functions
+ }
+
/// Returns the current aggregate_functions value
pub fn aggregate_functions(&mut self) -> &mut Option<Vec<Arc<AggregateUDF>>> {
&mut self.aggregate_functions
@@ -1841,6 +1905,7 @@ impl Debug for SessionStateBuilder {
.field("physical_optimizers", &self.physical_optimizers)
.field("table_functions", &self.table_functions)
.field("scalar_functions", &self.scalar_functions)
+ .field("higher_order_functions", &self.higher_order_functions)
.field("aggregate_functions", &self.aggregate_functions)
.field("window_functions", &self.window_functions)
.finish()
@@ -1951,6 +2016,10 @@ impl ContextProvider for SessionContextProvider<'_> {
self.state.scalar_functions().get(name).cloned()
}
+ fn get_higher_order_meta(&self, name: &str) -> Option<Arc<dyn HigherOrderUDF>> {
+ self.state.higher_order_functions().get(name).cloned()
+ }
+
fn get_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>> {
self.state.aggregate_functions().get(name).cloned()
}
@@ -1987,6 +2056,14 @@ impl ContextProvider for SessionContextProvider<'_> {
self.state.scalar_functions().keys().cloned().collect()
}
+ fn higher_order_function_names(&self) -> Vec<String> {
+ self.state
+ .higher_order_functions()
+ .keys()
+ .cloned()
+ .collect()
+ }
+
fn udaf_names(&self) -> Vec<String> {
self.state.aggregate_functions().keys().cloned().collect()
}
@@ -2026,6 +2103,16 @@ impl FunctionRegistry for SessionState {
})
}
+ fn higher_order_function(
+ &self,
+ name: &str,
+ ) -> datafusion_common::Result<Arc<dyn HigherOrderUDF>> {
+ self.higher_order_functions
+ .get(name)
+ .cloned()
+ .ok_or_else(|| plan_datafusion_err!("Higher Order Function {name} not found"))
+ }
+
fn udaf(&self, name: &str) -> datafusion_common::Result<Arc<AggregateUDF>> {
let result = self.aggregate_functions.get(name);
@@ -2053,6 +2140,19 @@ impl FunctionRegistry for SessionState {
Ok(self.scalar_functions.insert(udf.name().into(), udf))
}
+ fn register_higher_order_function(
+ &mut self,
+ function: Arc<dyn HigherOrderUDF>,
+ ) -> datafusion_common::Result<Option<Arc<dyn HigherOrderUDF>>> {
+ function.aliases().iter().for_each(|alias| {
+ self.higher_order_functions
+ .insert(alias.clone(), Arc::clone(&function));
+ });
+ Ok(self
+ .higher_order_functions
+ .insert(function.name().into(), function))
+ }
+
fn register_udaf(
&mut self,
udaf: Arc<AggregateUDF>,
@@ -2088,6 +2188,19 @@ impl FunctionRegistry for SessionState {
Ok(udf)
}
+ fn deregister_higher_order_function(
+ &mut self,
+ name: &str,
+ ) -> datafusion_common::Result<Option<Arc<dyn HigherOrderUDF>>> {
+ let function = self.higher_order_functions.remove(name);
+ if let Some(function) = &function {
+ for alias in function.aliases() {
+ self.higher_order_functions.remove(alias);
+ }
+ }
+ Ok(function)
+ }
+
fn deregister_udaf(
&mut self,
name: &str,
@@ -2134,6 +2247,10 @@ impl FunctionRegistry for SessionState {
Ok(())
}
+ fn higher_order_function_names(&self) -> HashSet<String> {
+ self.higher_order_functions.keys().cloned().collect()
+ }
+
fn udafs(&self) -> HashSet<String> {
self.aggregate_functions.keys().cloned().collect()
}
@@ -2176,6 +2293,7 @@ impl From<&SessionState> for TaskContext {
state.session_id.clone(),
state.config.clone(),
state.scalar_functions.clone(),
+ state.higher_order_functions.clone(),
state.aggregate_functions.clone(),
state.window_functions.clone(),
Arc::clone(&state.runtime_env),
@@ -2245,6 +2363,7 @@ mod tests {
use datafusion_common::config::Dialect;
use datafusion_execution::config::SessionConfig;
use datafusion_expr::Expr;
+ use datafusion_expr::HigherOrderUDF;
use datafusion_optimizer::Optimizer;
use datafusion_optimizer::optimizer::OptimizerRule;
use datafusion_physical_plan::display::DisplayableExecutionPlan;
@@ -2558,6 +2677,10 @@ mod tests {
self.state.scalar_functions().get(name).cloned()
}
+ fn get_higher_order_meta(&self, name: &str) -> Option<Arc<dyn HigherOrderUDF>> {
+ self.state.higher_order_functions().get(name).cloned()
+ }
+
fn get_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>> {
self.state.aggregate_functions().get(name).cloned()
}
@@ -2578,6 +2701,14 @@ mod tests {
self.state.scalar_functions().keys().cloned().collect()
}
+ fn higher_order_function_names(&self) -> Vec<String> {
+ self.state
+ .higher_order_functions()
+ .keys()
+ .cloned()
+ .collect()
+ }
+
fn udaf_names(&self) -> Vec<String> {
self.state.aggregate_functions().keys().cloned().collect()
}
diff --git a/datafusion/core/src/execution/session_state_defaults.rs b/datafusion/core/src/execution/session_state_defaults.rs
index 8ef041e5bf..5e85c1bbc5 100644
--- a/datafusion/core/src/execution/session_state_defaults.rs
+++ b/datafusion/core/src/execution/session_state_defaults.rs
@@ -37,7 +37,7 @@ use datafusion_execution::object_store::ObjectStoreUrl;
use datafusion_execution::runtime_env::RuntimeEnv;
use datafusion_expr::planner::ExprPlanner;
use datafusion_expr::registry::ExtensionTypeRegistrationRef;
-use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF};
+use datafusion_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF, WindowUDF};
use std::collections::HashMap;
use std::sync::Arc;
use url::Url;
@@ -113,6 +113,15 @@ impl SessionStateDefaults {
functions
}
+ /// returns the list of default [`HigherOrderUDF`]s
+ pub fn default_higher_order_functions() -> Vec<Arc<dyn HigherOrderUDF>> {
+ #[cfg(feature = "nested_expressions")]
+ return functions_nested::all_default_higher_order_functions();
+
+ #[cfg(not(feature = "nested_expressions"))]
+ return Vec::new();
+ }
+
/// returns the list of default [`AggregateUDF`]s
pub fn default_aggregate_functions() -> Vec<Arc<AggregateUDF>> {
functions_aggregate::all_default_aggregate_functions()
diff --git a/datafusion/core/tests/optimizer/mod.rs b/datafusion/core/tests/optimizer/mod.rs
index 72d50760f1..a461c6f6c5 100644
--- a/datafusion/core/tests/optimizer/mod.rs
+++ b/datafusion/core/tests/optimizer/mod.rs
@@ -30,8 +30,8 @@ use datafusion_common::tree_node::TransformedResult;
use datafusion_common::{DFSchema, Result, ScalarValue, TableReference, plan_err};
use datafusion_expr::interval_arithmetic::{Interval, NullableInterval};
use datafusion_expr::{
- AggregateUDF, BinaryExpr, Expr, ExprSchemable, LogicalPlan, Operator, ScalarUDF,
- TableSource, WindowUDF, col, lit,
+ AggregateUDF, BinaryExpr, Expr, ExprSchemable, HigherOrderUDF, LogicalPlan, Operator,
+ ScalarUDF, TableSource, WindowUDF, col, lit,
};
use datafusion_functions::core::expr_ext::FieldAccessor;
use datafusion_optimizer::analyzer::Analyzer;
@@ -216,6 +216,10 @@ impl ContextProvider for MyContextProvider {
self.udfs.get(name).cloned()
}
+ fn get_higher_order_meta(&self, _name: &str) -> Option<Arc<dyn HigherOrderUDF>> {
+ None
+ }
+
fn get_aggregate_meta(&self, _name: &str) -> Option<Arc<AggregateUDF>> {
None
}
@@ -236,6 +240,10 @@ impl ContextProvider for MyContextProvider {
Vec::new()
}
+ fn higher_order_function_names(&self) -> Vec<String> {
+ Vec::new()
+ }
+
fn udaf_names(&self) -> Vec<String> {
Vec::new()
}
diff --git a/datafusion/datasource-arrow/src/file_format.rs b/datafusion/datasource-arrow/src/file_format.rs
index f66275728c..9297486ad6 100644
--- a/datafusion/datasource-arrow/src/file_format.rs
+++ b/datafusion/datasource-arrow/src/file_format.rs
@@ -545,7 +545,9 @@ mod tests {
use datafusion_execution::runtime_env::RuntimeEnv;
use datafusion_expr::execution_props::ExecutionProps;
use datafusion_expr::registry::ExtensionTypeRegistryRef;
- use datafusion_expr::{AggregateUDF, Expr, LogicalPlan, ScalarUDF, WindowUDF};
+ use datafusion_expr::{
+ AggregateUDF, Expr, HigherOrderUDF, LogicalPlan, ScalarUDF, WindowUDF,
+ };
use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
use object_store::{chunked::ChunkedStore, memory::InMemory};
@@ -592,6 +594,10 @@ mod tests {
unimplemented!()
}
+ fn higher_order_functions(&self) -> &HashMap<String, Arc<dyn HigherOrderUDF>> {
+ unimplemented!()
+ }
+
fn aggregate_functions(&self) -> &HashMap<String, Arc<AggregateUDF>> {
unimplemented!()
}
diff --git a/datafusion/datasource/src/url.rs b/datafusion/datasource/src/url.rs
index 88f85eb456..4d7f5bf14c 100644
--- a/datafusion/datasource/src/url.rs
+++ b/datafusion/datasource/src/url.rs
@@ -518,7 +518,9 @@ mod tests {
use datafusion_execution::runtime_env::RuntimeEnv;
use datafusion_expr::execution_props::ExecutionProps;
use datafusion_expr::registry::ExtensionTypeRegistryRef;
- use datafusion_expr::{AggregateUDF, Expr, LogicalPlan, ScalarUDF, WindowUDF};
+ use datafusion_expr::{
+ AggregateUDF, Expr, HigherOrderUDF, LogicalPlan, ScalarUDF, WindowUDF,
+ };
use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
use datafusion_physical_plan::ExecutionPlan;
use object_store::{
@@ -1208,6 +1210,10 @@ mod tests {
unimplemented!()
}
+ fn higher_order_functions(&self) -> &HashMap<String, Arc<dyn HigherOrderUDF>> {
+ unimplemented!()
+ }
+
fn aggregate_functions(&self) -> &HashMap<String, Arc<AggregateUDF>> {
unimplemented!()
}
diff --git a/datafusion/execution/src/task.rs b/datafusion/execution/src/task.rs
index eadf3ea8ac..0de0c937f2 100644
--- a/datafusion/execution/src/task.rs
+++ b/datafusion/execution/src/task.rs
@@ -21,7 +21,7 @@ use crate::{
};
use datafusion_common::{Result, internal_datafusion_err, plan_datafusion_err};
use datafusion_expr::planner::ExprPlanner;
-use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF};
+use datafusion_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF, WindowUDF};
use std::collections::HashSet;
use std::{collections::HashMap, sync::Arc};
@@ -58,6 +58,8 @@ pub struct TaskContext {
session_config: SessionConfig,
/// Scalar functions associated with this task context
scalar_functions: HashMap<String, Arc<ScalarUDF>>,
+ /// Higher order functions associated with this task context
+ higher_order_functions: HashMap<String, Arc<dyn HigherOrderUDF>>,
/// Aggregate functions associated with this task context
aggregate_functions: HashMap<String, Arc<AggregateUDF>>,
/// Window functions associated with this task context
@@ -76,6 +78,7 @@ impl Default for TaskContext {
task_id: None,
session_config: SessionConfig::new(),
scalar_functions: HashMap::new(),
+ higher_order_functions: HashMap::new(),
aggregate_functions: HashMap::new(),
window_functions: HashMap::new(),
runtime,
@@ -89,11 +92,13 @@ impl TaskContext {
/// Most users will use [`SessionContext::task_ctx`] to create [`TaskContext`]s
///
/// [`SessionContext::task_ctx`]: https://docs.rs/datafusion/latest/datafusion/execution/context/struct.SessionContext.html#method.task_ctx
+ #[expect(clippy::too_many_arguments)]
pub fn new(
task_id: Option<String>,
session_id: String,
session_config: SessionConfig,
scalar_functions: HashMap<String, Arc<ScalarUDF>>,
+ higher_order_functions: HashMap<String, Arc<dyn HigherOrderUDF>>,
aggregate_functions: HashMap<String, Arc<AggregateUDF>>,
window_functions: HashMap<String, Arc<WindowUDF>>,
runtime: Arc<RuntimeEnv>,
@@ -103,6 +108,7 @@ impl TaskContext {
session_id,
session_config,
scalar_functions,
+ higher_order_functions,
aggregate_functions,
window_functions,
runtime,
@@ -138,6 +144,10 @@ impl TaskContext {
&self.scalar_functions
}
+ pub fn higher_order_functions(&self) -> &HashMap<String, Arc<dyn HigherOrderUDF>> {
+ &self.higher_order_functions
+ }
+
pub fn aggregate_functions(&self) -> &HashMap<String, Arc<AggregateUDF>> {
&self.aggregate_functions
}
@@ -172,6 +182,16 @@ impl FunctionRegistry for TaskContext {
})
}
+ fn higher_order_function(&self, name: &str) -> Result<Arc<dyn HigherOrderUDF>> {
+ let result = self.higher_order_functions.get(name);
+
+ result.cloned().ok_or_else(|| {
+ plan_datafusion_err!(
+ "There is no higher-order function named \"{name}\" in the TaskContext"
+ )
+ })
+ }
+
fn udaf(&self, name: &str) -> Result<Arc<AggregateUDF>> {
let result = self.aggregate_functions.get(name);
@@ -214,10 +234,27 @@ impl FunctionRegistry for TaskContext {
Ok(self.scalar_functions.insert(udf.name().into(), udf))
}
+ fn register_higher_order_function(
+ &mut self,
+ function: Arc<dyn HigherOrderUDF>,
+ ) -> Result<Option<Arc<dyn HigherOrderUDF>>> {
+ function.aliases().iter().for_each(|alias| {
+ self.higher_order_functions
+ .insert(alias.clone(), Arc::clone(&function));
+ });
+ Ok(self
+ .higher_order_functions
+ .insert(function.name().into(), function))
+ }
+
fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>> {
vec![]
}
+ fn higher_order_function_names(&self) -> HashSet<String> {
+ self.higher_order_functions.keys().cloned().collect()
+ }
+
fn udafs(&self) -> HashSet<String> {
self.aggregate_functions.keys().cloned().collect()
}
@@ -269,6 +306,7 @@ mod tests {
HashMap::default(),
HashMap::default(),
HashMap::default(),
+ HashMap::default(),
runtime,
);
@@ -301,6 +339,7 @@ mod tests {
HashMap::default(),
HashMap::default(),
HashMap::default(),
+ HashMap::default(),
runtime,
);
diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs
index c3dffe553d..582e8e41dd 100644
--- a/datafusion/expr/src/expr.rs
+++ b/datafusion/expr/src/expr.rs
@@ -26,8 +26,10 @@ use std::sync::Arc;
use crate::expr_fn::binary_expr;
use crate::function::WindowFunctionSimplification;
+use crate::higher_order_function::HigherOrderUDF;
use crate::logical_plan::Subquery;
-use crate::{AggregateUDF, Volatility};
+use crate::type_coercion::functions::value_fields_with_higher_order_udf;
+use crate::{AggregateUDF, LambdaParametersProgress, ValueOrLambda, Volatility};
use crate::{ExprSchemable, Operator, Signature, WindowFrame, WindowUDF};
use arrow::datatypes::{DataType, Field, FieldRef};
@@ -38,7 +40,8 @@ use datafusion_common::tree_node::{
Transformed, TransformedResult, TreeNode, TreeNodeContainer, TreeNodeRecursion,
};
use datafusion_common::{
- Column, DFSchema, HashMap, Result, ScalarValue, Spans, TableReference,
+ Column, DFSchema, ExprSchema, HashMap, Result, ScalarValue, Spans, TableReference,
+ plan_err,
};
use datafusion_expr_common::placement::ExpressionPlacement;
use datafusion_functions_window_common::field::WindowUDFFieldArgs;
@@ -404,6 +407,112 @@ pub enum Expr {
OuterReferenceColumn(FieldRef, Column),
/// Unnest expression
Unnest(Unnest),
+ /// Call a higher order function with a set of arguments.
+ ///
+ /// For example, `array_transform([1,2,3], v -> v+1)` would be equivalent to:
+ ///
+ /// ```text
+ /// HigherOrderFunction(array_transform)
+ /// ├── args[0]: Literal([1,2,3])
+ /// └── args[1]: Lambda
+ /// ├── params: ["v"]
+ /// └── body: BinaryExpr(+)
+ /// ├── LambdaVariable("v")
+ /// └── Literal(1)
+ /// ```
+ HigherOrderFunction(HigherOrderFunction),
+ /// A Lambda expression with a set of parameters names and a body
+ Lambda(Lambda),
+ /// A named reference to a lambda parameter
+ LambdaVariable(LambdaVariable),
+}
+
+/// Invoke a [`HigherOrderUDF`] with a set of arguments
+#[derive(Clone, Eq, PartialOrd, Debug)]
+pub struct HigherOrderFunction {
+ /// The function
+ pub func: Arc<dyn HigherOrderUDF>,
+ /// List of expressions to feed to the functions as arguments
+ pub args: Vec<Expr>,
+}
+
+impl HigherOrderFunction {
+ /// Create a new `HigherOrderFunction` from a [`HigherOrderUDF`]
+ pub fn new(func: Arc<dyn HigherOrderUDF>, args: Vec<Expr>) -> Self {
+ Self { func, args }
+ }
+
+ pub fn name(&self) -> &str {
+ self.func.name()
+ }
+
+ /// Invokes the inner function [`HigherOrderUDF::lambda_parameters`]
+ /// using the arguments of this invocation
+ pub fn lambda_parameters(
+ &self,
+ schema: &dyn ExprSchema,
+ ) -> Result<Vec<Vec<FieldRef>>> {
+ let args = self
+ .args
+ .iter()
+ .map(|e| match e {
+ Expr::Lambda(lambda) => {
+ Ok(ValueOrLambda::Lambda(Some(lambda.body.to_field(schema)?.1)))
+ }
+ _ => Ok(ValueOrLambda::Value(e.to_field(schema)?.1)),
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ let coerced_fields =
+ value_fields_with_higher_order_udf(&args, self.func.as_ref())?;
+
+ match self.func.lambda_parameters(0, &coerced_fields)? {
+ LambdaParametersProgress::Partial(_) => plan_err!(
+ "{} lambda_parameters returned a partial result when the return type of all it's lambdas were provided",
+ self.name()
+ ),
+ LambdaParametersProgress::Complete(items) => Ok(items),
+ }
+ }
+}
+
+impl Hash for HigherOrderFunction {
+ fn hash<H: Hasher>(&self, state: &mut H) {
+ self.func.hash(state);
+ self.args.hash(state);
+ }
+}
+
+impl PartialEq for HigherOrderFunction {
+ fn eq(&self, other: &Self) -> bool {
+ self.func.as_ref() == other.func.as_ref() && self.args == other.args
+ }
+}
+
+/// A named reference to a lambda parameter which includes it's own [`FieldRef`],
+/// which is used to implement [`ExprSchemable`], for example. Note the field must
+/// be set in order to create a physical lambda variable. A helper to automatically
+/// set them will be added in the future
+#[derive(Clone, PartialEq, PartialOrd, Eq, Debug, Hash)]
+pub struct LambdaVariable {
+ pub name: String,
+ pub field: Option<FieldRef>,
+ pub spans: Spans,
+}
+
+impl LambdaVariable {
+ /// Create a lambda variable from a name and a Field.
+ pub fn new(name: String, field: Option<FieldRef>) -> Self {
+ Self {
+ name,
+ field,
+ spans: Spans::new(),
+ }
+ }
+
+ pub fn spans_mut(&mut self) -> &mut Spans {
+ &mut self.spans
+ }
}
impl Default for Expr {
@@ -1288,6 +1397,25 @@ impl GroupingSet {
}
}
+/// A Lambda expression with a set of parameters names and a body
+#[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)]
+pub struct Lambda {
+ /// The parameters names
+ pub params: Vec<String>,
+ /// The body expression
+ 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 {
@@ -1620,6 +1748,9 @@ impl Expr {
#[expect(deprecated)]
Expr::Wildcard { .. } => "Wildcard",
Expr::Unnest { .. } => "Unnest",
+ Expr::HigherOrderFunction { .. } => "HigherOrderFunction",
+ Expr::Lambda { .. } => "Lambda",
+ Expr::LambdaVariable { .. } => "LambdaVariable",
}
}
@@ -2141,6 +2272,9 @@ impl Expr {
pub fn short_circuits(&self) -> bool {
match self {
Expr::ScalarFunction(ScalarFunction { func, .. }) => func.short_circuits(),
+ Expr::HigherOrderFunction(HigherOrderFunction { func, .. }) => {
+ func.short_circuits()
+ }
Expr::BinaryExpr(BinaryExpr { op, .. }) => {
matches!(op, Operator::And | Operator::Or)
}
@@ -2180,7 +2314,9 @@ impl Expr {
| Expr::Wildcard { .. }
| Expr::WindowFunction(..)
| Expr::Literal(..)
- | Expr::Placeholder(..) => false,
+ | Expr::Placeholder(..)
+ | Expr::Lambda(..)
+ | Expr::LambdaVariable(..) => false,
}
}
@@ -2780,6 +2916,20 @@ impl HashNode for Expr {
column.hash(state);
}
Expr::Unnest(Unnest { expr: _expr }) => {}
+ Expr::HigherOrderFunction(HigherOrderFunction { func, args: _args }) => {
+ func.hash(state);
+ }
+ Expr::Lambda(Lambda { params, body: _ }) => {
+ params.hash(state);
+ }
+ Expr::LambdaVariable(LambdaVariable {
+ name,
+ field,
+ spans: _,
+ }) => {
+ name.hash(state);
+ field.hash(state);
+ }
};
}
}
@@ -3098,6 +3248,25 @@ impl Display for SchemaDisplay<'_> {
}
}
}
+ Expr::HigherOrderFunction(HigherOrderFunction { func, args }) => {
+ match func.schema_name(args) {
+ Ok(name) => {
+ write!(f, "{name}")
+ }
+ Err(e) => {
+ write!(f, "got error from schema_name {e}")
+ }
+ }
+ }
+ Expr::Lambda(Lambda { params, body }) => {
+ write!(
+ f,
+ "({}) -> {}",
+ display_comma_separated(params),
+ SchemaDisplay(body)
+ )
+ }
+ Expr::LambdaVariable(c) => f.write_str(&c.name),
}
}
}
@@ -3278,6 +3447,9 @@ impl Display for SqlDisplay<'_> {
}
}
}
+ Expr::Lambda(Lambda { params, body }) => {
+ write!(f, "({}) -> {}", params.join(", "), SchemaDisplay(body))
+ }
_ => write!(f, "{}", self.0),
}
}
@@ -3595,6 +3767,13 @@ impl Display for Expr {
Expr::Unnest(Unnest { expr }) => {
write!(f, "{UNNEST_COLUMN_PREFIX}({expr})")
}
+ Expr::HigherOrderFunction(fun) => {
+ fmt_function(f, fun.name(), false, &fun.args, true)
+ }
+ Expr::Lambda(Lambda { params, body }) => {
+ write!(f, "({}) -> {body}", params.join(", "))
+ }
+ Expr::LambdaVariable(c) => f.write_str(&c.name),
}
}
}
diff --git a/datafusion/expr/src/expr_fn.rs b/datafusion/expr/src/expr_fn.rs
index ff90180d70..cec6b7ec05 100644
--- a/datafusion/expr/src/expr_fn.rs
+++ b/datafusion/expr/src/expr_fn.rs
@@ -18,8 +18,9 @@
//! Functions for creating logical expressions
use crate::expr::{
- AggregateFunction, BinaryExpr, Cast, Exists, GroupingSet, InList, InSubquery,
- NullTreatment, Placeholder, TryCast, Unnest, WildcardOptions, WindowFunction,
+ AggregateFunction, BinaryExpr, Cast, Exists, GroupingSet, InList, InSubquery, Lambda,
+ LambdaVariable, NullTreatment, Placeholder, TryCast, Unnest, WildcardOptions,
+ WindowFunction,
};
use crate::function::{
AccumulatorArgs, AccumulatorFactoryFunction, PartitionEvaluatorFactory,
@@ -719,6 +720,21 @@ pub fn interval_month_day_nano_lit(value: &str) -> Expr {
Expr::Literal(ScalarValue::IntervalMonthDayNano(interval), None)
}
+/// Create a lambda expression
+pub fn lambda(params: impl IntoIterator<Item = impl Into<String>>, body: Expr) -> Expr {
+ Expr::Lambda(Lambda::new(
+ params.into_iter().map(Into::into).collect(),
+ body,
+ ))
+}
+
+/// Create a lambda variable expression
+// todo: make this pub when support for optional field lands
+#[expect(unused)]
+fn lambda_var(name: impl Into<String>) -> Expr {
+ Expr::LambdaVariable(LambdaVariable::new(name.into(), None))
+}
+
/// Extensions for configuring [`Expr::AggregateFunction`] or [`Expr::WindowFunction`]
///
/// Adds methods to [`Expr`] that make it easy to set optional options
diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs
index b613db4b6a..c989bab304 100644
--- a/datafusion/expr/src/expr_schema.rs
+++ b/datafusion/expr/src/expr_schema.rs
@@ -16,18 +16,22 @@
// under the License.
use super::{Between, Expr, Like, predicate_bounds};
+use crate::ValueOrLambda;
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, LambdaVariable};
+use crate::higher_order_function::HigherOrderReturnFieldArgs;
+use crate::type_coercion::functions::value_fields_with_higher_order_udf_and_lambdas;
use crate::type_coercion::functions::{UDFCoercionExt, fields_with_udf};
use crate::udf::ReturnFieldArgs;
use crate::{LogicalPlan, Projection, Subquery, WindowFunctionDefinition, utils};
use arrow::compute::can_cast_types;
-use arrow::datatypes::{DataType, Field, FieldRef};
+use arrow::datatypes::FieldRef;
+use arrow::datatypes::{DataType, Field};
use datafusion_common::datatype::FieldExt;
-use datafusion_common::metadata::FieldMetadata;
use datafusion_common::{
Column, DataFusionError, ExprSchema, Result, ScalarValue, Spans, TableReference,
not_impl_err, plan_datafusion_err, plan_err,
@@ -217,6 +221,16 @@ impl ExprSchemable for Expr {
// Grouping sets do not really have a type and do not appear in projections
Ok(DataType::Null)
}
+ Expr::HigherOrderFunction(_func) => {
+ Ok(self.to_field(schema)?.1.data_type().clone())
+ }
+ Expr::Lambda(_lambda) => Ok(DataType::Null),
+ Expr::LambdaVariable(LambdaVariable { field, .. }) => match field {
+ Some(f) => Ok(f.data_type().clone()),
+ // If the lambda variable's field hasn't been specified, treat it as
+ // null (unspecified lambda variables generate an error during planning)
+ None => Ok(DataType::Null),
+ },
}
}
@@ -370,6 +384,16 @@ impl ExprSchemable for Expr {
// in projections
Ok(true)
}
+ Expr::HigherOrderFunction(_func) => {
+ Ok(self.to_field(input_schema)?.1.is_nullable())
+ }
+ Expr::Lambda(_lambda) => Ok(true),
+ Expr::LambdaVariable(LambdaVariable { field, .. }) => match field {
+ Some(f) => Ok(f.is_nullable()),
+ // If the lambda variable's field hasn't been specified, treat it as
+ // null (unspecified lambda variables generate an error during planning)
+ None => Ok(true),
+ },
}
}
@@ -585,6 +609,9 @@ impl ExprSchemable for Expr {
cast_output_field(&src, field.data_type(), true)
})
}
+ Expr::LambdaVariable(LambdaVariable {
+ field: Some(field), ..
+ }) => Ok(Arc::clone(field).renamed(&schema_name)),
Expr::Like(_)
| Expr::SimilarTo(_)
| Expr::Not(_)
@@ -596,11 +623,51 @@ impl ExprSchemable for Expr {
| Expr::Wildcard { .. }
| Expr::GroupingSet(_)
| Expr::Placeholder(_)
- | Expr::Unnest(_) => Ok(Arc::new(Field::new(
+ | Expr::Unnest(_)
+ | Expr::Lambda(_)
+ | Expr::LambdaVariable(_) => Ok(Arc::new(Field::new(
&schema_name,
self.get_type(schema)?,
self.nullable(schema)?,
))),
+ Expr::HigherOrderFunction(func) => {
+ let arg_fields = func
+ .args
+ .iter()
+ .map(|arg| match arg {
+ Expr::Lambda(Lambda { params: _, body }) => {
+ // use the name of the lambda instead of just the body to help with debugging
+ Ok(ValueOrLambda::Lambda(Arc::new(Field::new(
+ arg.qualified_name().1,
+ body.get_type(schema)?,
+ body.nullable(schema)?,
+ ))))
+ }
+ _ => Ok(ValueOrLambda::Value(arg.to_field(schema)?.1)),
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ let new_fields = value_fields_with_higher_order_udf_and_lambdas(
+ &arg_fields,
+ func.func.as_ref(),
+ )?;
+
+ let arguments = func
+ .args
+ .iter()
+ .map(|e| match e {
+ Expr::Literal(sv, _) => Some(sv),
+ _ => None,
+ })
+ .collect::<Vec<_>>();
+
+ let args = HigherOrderReturnFieldArgs {
+ arg_fields: &new_fields,
+ scalar_arguments: &arguments,
+ };
+
+ func.func.return_field_from_args(args)
+ }
}?;
Ok((
diff --git a/datafusion/expr/src/higher_order_function.rs b/datafusion/expr/src/higher_order_function.rs
new file mode 100644
index 0000000000..40eda950ec
--- /dev/null
+++ b/datafusion/expr/src/higher_order_function.rs
@@ -0,0 +1,772 @@
+// 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.
+
+//! [`HigherOrderUDF`]: User Defined Higher Order Functions
+
+use crate::expr::schema_name_from_exprs_comma_separated_without_space;
+use crate::{ColumnarValue, Documentation, Expr};
+use arrow::array::{ArrayRef, RecordBatch};
+use arrow::datatypes::{DataType, FieldRef, Schema};
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::{Result, ScalarValue, not_impl_err};
+use datafusion_expr_common::dyn_eq::{DynEq, DynHash};
+use datafusion_expr_common::signature::Volatility;
+use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
+use std::any::Any;
+use std::cmp::Ordering;
+use std::fmt::Debug;
+use std::hash::{Hash, Hasher};
+use std::sync::Arc;
+
+/// The types of arguments for which a function has implementations.
+///
+/// [`HigherOrderTypeSignature`] **DOES NOT** define the types that a user query could call the
+/// function with. DataFusion will automatically coerce (cast) argument types to
+/// one of the supported function signatures, if possible.
+///
+/// # Overview
+/// Functions typically provide implementations for a small number of different
+/// argument [`DataType`]s, rather than all possible combinations. If a user
+/// calls a function with arguments that do not match any of the declared types,
+/// DataFusion will attempt to automatically coerce (add casts to) function
+/// arguments so they match the [`HigherOrderTypeSignature`]. See the [`type_coercion`] module
+/// for more details
+///
+/// [`type_coercion`]: crate::type_coercion
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
+pub enum HigherOrderTypeSignature {
+ /// The acceptable signature and coercions rules are special for this
+ /// function.
+ ///
+ /// If this signature is specified,
+ /// DataFusion will call [`HigherOrderUDF::coerce_value_types`] to prepare argument types.
+ UserDefined,
+ /// One or more lambdas or arguments with arbitrary types
+ VariadicAny,
+ /// The specified number of lambdas or arguments with arbitrary types.
+ Any(usize),
+}
+
+/// Provides information necessary for calling a higher order function.
+///
+/// - [`HigherOrderTypeSignature`] defines the argument types that a function has implementations
+/// for.
+///
+/// - [`Volatility`] defines how the output of the function changes with the input.
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
+pub struct HigherOrderSignature {
+ /// The data types that the function accepts. See [HigherOrderTypeSignature] for more information.
+ pub type_signature: HigherOrderTypeSignature,
+ /// The volatility of the function. See [Volatility] for more information.
+ pub volatility: Volatility,
+ /// Whether [HigherOrderUDF::coerce_values_for_lambdas] should be called
+ pub coerce_values_for_lambdas: bool,
+ /// The max number of times to call [HigherOrderUDF::lambda_parameters] before raising an error.
+ /// Used to guard against implementations that causes an infinite loop by endlessly returning
+ /// [LambdaParametersProgress::Partial]. Defaults to 256
+ pub lambda_parameters_max_iterations: usize,
+}
+
+const LAMBDA_PARAMETERS_MAX_ITERATIONS: usize = 256;
+
+impl HigherOrderSignature {
+ /// Creates a new `HigherOrderSignature` from a given type signature and volatility.
+ pub fn new(type_signature: HigherOrderTypeSignature, volatility: Volatility) -> Self {
+ HigherOrderSignature {
+ type_signature,
+ volatility,
+ coerce_values_for_lambdas: false,
+ lambda_parameters_max_iterations: LAMBDA_PARAMETERS_MAX_ITERATIONS,
+ }
+ }
+
+ /// User-defined coercion rules for the function.
+ pub fn user_defined(volatility: Volatility) -> Self {
+ Self {
+ type_signature: HigherOrderTypeSignature::UserDefined,
+ volatility,
+ coerce_values_for_lambdas: false,
+ lambda_parameters_max_iterations: LAMBDA_PARAMETERS_MAX_ITERATIONS,
+ }
+ }
+
+ /// An arbitrary number of lambdas or arguments of any type.
+ pub fn variadic_any(volatility: Volatility) -> Self {
+ Self {
+ type_signature: HigherOrderTypeSignature::VariadicAny,
+ volatility,
+ coerce_values_for_lambdas: false,
+ lambda_parameters_max_iterations: LAMBDA_PARAMETERS_MAX_ITERATIONS,
+ }
+ }
+
+ /// A specified number of arguments of any type
+ pub fn any(arg_count: usize, volatility: Volatility) -> Self {
+ Self {
+ type_signature: HigherOrderTypeSignature::Any(arg_count),
+ volatility,
+ coerce_values_for_lambdas: false,
+ lambda_parameters_max_iterations: LAMBDA_PARAMETERS_MAX_ITERATIONS,
+ }
+ }
+
+ /// Set [Self::coerce_values_for_lambdas] to true to indicate that [HigherOrderUDF::coerce_values_for_lambdas]
+ /// should be called
+ pub fn with_coerce_values_for_lambdas(mut self) -> Self {
+ self.coerce_values_for_lambdas = true;
+
+ self
+ }
+}
+
+impl PartialEq for dyn HigherOrderUDF {
+ fn eq(&self, other: &Self) -> bool {
+ self.dyn_eq(other as _)
+ }
+}
+
+impl PartialOrd for dyn HigherOrderUDF {
+ fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+ let mut cmp = self.name().cmp(other.name());
+ if cmp == Ordering::Equal {
+ cmp = self.signature().partial_cmp(other.signature())?;
+ }
+ if cmp == Ordering::Equal {
+ cmp = self.aliases().partial_cmp(other.aliases())?;
+ }
+ // Contract for PartialOrd and PartialEq consistency requires that
+ // a == b if and only if partial_cmp(a, b) == Some(Equal).
+ if cmp == Ordering::Equal && self != other {
+ // Functions may have other properties besides name and signature
+ // that differentiate two instances (e.g. type, or arbitrary parameters).
+ // We cannot return Some(Equal) in such case.
+ return None;
+ }
+ debug_assert!(
+ cmp == Ordering::Equal || self != other,
+ "Detected incorrect implementation of PartialEq when comparing functions: '{}' and '{}'. \
+ The functions compare as equal, but they are not equal based on general properties that \
+ the PartialOrd implementation observes,",
+ self.name(),
+ other.name()
+ );
+ Some(cmp)
+ }
+}
+
+impl Eq for dyn HigherOrderUDF {}
+
+impl Hash for dyn HigherOrderUDF {
+ fn hash<H: Hasher>(&self, state: &mut H) {
+ self.dyn_hash(state)
+ }
+}
+
+/// Arguments passed to [`HigherOrderUDF::invoke_with_args`] when invoking a
+/// higher order function.
+#[derive(Debug, Clone)]
+pub struct HigherOrderFunctionArgs {
+ /// The evaluated arguments and lambdas to the function
+ pub args: Vec<ValueOrLambda<ColumnarValue, LambdaArgument>>,
+ /// Field associated with each arg, if it exists
+ /// For lambdas, it will be the field of the result of
+ /// the lambda if evaluated with the parameters
+ /// returned from [`HigherOrderUDF::lambda_parameters`]
+ pub arg_fields: Vec<ValueOrLambda<FieldRef, FieldRef>>,
+ /// The number of rows in record batch being evaluated
+ pub number_rows: usize,
+ /// The return field of the higher order function returned
+ /// (from `return_field_from_args`) when creating the
+ /// physical expression from the logical expression
+ pub return_field: FieldRef,
+ /// The config options at execution time
+ pub config_options: Arc<ConfigOptions>,
+}
+
+impl HigherOrderFunctionArgs {
+ /// The return type of the function. See [`Self::return_field`] for more
+ /// details.
+ pub fn return_type(&self) -> &DataType {
+ self.return_field.data_type()
+ }
+}
+
+/// A lambda argument to a HigherOrderFunction
+#[derive(Clone, Debug)]
+pub struct LambdaArgument {
+ /// The parameters defined in this lambda
+ ///
+ /// For example, for `array_transform([2], v -> -v)`,
+ /// this will be `vec![Field::new("v", DataType::Int32, true)]`
+ params: Vec<FieldRef>,
+ /// The body of the lambda
+ ///
+ /// For example, for `array_transform([2], v -> -v)`,
+ /// this will be the physical expression of `-v`
+ body: Arc<dyn PhysicalExpr>,
+}
+
+impl LambdaArgument {
+ pub fn new(params: Vec<FieldRef>, body: Arc<dyn PhysicalExpr>) -> Self {
+ Self { params, body }
+ }
+
+ /// Evaluate this lambda
+ /// `args` should evaluate to the value of each parameter
+ /// of the correspondent lambda returned in [HigherOrderUDF::lambda_parameters].
+ pub fn evaluate(
+ &self,
+ args: &[&dyn Fn() -> Result<ArrayRef>],
+ ) -> Result<ColumnarValue> {
+ let columns = args
+ .iter()
+ .take(self.params.len())
+ .map(|arg| arg())
+ .collect::<Result<_>>()?;
+
+ let schema = Arc::new(Schema::new(self.params.clone()));
+
+ let batch = RecordBatch::try_new(schema, columns)?;
+
+ self.body.evaluate(&batch)
+ }
+}
+
+/// Information about arguments passed to the function
+///
+/// This structure contains metadata about how the function was called
+/// such as the type of the arguments, any scalar arguments and if the
+/// arguments can (ever) be null
+///
+/// See [`HigherOrderUDF::return_field_from_args`] for more information
+#[derive(Clone, Debug)]
+pub struct HigherOrderReturnFieldArgs<'a> {
+ /// The data types of the arguments to the function
+ ///
+ /// If argument `i` to the function is a lambda, it will be the field of the result of the
+ /// lambda if evaluated with the parameters returned from [`HigherOrderUDF::lambda_parameters`]
+ ///
+ /// For example, with `array_transform([1], v -> v == 5)`
+ /// this field will be
+ /// ```ignore
+ /// [
+ /// ValueOrLambda::Value(Field::new("", DataType::new_list(DataType::Int32, true), true)),
+ /// ValueOrLambda::Lambda(Field::new("", DataType::Boolean, true))
+ /// ]
+ /// ```
+ pub arg_fields: &'a [ValueOrLambda<FieldRef, FieldRef>],
+ /// Is argument `i` to the function a scalar (constant)?
+ ///
+ /// If the argument `i` is not a scalar, it will be None
+ ///
+ /// For example, if a function is called like `array_transform([1], v -> v == 5)`
+ /// this field will be `[Some(ScalarValue::List(...), None]`
+ pub scalar_arguments: &'a [Option<&'a ScalarValue>],
+}
+
+/// An argument to a higher order function
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub enum ValueOrLambda<V, L> {
+ /// A value with associated data
+ Value(V),
+ /// A lambda with associated data
+ Lambda(L),
+}
+
+/// Represents a step during the resolution of the parameters of all lambdas of a given
+/// higher-order function via [HigherOrderUDF::lambda_parameters]. It's valid that the
+/// fields of a given lambda changes between steps, and is up to the implementation to
+/// provide during the function evaluation the parameters that matches the fields returned
+/// at the [LambdaParametersProgress::Complete] step. See [HigherOrderUDF::lambda_parameters]
+/// docs for more details
+pub enum LambdaParametersProgress {
+ /// The parameters of some lambdas are unknown due to a dependency on another lambda output field
+ /// or are placeholders due to a dependency on it's own output field. It's perfectly valid to
+ /// contain only `Some`'s and not a single `None`, representing lambdas that depends only on itself
+ /// and not on others. [HigherOrderUDF::lambda_parameters] will be called again with the output
+ /// field of all lambdas with known parameters.
+ Partial(Vec<Option<Vec<FieldRef>>>),
+ /// There are no unmet dependencies and all parameters are known, [HigherOrderUDF::lambda_parameters]
+ /// will not be called again
+ Complete(Vec<Vec<FieldRef>>),
+}
+
+/// Trait for implementing user defined higher order functions.
+///
+/// This trait exposes the full API for implementing user defined functions and
+/// can be used to implement any function.
+///
+/// See [`array_transform.rs`] for a commented complete implementation
+///
+/// [`array_transform.rs`]: https://github.com/apache/datafusion/blob/main/datafusion/functions-nested/src/array_transform.rs
+pub trait HigherOrderUDF: Debug + DynEq + DynHash + Send + Sync + Any {
+ /// Returns this function's name
+ fn name(&self) -> &str;
+
+ /// Returns any aliases (alternate names) for this function.
+ ///
+ /// Aliases can be used to invoke the same function using different names.
+ /// For example in some databases `now()` and `current_timestamp()` are
+ /// aliases for the same function. This behavior can be obtained by
+ /// returning `current_timestamp` as an alias for the `now` function.
+ ///
+ /// Note: `aliases` should only include names other than [`Self::name`].
+ /// Defaults to `[]` (no aliases)
+ fn aliases(&self) -> &[String] {
+ &[]
+ }
+
+ /// Returns the name of the column this expression would create
+ ///
+ /// See [`Expr::schema_name`] for details
+ fn schema_name(&self, args: &[Expr]) -> Result<String> {
+ Ok(format!(
+ "{}({})",
+ self.name(),
+ schema_name_from_exprs_comma_separated_without_space(args)?
+ ))
+ }
+
+ /// Returns a [`HigherOrderSignature`] describing the argument types for which this
+ /// function has an implementation, and the function's [`Volatility`].
+ ///
+ /// See [`HigherOrderSignature`] for more details on argument type handling
+ /// and [`Self::return_field_from_args`] for computing the return type.
+ ///
+ /// [`Volatility`]: datafusion_expr_common::signature::Volatility
+ fn signature(&self) -> &HigherOrderSignature;
+
+ /// Return the field of all the parameters supported by the lambdas in `fields`.
+ /// If a lambda support multiple parameters, all should be returned, regardless of
+ /// whether they are used or not on a particular invocation
+ ///
+ /// Tip: If you have a [`HigherOrderFunction`] invocation, you can call the helper
+ /// [`HigherOrderFunction::lambda_parameters`] instead of this method directly
+ ///
+ /// The name of the returned fields are ignored.
+ ///
+ /// This function is repeatedelly called until [LambdaParametersProgress::Complete] is returned, with
+ /// `step` increased by one at each invocation, starting at 0.
+ ///
+ /// For functions which all lambda parameters depend only on the field of it's value arguments,
+ /// this can return [LambdaParametersProgress::Complete] at step 0. Taking as an example a strict
+ /// array_reduce with the signature `(arr: [V], initial_value: I, (I, V) -> I, (I) -> O) -> O`, which
+ /// requires it's initial value to be the exact same type of it's merge output, which is also the
+ /// parameter of it's finish lambda, the expression
+ ///
+ /// `array_reduce([1.2, 2.1], 0.0, (acc, v) -> acc + v + 1.5, v -> v > 5.1)`
+ ///
+ /// would result in this function being called as the following:
+ ///
+ /// ```ignore
+ /// let lambda_parameters = array_reduce.lambda_parameters(
+ /// 0,
+ /// &[
+ /// // the Field of the literal `[1.2, 2.1]`, the array being reduced
+ /// ValueOrLambda::Value(Arc::new(Field::new("", DataType::new_list(DataType::Float32, true), true))),
+ /// // the Field of the literal `0.0`, the initial value
+ /// ValueOrLambda::Value(Arc::new(Field::new("", DataType::Float32, true))),
+ /// // the Field of the output of the merge lambda, which is unknown at this point because it depends
+ /// // on the return of this call
+ /// ValueOrLambda::Lambda(None),
+ /// // the Field of the output of the finish lambda, unknown for the same reason as above
+ /// ValueOrLambda::Lambda(None),
+ /// ])?;
+ ///
+ /// assert_eq!(
+ /// lambda_parameters,
+ /// LambdaParametersProgress::Complete(vec![
+ /// // the finish lambda supported parameters, regardless of how many are actually used
+ /// vec![
+ /// // the accumulator which is the field of the initial value
+ /// Arc::new(Field::new("ignored_name", DataType::Float32, true)),
+ /// // the array values being reduced
+ /// Arc::new(Field::new("", DataType::Float32, true)),
+ /// ],
+ /// // the merge lambda supported parameters
+ /// vec![
+ /// // the reduced value which is the field of the initial value
+ /// Arc::new(Field::new("ignored_name", DataType::Float32, true)),
+ /// ],
+ /// ])
+ /// );
+ /// ```
+ ///
+ /// For functions which lambda parameters depends on the output of other lambdas, or on their own lambda,
+ /// this can return [LambdaParametersProgress::Partial] until all dependencies are met. Note that for
+ /// lambda with cyclic dependencies, you likely want to use [HigherOrderUDF::coerce_values_for_lambdas] too.
+ /// Take as an example a flexible array_reduce with the signature `(arr: [V], initial_value: I, (ACC, V) -> ACC, (ACC) -> O) -> O`.
+ /// It has a cyclic dependency in the merge lambda, and a dependency of the finish lambda in the merge lambda,
+ /// and only requires the initial value to be *coercible* to the output of the merge lambda, which is defined by
+ /// it's [HigherOrderUDF::coerce_values_for_lambdas] implementation. The expression
+ ///
+ /// `array_reduce([1.2, 2.1], 0, (acc, v) -> acc + v + 1.5, v -> v > 5.1)`
+ ///
+ /// would result in this function being called as the following:
+ ///
+ /// ```ignore
+ /// let lambda_parameters = array_reduce.lambda_parameters(
+ /// 0,
+ /// &[
+ /// // the Field of the literal `[1.2, 2.1]`, the array being reduced
+ /// ValueOrLambda::Value(Arc::new(Field::new("", DataType::new_list(DataType::Float32, true), true))),
+ /// // the Field of the literal `0`, the initial value
+ /// ValueOrLambda::Value(Arc::new(Field::new("", DataType::Int32, true))),
+ /// // the Field of the output of the merge lambda, which is unknown at this point because it depends on
+ /// // the return this call
+ /// ValueOrLambda::Lambda(None),
+ /// // the Field of the output of the finish lambda, unknown for the same reason as above
+ /// ValueOrLambda::Lambda(None),
+ /// ])?;
+ ///
+ /// assert_eq!(
+ /// lambda_parameters,
+ /// LambdaParametersProgress::Partial(vec![
+ /// // the finish lambda supported parameters, regardless of how many are actually used
+ /// Some(vec![
+ /// // at step 0, use the field of the initial value
+ /// Arc::new(Field::new("ignored_name", DataType::Int32, true)),
+ /// // the array values being reduced
+ /// Arc::new(Field::new("", DataType::Float32, true)),
+ /// ]),
+ /// // the merge lambda supported parameters, unknown at this point due to dependency on the merge output
+ /// None,
+ /// ])
+ /// );
+ ///
+ /// let lambda_parameters = array_reduce.lambda_parameters(
+ /// 1,
+ /// &[
+ /// // the Field of the literal `[1.2, 2.1]`, the array being reduced
+ /// ValueOrLambda::Value(Arc::new(Field::new("", DataType::new_list(DataType::Float32, true), true))),
+ /// // the Field of the literal `0`, the initial value
+ /// ValueOrLambda::Value(Arc::new(Field::new("", DataType::Int32, true))),
+ /// // the Field of the output of the merge lambda, which could be inferred to be a Float32 based on the
+ /// // returned values of the previous step
+ /// ValueOrLambda::Value(Arc::new(Field::new("", DataType::Float32, true))),
+ /// // the Field of the output of the finish lambda, which is unknown at this point because it depends
+ /// // on the return of this call
+ /// ValueOrLambda::Lambda(None),
+ /// ])?;
+ ///
+ /// assert_eq!(
+ /// lambda_parameters,
+ /// LambdaParametersProgress::Complete(vec![
+ /// // the finish lambda supported parameters, regardless of how many are actually used
+ /// vec![
+ /// // the finish lambda own output now used as it's accumulator
+ /// Arc::new(Field::new("ignored_name", DataType::Float32, true)),
+ /// // the array values being reduced
+ /// Arc::new(Field::new("", DataType::Float32, true)),
+ /// ],
+ /// // the merge lambda supported parameters, which is the output of the merge lambda,
+ /// vec![
+ /// // the output of the merge lambda
+ /// Arc::new(Field::new("", DataType::Float32, true)),
+ /// ],
+ /// ])
+ /// );
+ ///
+ /// let coerce_to = array_reduce.coerce_values_for_lambdas(&[
+ /// // the literal `[1.2, 2.1]` data type, the array being reduced
+ /// ValueOrLambda::Value(DataType::new_list(DataType::Float32, true)),
+ /// // the literal `0` data type, the initial value
+ /// ValueOrLambda::Value(DataType::Int32),
+ /// // the output data type of the merge lambda
+ /// ValueOrLambda::Lambda(DataType::Float32),
+ /// // the output data type of the finish lambda
+ /// ValueOrLambda::Lambda(DataType::Boolean),
+ /// ])?;
+ ///
+ /// assert_eq!(
+ /// coerce_to,
+ /// vec![
+ /// // return the same type for the array being reduced
+ /// DataType::new_list(DataType::Float32, true),
+ /// // coerce the initial value to the output of the merge lambda
+ /// DataType::Float32,
+ /// ]
+ /// );
+ ///
+ /// ```
+ ///
+ /// Note this may also be called at step 0 with all lambda outputs already set, and in that case,
+ /// [LambdaParametersProgress::Complete] must be returned
+ ///
+ /// The implementation can assume that some other part of the code has coerced
+ /// the actual argument types to match [`Self::signature`], except the coercion defined by
+ /// [Self::coerce_values_for_lambdas], if applicable.
+ ///
+ /// [`HigherOrderFunction`]: crate::expr::HigherOrderFunction
+ /// [`HigherOrderFunction::lambda_parameters`]: crate::expr::HigherOrderFunction::lambda_parameters
+ fn lambda_parameters(
+ &self,
+ step: usize,
+ fields: &[ValueOrLambda<FieldRef, Option<FieldRef>>],
+ ) -> Result<LambdaParametersProgress>;
+
+ /// Coerce value arguments of a function call to types that the function can evaluate also taking into
+ /// account the *output type of it's lambdas*. This differs from [HigherOrderUDF::coerce_value_types]
+ /// that only has access to the type of it's value arguments because it's called before the output type
+ /// of lambdas are known. So that this method is called, the function must have it's
+ /// [HigherOrderSignature::coerce_values_for_lambdas] set to true
+ ///
+ /// See the [type coercion module](crate::type_coercion)
+ /// documentation for more details on type coercion
+ ///
+ /// # Parameters
+ /// * `fields`: The argument types of the value arguments of this function, or the output type of lambdas
+ ///
+ /// # Return value
+ /// A Vec with the same number of [ValueOrLambda::Value] in `fields`. DataFusion will `CAST` the
+ /// function call arguments to these specific types.
+ ///
+ /// For example, a flexible array_reduce implementation (see [Self::lambda_parameters] docs), when working
+ /// with the expression below, may want to coerce it's initial value argument, the *integer* `0`,
+ /// to match the output it's merge function, which is a *float*:
+ ///
+ /// `array_reduce([1.2, 2.1], 0, (acc, v) -> acc + v + 1.5, v -> v > 2.0)`
+ fn coerce_values_for_lambdas(
+ &self,
+ _fields: &[ValueOrLambda<DataType, DataType>],
+ ) -> Result<Vec<DataType>> {
+ not_impl_err!(
+ "{} coerce_values_for_lambdas is not implemented",
+ self.name()
+ )
+ }
+
+ /// What type will be returned by this function, given the arguments?
+ ///
+ /// The implementation can assume that some other part of the code has coerced
+ /// the actual argument types to match [`Self::signature`], including the coercion
+ /// defined by [Self::coerce_values_for_lambdas], if applicable.
+ ///
+ /// # Example creating `Field`
+ ///
+ /// Note the name of the `Field` is ignored, except for structured types such as
+ /// `DataType::Struct`.
+ ///
+ /// ```rust
+ /// # use std::sync::Arc;
+ /// # use arrow::datatypes::{DataType, Field, FieldRef};
+ /// # use datafusion_common::Result;
+ /// # use datafusion_expr::HigherOrderReturnFieldArgs;
+ /// # struct Example{}
+ /// # impl Example {
+ /// fn return_field_from_args(&self, args: HigherOrderReturnFieldArgs) -> Result<FieldRef> {
+ /// let field = Arc::new(Field::new("ignored_name", DataType::Int32, true));
+ /// Ok(field)
+ /// }
+ /// # }
+ /// ```
+ fn return_field_from_args(
+ &self,
+ args: HigherOrderReturnFieldArgs,
+ ) -> Result<FieldRef>;
+
+ /// Whether List or LargeList arguments should have it's non-empty null
+ /// sublists cleaned with [remove_list_null_values] before invoking this function
+ ///
+ /// The default implementation always returns true and should only be implemented
+ /// if you want to handle non-empty null sublists yourself
+ ///
+ /// [remove_list_null_values]: datafusion_common::utils::remove_list_null_values
+ // todo: extend this to listview and maps when remove_list_null_values supports it
+ fn clear_null_values(&self) -> bool {
+ true
+ }
+
+ /// Invoke the function returning the appropriate result.
+ ///
+ /// # Performance
+ ///
+ /// For the best performance, the implementations should handle the common case
+ /// when one or more of their arguments are constant values (aka
+ /// [`ColumnarValue::Scalar`]).
+ ///
+ /// [`ColumnarValue::values_to_arrays`] can be used to convert the arguments
+ /// to arrays, which will likely be simpler code, but be slower.
+ fn invoke_with_args(&self, args: HigherOrderFunctionArgs) -> Result<ColumnarValue>;
+
+ /// Returns true if some of this `exprs` subexpressions may not be evaluated
+ /// and thus any side effects (like divide by zero) may not be encountered.
+ ///
+ /// Setting this to true prevents certain optimizations such as common
+ /// subexpression elimination
+ ///
+ /// When overriding this function to return `true`, [HigherOrderUDF::conditional_arguments] can also be
+ /// overridden to report more accurately which arguments are eagerly evaluated and which ones
+ /// lazily.
+ fn short_circuits(&self) -> bool {
+ false
+ }
+
+ /// Determines which of the arguments passed to *this higher-order function*
+ /// are evaluated eagerly and which may be evaluated lazily. Note that this
+ /// does *not* applies to the arguments that *lambda functions* pass to it's
+ /// body expression
+ ///
+ /// If this function returns `None`, all arguments are eagerly evaluated.
+ /// Returning `None` is a micro optimization that saves a needless `Vec`
+ /// allocation.
+ ///
+ /// If the function returns `Some`, returns (`eager`, `lazy`) where `eager`
+ /// are the arguments that are always evaluated, and `lazy` are the
+ /// arguments that may be evaluated lazily (i.e. may not be evaluated at all
+ /// in some cases).
+ ///
+ /// Implementations must ensure that the two returned `Vec`s are disjunct,
+ /// and that each argument from `args` is present in one the two `Vec`s.
+ ///
+ /// When overriding this function, [HigherOrderUDF::short_circuits] must
+ /// be overridden to return `true`.
+ fn conditional_arguments<'a>(
+ &self,
+ args: &'a [Expr],
+ ) -> Option<(Vec<&'a Expr>, Vec<&'a Expr>)> {
+ if self.short_circuits() {
+ Some((vec![], args.iter().collect()))
+ } else {
+ None
+ }
+ }
+
+ /// Coerce value arguments of a function call to types that the function can evaluate.
+ /// Note that if you need to coerce values based on the output type of lambdas, you
+ /// must use [HigherOrderUDF::coerce_values_for_lambdas], as this function is used before
+ /// the output type of lambdas are known
+ ///
+ /// See the [type coercion module](crate::type_coercion)
+ /// documentation for more details on type coercion
+ ///
+ /// For example, if your function requires a contiguous list argument, but the user calls
+ /// it like `my_func(c, v -> v+2)` (i.e. with `c` as a ListView), coerce_types can return `[DataType::List(..)]`
+ /// to ensure the argument is converted to a List
+ ///
+ /// # Parameters
+ /// * `arg_types`: The argument types of the value arguments of this function, excluding lambdas
+ ///
+ /// # Return value
+ /// A Vec the same length as `arg_types`. DataFusion will `CAST` the function call
+ /// arguments to these specific types.
+ fn coerce_value_types(&self, _arg_types: &[DataType]) -> Result<Vec<DataType>> {
+ not_impl_err!(
+ "Function {} does not implement coerce_value_types",
+ self.name()
+ )
+ }
+
+ /// Returns the documentation for this HigherOrderUDF.
+ ///
+ /// Documentation can be accessed programmatically as well as generating
+ /// publicly facing documentation.
+ fn documentation(&self) -> Option<&Documentation> {
+ None
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use datafusion_expr_common::signature::Volatility;
+
+ use super::*;
+ use std::hash::DefaultHasher;
+
+ #[derive(Debug, PartialEq, Eq, Hash)]
+ struct TestHigherOrderUDF {
+ name: &'static str,
+ field: &'static str,
+ signature: HigherOrderSignature,
+ }
+ impl HigherOrderUDF for TestHigherOrderUDF {
+ fn name(&self) -> &str {
+ self.name
+ }
+
+ fn signature(&self) -> &HigherOrderSignature {
+ &self.signature
+ }
+
+ fn lambda_parameters(
+ &self,
+ _step: usize,
+ _fields: &[ValueOrLambda<FieldRef, Option<FieldRef>>],
+ ) -> Result<LambdaParametersProgress> {
+ unimplemented!()
+ }
+
+ fn return_field_from_args(
+ &self,
+ _args: HigherOrderReturnFieldArgs,
+ ) -> Result<FieldRef> {
+ unimplemented!()
+ }
+
+ fn invoke_with_args(
+ &self,
+ _args: HigherOrderFunctionArgs,
+ ) -> Result<ColumnarValue> {
+ unimplemented!()
+ }
+ }
+
+ // PartialEq and Hash must be consistent, and also PartialEq and PartialOrd
+ // must be consistent, so they are tested together.
+ #[test]
+ fn test_partial_eq_hash_and_partial_ord() {
+ // A parameterized function
+ let f = test_func("foo", "a");
+
+ // Same like `f`, different instance
+ let f2 = test_func("foo", "a");
+ assert_eq!(&f, &f2);
+ assert_eq!(hash(&f), hash(&f2));
+ assert_eq!(f.partial_cmp(&f2), Some(Ordering::Equal));
+
+ // Different parameter
+ let b = test_func("foo", "b");
+ assert_ne!(&f, &b);
+ assert_ne!(hash(&f), hash(&b)); // hash can collide for different values but does not collide in this test
+ assert_eq!(f.partial_cmp(&b), None);
+
+ // Different name
+ let o = test_func("other", "a");
+ assert_ne!(&f, &o);
+ assert_ne!(hash(&f), hash(&o)); // hash can collide for different values but does not collide in this test
+ assert_eq!(f.partial_cmp(&o), Some(Ordering::Less));
+
+ // Different name and parameter
+ assert_ne!(&b, &o);
+ assert_ne!(hash(&b), hash(&o)); // hash can collide for different values but does not collide in this test
+ assert_eq!(b.partial_cmp(&o), Some(Ordering::Less));
+ }
+
+ fn test_func(name: &'static str, parameter: &'static str) -> Arc<dyn HigherOrderUDF> {
+ Arc::new(TestHigherOrderUDF {
+ name,
+ field: parameter,
+ signature: HigherOrderSignature::variadic_any(Volatility::Immutable),
+ })
+ }
+
+ fn hash<T: Hash>(value: &T) -> u64 {
+ let hasher = &mut DefaultHasher::new();
+ value.hash(hasher);
+ hasher.finish()
+ }
+}
diff --git a/datafusion/expr/src/lib.rs b/datafusion/expr/src/lib.rs
index 51b7cf8638..92e6b13330 100644
--- a/datafusion/expr/src/lib.rs
+++ b/datafusion/expr/src/lib.rs
@@ -37,6 +37,7 @@
extern crate core;
+mod higher_order_function;
mod literal;
mod operation;
mod partition_evaluator;
@@ -112,6 +113,11 @@ pub use function::{
AccumulatorFactoryFunction, PartitionEvaluatorFactory, ReturnTypeFunction,
ScalarFunctionImplementation, StateTypeFunction,
};
+pub use higher_order_function::{
+ HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, HigherOrderSignature,
+ HigherOrderTypeSignature, HigherOrderUDF, LambdaArgument, LambdaParametersProgress,
+ ValueOrLambda,
+};
pub use literal::{
Literal, TimestampLiteral, lit, lit_timestamp_nano, lit_with_metadata,
};
diff --git a/datafusion/expr/src/planner.rs b/datafusion/expr/src/planner.rs
index 197ac8c035..d69f4ac5fe 100644
--- a/datafusion/expr/src/planner.rs
+++ b/datafusion/expr/src/planner.rs
@@ -24,8 +24,8 @@ use crate::expr::NullTreatment;
#[cfg(feature = "sql")]
use crate::logical_plan::LogicalPlan;
use crate::{
- AggregateUDF, Expr, GetFieldAccess, ScalarUDF, SortExpr, TableSource, WindowFrame,
- WindowFunctionDefinition, WindowUDF,
+ AggregateUDF, Expr, GetFieldAccess, HigherOrderUDF, ScalarUDF, SortExpr, TableSource,
+ WindowFrame, WindowFunctionDefinition, WindowUDF,
};
use arrow::datatypes::{DataType, Field, FieldRef, SchemaRef};
use datafusion_common::datatype::DataTypeExt;
@@ -103,6 +103,9 @@ pub trait ContextProvider {
/// Return the scalar function with a given name, if any
fn get_function_meta(&self, name: &str) -> Option<Arc<ScalarUDF>>;
+ /// Return the higher order function with a given name, if any
+ fn get_higher_order_meta(&self, name: &str) -> Option<Arc<dyn HigherOrderUDF>>;
+
/// Return the aggregate function with a given name, if any
fn get_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>>;
@@ -131,6 +134,9 @@ pub trait ContextProvider {
/// Return all scalar function names
fn udf_names(&self) -> Vec<String>;
+ /// Return all higher order function names
+ fn higher_order_function_names(&self) -> Vec<String>;
+
/// Return all aggregate function names
fn udaf_names(&self) -> Vec<String>;
diff --git a/datafusion/expr/src/registry.rs b/datafusion/expr/src/registry.rs
index 4be163a4bd..f03cc5936c 100644
--- a/datafusion/expr/src/registry.rs
+++ b/datafusion/expr/src/registry.rs
@@ -18,6 +18,7 @@
//! FunctionRegistry trait
use crate::expr_rewriter::FunctionRewrite;
+use crate::higher_order_function::HigherOrderUDF;
use crate::planner::ExprPlanner;
use crate::{AggregateUDF, ScalarUDF, UserDefinedLogicalNode, WindowUDF};
use arrow::datatypes::Field;
@@ -40,6 +41,9 @@ pub trait FunctionRegistry {
/// Returns names of all available scalar user defined functions.
fn udfs(&self) -> HashSet<String>;
+ /// Returns names of all available higher order user defined functions.
+ fn higher_order_function_names(&self) -> HashSet<String>;
+
/// Returns names of all available aggregate user defined functions.
fn udafs(&self) -> HashSet<String>;
@@ -50,6 +54,10 @@ pub trait FunctionRegistry {
/// `name`.
fn udf(&self, name: &str) -> Result<Arc<ScalarUDF>>;
+ /// Returns a reference to the user defined higher order function named
+ /// `name`.
+ fn higher_order_function(&self, name: &str) -> Result<Arc<dyn HigherOrderUDF>>;
+
/// Returns a reference to the user defined aggregate function (udaf) named
/// `name`.
fn udaf(&self, name: &str) -> Result<Arc<AggregateUDF>>;
@@ -66,6 +74,17 @@ pub trait FunctionRegistry {
fn register_udf(&mut self, _udf: Arc<ScalarUDF>) -> Result<Option<Arc<ScalarUDF>>> {
not_impl_err!("Registering ScalarUDF")
}
+ /// Registers a new [`HigherOrderUDF`], returning any previously registered
+ /// implementation.
+ ///
+ /// Returns an error (the default) if the function can not be registered,
+ /// for example if the registry is read only.
+ fn register_higher_order_function(
+ &mut self,
+ _function: Arc<dyn HigherOrderUDF>,
+ ) -> Result<Option<Arc<dyn HigherOrderUDF>>> {
+ not_impl_err!("Registering HigherOrderUDF")
+ }
/// Registers a new [`AggregateUDF`], returning any previously registered
/// implementation.
///
@@ -95,6 +114,18 @@ pub trait FunctionRegistry {
not_impl_err!("Deregistering ScalarUDF")
}
+ /// Deregisters a [`HigherOrderUDF`], returning the implementation that was
+ /// deregistered.
+ ///
+ /// Returns an error (the default) if the function can not be deregistered,
+ /// for example if the registry is read only.
+ fn deregister_higher_order_function(
+ &mut self,
+ _name: &str,
+ ) -> Result<Option<Arc<dyn HigherOrderUDF>>> {
+ not_impl_err!("Deregistering HigherOrderUDF")
+ }
+
/// Deregisters a [`AggregateUDF`], returning the implementation that was
/// deregistered.
///
@@ -166,6 +197,8 @@ pub struct MemoryFunctionRegistry {
udafs: HashMap<String, Arc<AggregateUDF>>,
/// Window Functions
udwfs: HashMap<String, Arc<WindowUDF>>,
+ /// Higher Order Functions
+ higher_order_functions: HashMap<String, Arc<dyn HigherOrderUDF>>,
}
impl MemoryFunctionRegistry {
@@ -186,6 +219,13 @@ impl FunctionRegistry for MemoryFunctionRegistry {
.ok_or_else(|| plan_datafusion_err!("Function {name} not found"))
}
+ fn higher_order_function(&self, name: &str) -> Result<Arc<dyn HigherOrderUDF>> {
+ self.higher_order_functions
+ .get(name)
+ .cloned()
+ .ok_or_else(|| plan_datafusion_err!("Higher Order Function {name} not found"))
+ }
+
fn udaf(&self, name: &str) -> Result<Arc<AggregateUDF>> {
self.udafs
.get(name)
@@ -203,6 +243,14 @@ impl FunctionRegistry for MemoryFunctionRegistry {
fn register_udf(&mut self, udf: Arc<ScalarUDF>) -> Result<Option<Arc<ScalarUDF>>> {
Ok(self.udfs.insert(udf.name().to_string(), udf))
}
+ fn register_higher_order_function(
+ &mut self,
+ function: Arc<dyn HigherOrderUDF>,
+ ) -> Result<Option<Arc<dyn HigherOrderUDF>>> {
+ Ok(self
+ .higher_order_functions
+ .insert(function.name().into(), function))
+ }
fn register_udaf(
&mut self,
udaf: Arc<AggregateUDF>,
@@ -217,6 +265,10 @@ impl FunctionRegistry for MemoryFunctionRegistry {
vec![]
}
+ fn higher_order_function_names(&self) -> HashSet<String> {
+ self.higher_order_functions.keys().cloned().collect()
+ }
+
fn udafs(&self) -> HashSet<String> {
self.udafs.keys().cloned().collect()
}
diff --git a/datafusion/expr/src/tree_node.rs b/datafusion/expr/src/tree_node.rs
index f43b138a28..010441b5a2 100644
--- a/datafusion/expr/src/tree_node.rs
+++ b/datafusion/expr/src/tree_node.rs
@@ -17,16 +17,20 @@
//! Tree node implementation for Logical Expressions
-use crate::Expr;
-use crate::expr::{
- AggregateFunction, AggregateFunctionParams, Alias, Between, BinaryExpr, Case, Cast,
- GroupingSet, InList, InSubquery, Like, Placeholder, ScalarFunction, SetComparison,
- TryCast, Unnest, WindowFunction, WindowFunctionParams,
+use crate::{
+ Expr,
+ expr::{
+ AggregateFunction, AggregateFunctionParams, Alias, Between, BinaryExpr, Case,
+ Cast, GroupingSet, HigherOrderFunction, InList, InSubquery, Lambda, Like,
+ Placeholder, ScalarFunction, SetComparison, TryCast, Unnest, WindowFunction,
+ WindowFunctionParams,
+ },
};
-
-use datafusion_common::Result;
-use datafusion_common::tree_node::{
- Transformed, TreeNode, TreeNodeContainer, TreeNodeRecursion, TreeNodeRefContainer,
+use datafusion_common::{
+ Result,
+ tree_node::{
+ Transformed, TreeNode, TreeNodeContainer, TreeNodeRecursion, TreeNodeRefContainer,
+ },
};
/// Implementation of the [`TreeNode`] trait
@@ -78,7 +82,8 @@ impl TreeNode for Expr {
| Expr::Exists { .. }
| Expr::ScalarSubquery(_)
| Expr::Wildcard { .. }
- | Expr::Placeholder(_) => Ok(TreeNodeRecursion::Continue),
+ | Expr::Placeholder(_)
+ | Expr::LambdaVariable(_) => Ok(TreeNodeRecursion::Continue),
Expr::BinaryExpr(BinaryExpr { left, right, .. }) => {
(left, right).apply_ref_elements(f)
}
@@ -107,6 +112,8 @@ impl TreeNode for Expr {
Expr::InList(InList { expr, list, .. }) => {
(expr, list).apply_ref_elements(f)
}
+ Expr::HigherOrderFunction(HigherOrderFunction { func: _, args}) => args.apply_elements(f),
+ Expr::Lambda (Lambda{ params: _, body}) => body.apply_elements(f)
}
}
@@ -128,7 +135,8 @@ impl TreeNode for Expr {
| Expr::Exists { .. }
| Expr::ScalarSubquery(_)
| Expr::ScalarVariable(_, _)
- | Expr::Literal(_, _) => Transformed::no(self),
+ | Expr::Literal(_, _)
+ | Expr::LambdaVariable(_) => Transformed::no(self),
Expr::SetComparison(SetComparison {
expr,
subquery,
@@ -330,6 +338,14 @@ impl TreeNode for Expr {
.update_data(|(new_expr, new_list)| {
Expr::InList(InList::new(new_expr, new_list, negated))
}),
+ Expr::HigherOrderFunction(HigherOrderFunction { func, args }) => {
+ args.map_elements(f)?.update_data(|args| {
+ Expr::HigherOrderFunction(HigherOrderFunction { func, args })
+ })
+ }
+ Expr::Lambda(Lambda { params, body }) => body
+ .map_elements(f)?
+ .update_data(|body| Expr::Lambda(Lambda { params, body })),
})
}
}
diff --git a/datafusion/expr/src/type_coercion/functions.rs b/datafusion/expr/src/type_coercion/functions.rs
index 7b6cb7b35d..b2d7a860eb 100644
--- a/datafusion/expr/src/type_coercion/functions.rs
+++ b/datafusion/expr/src/type_coercion/functions.rs
@@ -16,12 +16,16 @@
// under the License.
use super::binary::binary_numeric_coercion;
-use crate::{AggregateUDF, ScalarUDF, Signature, TypeSignature, WindowUDF};
+use crate::{
+ AggregateUDF, HigherOrderTypeSignature, HigherOrderUDF, ScalarUDF, Signature,
+ TypeSignature, ValueOrLambda, WindowUDF,
+};
use arrow::datatypes::{Field, FieldRef};
use arrow::{
compute::can_cast_types,
datatypes::{DataType, TimeUnit},
};
+use datafusion_common::internal_datafusion_err;
use datafusion_common::types::LogicalType;
use datafusion_common::utils::{
ListCoercion, base_type, coerced_fixed_size_list_to_list,
@@ -148,6 +152,150 @@ pub fn fields_with_udf<F: UDFCoercionExt>(
.collect())
}
+/// Performs type coercion for higher order function arguments.
+///
+/// For value arguments, returns the field to which each
+/// argument must be coerced to match `signature`.
+/// For lambda arguments, returns a clone of the associated data
+///
+/// Note this does not invokes [HigherOrderUDF::coerce_values_for_lambdas]
+/// if requested by the function signature. If that's required, use
+/// [value_fields_with_higher_order_udf_and_lambdas] instead
+///
+/// For more details on coercion in general, please see the
+/// [`type_coercion`](crate::type_coercion) module.
+pub fn value_fields_with_higher_order_udf<L: Clone>(
+ current_fields: &[ValueOrLambda<FieldRef, L>],
+ func: &dyn HigherOrderUDF,
+) -> Result<Vec<ValueOrLambda<FieldRef, L>>> {
+ match func.signature().type_signature {
+ HigherOrderTypeSignature::UserDefined => {
+ let arg_types = current_fields
+ .iter()
+ .filter_map(|p| match p {
+ ValueOrLambda::Value(field) => Some(field.data_type().clone()),
+ ValueOrLambda::Lambda(_) => None,
+ })
+ .collect::<Vec<_>>();
+
+ let coerced_types = func.coerce_value_types(&arg_types)?;
+
+ if coerced_types.len() != arg_types.len() {
+ return plan_err!(
+ "{} coerce_value_types should have returned {} items but returned {}",
+ func.name(),
+ arg_types.len(),
+ coerced_types.len()
+ );
+ }
+
+ // coerced_types has been partitioned from current_fields
+ // and refers only to values and not to lambdas, so instead
+ // of zipping them, we iterate over current_fields and only
+ // consume from coerced_types when a given argument is a value
+ // to reconstruct the arguments list with the correct order
+ // this supports any value and lambda positioning including
+ // multiple lambdas interleaved with values
+ let mut coerced_types = coerced_types.into_iter();
+
+ current_fields
+ .iter()
+ .map(|current_field| match current_field {
+ ValueOrLambda::Value(field) => {
+ let data_type = coerced_types.next().ok_or_else(|| {
+ internal_datafusion_err!(
+ "coerced_types len should have been checked above"
+ )
+ })?;
+
+ Ok(ValueOrLambda::Value(Arc::new(
+ field.as_ref().clone().with_data_type(data_type),
+ )))
+ }
+ ValueOrLambda::Lambda(lambda) => {
+ Ok(ValueOrLambda::Lambda(lambda.clone()))
+ }
+ })
+ .collect()
+ }
+ HigherOrderTypeSignature::VariadicAny => Ok(current_fields.to_vec()),
+ HigherOrderTypeSignature::Any(number) => {
+ if current_fields.len() != number {
+ return plan_err!(
+ "The function '{}' expected {number} arguments but received {}",
+ func.name(),
+ current_fields.len()
+ );
+ }
+
+ Ok(current_fields.to_vec())
+ }
+ }
+}
+
+/// Performs type coercion for higher order function arguments,
+/// including those defined by [HigherOrderUDF::coerce_values_for_lambdas],
+/// if defined by the signature. Note that compared to
+/// [value_fields_with_higher_order_udf], this function requires
+/// the [ValueOrLambda::Lambda] variant to contain the output field of the lambda.
+///
+/// For value arguments, returns the field to which each
+/// argument must be coerced to match `signature`.
+/// For lambda arguments, returns a clone of the output field
+///
+/// For more details on coercion in general, please see the
+/// [`type_coercion`](crate::type_coercion) module.
+pub fn value_fields_with_higher_order_udf_and_lambdas(
+ current_fields: &[ValueOrLambda<FieldRef, FieldRef>],
+ func: &dyn HigherOrderUDF,
+) -> Result<Vec<ValueOrLambda<FieldRef, FieldRef>>> {
+ let mut new_fields = value_fields_with_higher_order_udf(current_fields, func)?;
+
+ if func.signature().coerce_values_for_lambdas {
+ let new_types = new_fields
+ .iter()
+ .map(|f| match f {
+ ValueOrLambda::Value(f) => ValueOrLambda::Value(f.data_type().clone()),
+ ValueOrLambda::Lambda(f) => ValueOrLambda::Lambda(f.data_type().clone()),
+ })
+ .collect::<Vec<_>>();
+
+ let mut new_value_types = func.coerce_values_for_lambdas(&new_types)?.into_iter();
+
+ let value_types_count = new_types
+ .iter()
+ .filter(|e| matches!(e, ValueOrLambda::Value(_)))
+ .count();
+
+ if new_value_types.len() != value_types_count {
+ return plan_err!(
+ "{} coerce_values_for_lambdas returned {} values but {value_types_count} expected",
+ func.name(),
+ new_value_types.len()
+ );
+ }
+
+ for new_field in &mut new_fields {
+ match new_field {
+ ValueOrLambda::Value(value) => {
+ let coerce_to = new_value_types.next().ok_or_else(|| {
+ internal_datafusion_err!(
+ "new_value_types len should have been checked above"
+ )
+ })?;
+
+ if value.data_type() != &coerce_to {
+ Arc::make_mut(value).set_data_type(coerce_to);
+ }
+ }
+ ValueOrLambda::Lambda(_) => {}
+ }
+ }
+ };
+
+ Ok(new_fields)
+}
+
/// Performs type coercion for scalar function arguments.
///
/// Returns the data types to which each argument must be coerced to
@@ -943,7 +1091,10 @@ fn coerced_from<'a>(
#[cfg(test)]
mod tests {
- use crate::Volatility;
+ use crate::{
+ HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, HigherOrderSignature,
+ Volatility,
+ };
use super::*;
use arrow::datatypes::IntervalUnit;
@@ -951,7 +1102,10 @@ mod tests {
assert_contains,
types::{logical_binary, logical_int64},
};
- use datafusion_expr_common::signature::{Coercion, TypeSignatureClass};
+ use datafusion_expr_common::{
+ columnar_value::ColumnarValue,
+ signature::{Coercion, TypeSignatureClass},
+ };
#[test]
fn test_string_conversion() {
@@ -1616,4 +1770,208 @@ mod tests {
Ok(())
}
+
+ #[derive(Debug, PartialEq, Eq, Hash)]
+ struct MockHigherOrderUDF {
+ signature: HigherOrderSignature,
+ coerced_value_types: Vec<DataType>,
+ }
+
+ impl HigherOrderUDF for MockHigherOrderUDF {
+ fn name(&self) -> &str {
+ "mock_higher_order_function"
+ }
+
+ fn signature(&self) -> &HigherOrderSignature {
+ &self.signature
+ }
+
+ fn coerce_value_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
+ if arg_types.len() != 1 {
+ return plan_err!(
+ "mock_higher_order_function expects 1 value arguments, got {}",
+ arg_types.len()
+ );
+ }
+ Ok(self.coerced_value_types.clone())
+ }
+
+ fn coerce_values_for_lambdas(
+ &self,
+ fields: &[ValueOrLambda<DataType, DataType>],
+ ) -> Result<Vec<DataType>> {
+ // thoerical impl of array_reduce without finish
+ let [
+ ValueOrLambda::Value(list),
+ ValueOrLambda::Value(_initial),
+ ValueOrLambda::Lambda(merge),
+ ] = fields
+ else {
+ unreachable!()
+ };
+
+ Ok(vec![list.clone(), merge.clone()])
+ }
+
+ fn lambda_parameters(
+ &self,
+ _step: usize,
+ _fields: &[ValueOrLambda<FieldRef, Option<FieldRef>>],
+ ) -> Result<crate::LambdaParametersProgress> {
+ unimplemented!("mock_higher_order_function")
+ }
+
+ fn return_field_from_args(
+ &self,
+ _args: HigherOrderReturnFieldArgs,
+ ) -> Result<FieldRef> {
+ unimplemented!("mock_higher_order_function")
+ }
+
+ fn invoke_with_args(
+ &self,
+ _args: HigherOrderFunctionArgs,
+ ) -> Result<ColumnarValue> {
+ unimplemented!("mock_higher_order_function")
+ }
+ }
+
+ #[test]
+ fn test_higher_order_function_user_defined_type_coercion() {
+ let fun = MockHigherOrderUDF {
+ signature: HigherOrderSignature::user_defined(Volatility::Immutable),
+ coerced_value_types: vec![DataType::new_large_list(DataType::Int32, false)],
+ };
+
+ let new_fields = value_fields_with_higher_order_udf(
+ &[
+ ValueOrLambda::Value(Arc::new(Field::new_list(
+ "",
+ Field::new_list_field(DataType::Int32, false),
+ false,
+ ))),
+ ValueOrLambda::Lambda(()),
+ ],
+ &fun,
+ )
+ .unwrap();
+
+ // from List(Int32) to LargeList(Int32)
+ assert_eq!(
+ new_fields,
+ vec![
+ ValueOrLambda::Value(Arc::new(Field::new_large_list(
+ "",
+ Field::new_list_field(DataType::Int32, false),
+ false
+ ))),
+ ValueOrLambda::Lambda(()),
+ ]
+ )
+ }
+
+ #[test]
+ fn test_higher_order_function_coerce_values_for_lambdas() {
+ let fun = MockHigherOrderUDF {
+ signature: HigherOrderSignature::variadic_any(Volatility::Immutable)
+ .with_coerce_values_for_lambdas(),
+ coerced_value_types: vec![],
+ };
+
+ let new_fields = value_fields_with_higher_order_udf_and_lambdas(
+ &[
+ ValueOrLambda::Value(Arc::new(Field::new_list(
+ "",
+ Field::new_list_field(DataType::Float32, true),
+ true,
+ ))),
+ ValueOrLambda::Value(Arc::new(Field::new("", DataType::Int32, true))),
+ ValueOrLambda::Lambda(Arc::new(Field::new("", DataType::Float32, true))),
+ ],
+ &fun,
+ )
+ .unwrap();
+
+ // second parameter from Int32 to Float32
+ assert_eq!(
+ new_fields,
+ vec![
+ ValueOrLambda::Value(Arc::new(Field::new_list(
+ "",
+ Field::new_list_field(DataType::Float32, true),
+ true,
+ ))),
+ ValueOrLambda::Value(Arc::new(Field::new("", DataType::Float32, true))),
+ ValueOrLambda::Lambda(Arc::new(Field::new("", DataType::Float32, true))),
+ ]
+ )
+ }
+
+ #[test]
+ fn test_higher_order_function_user_defined_type_coercion_bad_args() {
+ let fun = MockHigherOrderUDF {
+ signature: HigherOrderSignature::user_defined(Volatility::Immutable),
+ coerced_value_types: vec![DataType::Int32],
+ };
+
+ let err = value_fields_with_higher_order_udf::<()>(&[], &fun).unwrap_err();
+
+ assert_contains!(
+ err.to_string(),
+ "mock_higher_order_function expects 1 value arguments, got 0"
+ );
+ }
+
+ #[test]
+ fn test_higher_order_function_faulty_user_defined_type_coercion() {
+ let fun = MockHigherOrderUDF {
+ signature: HigherOrderSignature::user_defined(Volatility::Immutable),
+ coerced_value_types: vec![DataType::Int32, DataType::Int32],
+ };
+
+ let err = value_fields_with_higher_order_udf::<()>(
+ &[ValueOrLambda::Value(Arc::new(Field::new(
+ "",
+ DataType::Int32,
+ false,
+ )))],
+ &fun,
+ )
+ .unwrap_err();
+
+ assert_contains!(
+ err.to_string(),
+ "mock_higher_order_function coerce_value_types should have returned 1 items but returned 2"
+ );
+ }
+
+ #[test]
+ fn test_higher_order_function_any_signature() {
+ let fun = MockHigherOrderUDF {
+ signature: HigherOrderSignature::any(1, Volatility::Immutable),
+ coerced_value_types: vec![],
+ };
+
+ let new_fields =
+ value_fields_with_higher_order_udf(&[ValueOrLambda::Lambda(())], &fun)
+ .unwrap();
+
+ // no coercion, just number of args checked
+ assert_eq!(new_fields, vec![ValueOrLambda::Lambda(())])
+ }
+
+ #[test]
+ fn test_higher_order_function_any_signature_bad_args() {
+ let fun = MockHigherOrderUDF {
+ signature: HigherOrderSignature::any(1, Volatility::Immutable),
+ coerced_value_types: vec![],
+ };
+
+ let err = value_fields_with_higher_order_udf::<()>(&[], &fun).unwrap_err();
+
+ assert_contains!(
+ err.to_string(),
+ "The function 'mock_higher_order_function' expected 1 arguments but received 0"
+ );
+ }
}
diff --git a/datafusion/expr/src/udf_eq.rs b/datafusion/expr/src/udf_eq.rs
index 3eed4337ee..5fb0266aef 100644
--- a/datafusion/expr/src/udf_eq.rs
+++ b/datafusion/expr/src/udf_eq.rs
@@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.
-use crate::{AggregateUDFImpl, ScalarUDFImpl, WindowUDFImpl};
+use crate::{AggregateUDFImpl, HigherOrderUDF, ScalarUDFImpl, WindowUDFImpl};
use std::any::Any;
use std::fmt::Debug;
use std::hash::{DefaultHasher, Hash, Hasher};
@@ -94,6 +94,18 @@ impl UdfPointer for Arc<dyn ScalarUDFImpl + '_> {
}
}
+impl UdfPointer for Arc<dyn HigherOrderUDF + '_> {
+ fn equals(&self, other: &Self::Target) -> bool {
+ self.as_ref().dyn_eq(other)
+ }
+
+ fn hash_value(&self) -> u64 {
+ let hasher = &mut DefaultHasher::new();
+ self.as_ref().dyn_hash(hasher);
+ hasher.finish()
+ }
+}
+
impl UdfPointer for Arc<dyn AggregateUDFImpl + '_> {
fn equals(&self, other: &(dyn AggregateUDFImpl + '_)) -> bool {
self.as_ref().dyn_eq(other)
diff --git a/datafusion/expr/src/utils.rs b/datafusion/expr/src/utils.rs
index 9439d8c590..fdb4b6de78 100644
--- a/datafusion/expr/src/utils.rs
+++ b/datafusion/expr/src/utils.rs
@@ -316,7 +316,10 @@ pub fn expr_to_columns(expr: &Expr, accum: &mut HashSet<Column>) -> Result<()> {
| Expr::ScalarSubquery(_)
| Expr::Wildcard { .. }
| Expr::Placeholder(_)
- | Expr::OuterReferenceColumn { .. } => {}
+ | Expr::OuterReferenceColumn { .. }
+ | Expr::HigherOrderFunction(_)
+ | Expr::Lambda(_)
+ | Expr::LambdaVariable(_) => {}
}
Ok(TreeNodeRecursion::Continue)
})
diff --git a/datafusion/ffi/src/execution/task_ctx.rs b/datafusion/ffi/src/execution/task_ctx.rs
index d977e56029..0a48a5fe5a 100644
--- a/datafusion/ffi/src/execution/task_ctx.rs
+++ b/datafusion/ffi/src/execution/task_ctx.rs
@@ -15,6 +15,7 @@
// specific language governing permissions and limitations
// under the License.
+use std::collections::HashMap;
use std::ffi::c_void;
use std::sync::Arc;
@@ -234,6 +235,7 @@ impl From<FFI_TaskContext> for Arc<TaskContext> {
session_id,
session_config,
scalar_functions,
+ HashMap::new(),
aggregate_functions,
window_functions,
runtime,
diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs
index 732be5b5b9..dfc9d1c7df 100644
--- a/datafusion/ffi/src/session/mod.rs
+++ b/datafusion/ffi/src/session/mod.rs
@@ -32,8 +32,8 @@ use datafusion_execution::runtime_env::RuntimeEnv;
use datafusion_expr::execution_props::ExecutionProps;
use datafusion_expr::registry::{ExtensionTypeRegistryRef, MemoryExtensionTypeRegistry};
use datafusion_expr::{
- AggregateUDF, AggregateUDFImpl, Expr, LogicalPlan, ScalarUDF, ScalarUDFImpl,
- WindowUDF, WindowUDFImpl,
+ AggregateUDF, AggregateUDFImpl, Expr, HigherOrderUDF, LogicalPlan, ScalarUDF,
+ ScalarUDFImpl, WindowUDF, WindowUDFImpl,
};
use datafusion_physical_expr::PhysicalExpr;
use datafusion_physical_plan::ExecutionPlan;
@@ -378,6 +378,7 @@ pub struct ForeignSession {
session: FFI_SessionRef,
config: SessionConfig,
scalar_functions: HashMap<String, Arc<ScalarUDF>>,
+ higher_order_functions: HashMap<String, Arc<dyn HigherOrderUDF>>,
aggregate_functions: HashMap<String, Arc<AggregateUDF>>,
window_functions: HashMap<String, Arc<WindowUDF>>,
extension_types: ExtensionTypeRegistryRef,
@@ -447,6 +448,7 @@ impl TryFrom<&FFI_SessionRef> for ForeignSession {
config,
table_options,
scalar_functions,
+ higher_order_functions: HashMap::new(),
aggregate_functions,
window_functions,
extension_types: Arc::new(MemoryExtensionTypeRegistry::default()),
@@ -588,6 +590,10 @@ impl Session for ForeignSession {
&self.scalar_functions
}
+ fn higher_order_functions(&self) -> &HashMap<String, Arc<dyn HigherOrderUDF>> {
+ &self.higher_order_functions
+ }
+
fn aggregate_functions(&self) -> &HashMap<String, Arc<AggregateUDF>> {
&self.aggregate_functions
}
diff --git a/datafusion/functions-nested/Cargo.toml b/datafusion/functions-nested/Cargo.toml
index d2204a7be7..b25aa0b9bc 100644
--- a/datafusion/functions-nested/Cargo.toml
+++ b/datafusion/functions-nested/Cargo.toml
@@ -67,6 +67,8 @@ memchr = { workspace = true }
arrow = { workspace = true, features = ["test_utils"] }
criterion = { workspace = true, features = ["async_tokio"] }
rand = { workspace = true }
+# used to test array_transform
+datafusion-physical-expr = { workspace = true }
[[bench]]
harness = false
diff --git a/datafusion/functions-nested/src/array_transform.rs b/datafusion/functions-nested/src/array_transform.rs
new file mode 100644
index 0000000000..0dcc7a7613
--- /dev/null
+++ b/datafusion/functions-nested/src/array_transform.rs
@@ -0,0 +1,392 @@
+// 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.
+
+//! [`HigherOrderUDF`] definitions for array_transform function.
+
+use arrow::{
+ array::{Array, ArrayRef, AsArray, LargeListArray, ListArray},
+ datatypes::{DataType, Field, FieldRef},
+};
+use datafusion_common::{
+ Result, ScalarValue, exec_err, plan_err,
+ utils::{adjust_offsets_for_slice, list_values, take_function_args},
+};
+use datafusion_expr::{
+ ColumnarValue, Documentation, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs,
+ HigherOrderSignature, HigherOrderUDF, LambdaParametersProgress, ValueOrLambda,
+ Volatility,
+};
+use datafusion_macros::user_doc;
+use std::{fmt::Debug, sync::Arc};
+
+make_higher_order_function_expr_and_func!(
+ ArrayTransform,
+ array_transform,
+ array lambda,
+ "transforms the values of an array",
+ array_transform_higher_order_function
+);
+
+#[user_doc(
+ doc_section(label = "Array Functions"),
+ description = "transforms the values of an array",
+ syntax_example = "array_transform(array, x -> x*2)",
+ sql_example = r#"```sql
+> select array_transform([1, 2, 3, 4, 5], x -> x*2);
++-------------------------------------------+
+| array_transform([1, 2, 3, 4, 5], x -> x*2) |
++-------------------------------------------+
+| [2, 4, 6, 8, 10] |
++-------------------------------------------+
+```"#,
+ argument(
+ name = "array",
+ description = "Array expression. Can be a constant, column, or function, and any combination of array operators."
+ ),
+ argument(name = "lambda", description = "Lambda")
+)]
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct ArrayTransform {
+ signature: HigherOrderSignature,
+ aliases: Vec<String>,
+}
+
+impl Default for ArrayTransform {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ArrayTransform {
+ pub fn new() -> Self {
+ Self {
+ signature: HigherOrderSignature::user_defined(Volatility::Immutable),
+ aliases: vec![String::from("list_transform")],
+ }
+ }
+}
+
+impl HigherOrderUDF for ArrayTransform {
+ fn name(&self) -> &str {
+ "array_transform"
+ }
+
+ fn aliases(&self) -> &[String] {
+ &self.aliases
+ }
+
+ fn signature(&self) -> &HigherOrderSignature {
+ &self.signature
+ }
+
+ fn coerce_value_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
+ let list = if arg_types.len() == 1 {
+ &arg_types[0]
+ } else {
+ return plan_err!(
+ "{} function requires 1 value arguments, got {}",
+ self.name(),
+ arg_types.len()
+ );
+ };
+
+ let coerced = match list {
+ DataType::List(_) | DataType::LargeList(_) => list.clone(),
+ DataType::ListView(field) | DataType::FixedSizeList(field, _) => {
+ DataType::List(Arc::clone(field))
+ }
+ DataType::LargeListView(field) => DataType::LargeList(Arc::clone(field)),
+ _ => {
+ return plan_err!(
+ "{} expected a list as first argument, got {}",
+ self.name(),
+ list
+ );
+ }
+ };
+
+ Ok(vec![coerced])
+ }
+
+ fn lambda_parameters(
+ &self,
+ _step: usize,
+ fields: &[ValueOrLambda<FieldRef, Option<FieldRef>>],
+ ) -> Result<LambdaParametersProgress> {
+ let (list, _lambda) = value_lambda_pair(self.name(), fields)?;
+
+ let field = match list.data_type() {
+ DataType::List(field) => field,
+ DataType::LargeList(field) => field,
+ _ => return plan_err!("expected list, got {list}"),
+ };
+
+ // we don't need to check whether the lambda contains more than two parameters,
+ // e.g. array_transform([], (v, i, j) -> v+i+j), as datafusion will do that for us
+ Ok(LambdaParametersProgress::Complete(vec![vec![Arc::clone(
+ field,
+ )]]))
+ }
+
+ fn return_field_from_args(
+ &self,
+ args: HigherOrderReturnFieldArgs,
+ ) -> Result<Arc<Field>> {
+ let (list, lambda) = value_lambda_pair(self.name(), args.arg_fields)?;
+
+ //TODO: should metadata be copied into the transformed array?
+
+ // lambda is the resulting field of executing the lambda body
+ // with the parameters returned in lambda_parameters
+ let field = Arc::new(Field::new(
+ Field::LIST_FIELD_DEFAULT_NAME,
+ lambda.data_type().clone(),
+ lambda.is_nullable(),
+ ));
+
+ let return_type = match list.data_type() {
+ DataType::List(_) => DataType::List(field),
+ DataType::LargeList(_) => DataType::LargeList(field),
+ other => plan_err!("expected list, got {other}")?,
+ };
+
+ Ok(Arc::new(Field::new("", return_type, list.is_nullable())))
+ }
+
+ fn invoke_with_args(&self, args: HigherOrderFunctionArgs) -> Result<ColumnarValue> {
+ let (list, lambda) = value_lambda_pair(self.name(), &args.args)?;
+
+ let list_array = list.to_array(args.number_rows)?;
+
+ // Fast path for fully null input array
+ if list_array.null_count() == list_array.len() {
+ return Ok(ColumnarValue::Scalar(ScalarValue::try_new_null(
+ args.return_type(),
+ )?));
+ }
+
+ // as per list_values docs, if list_array is sliced, list_values will be sliced too,
+ // so before constructing the transformed array below, we must adjust the list offsets with
+ // adjust_offsets_for_slice
+ let list_values = list_values(&list_array)?;
+
+ // fast path: when every sublist is empty and non-null we can return a scalar of an non-null empty sublist.
+ // If every sublist is null have already been handled above
+ if list_values.is_empty()
+ && list_array.null_count() == 0
+ && matches!(
+ args.return_type(),
+ DataType::List(_) | DataType::LargeList(_)
+ )
+ {
+ return Ok(ColumnarValue::Scalar(ScalarValue::new_default(
+ args.return_type(),
+ )?));
+ }
+
+ // by passing closures, lambda.evaluate can evaluate only those actually needed
+ let values_param = || Ok(Arc::clone(&list_values));
+
+ // call the transforming lambda
+ let transformed_values = lambda
+ .evaluate(&[&values_param])?
+ .into_array(list_values.len())?;
+
+ let field = match args.return_field.data_type() {
+ DataType::List(field) | DataType::LargeList(field) => Arc::clone(field),
+ _ => {
+ return exec_err!(
+ "{} expected ScalarFunctionArgs.return_field to be a list, got {}",
+ self.name(),
+ args.return_field
+ );
+ }
+ };
+
+ let transformed_list = match list_array.data_type() {
+ DataType::List(_) => {
+ let list = list_array.as_list();
+
+ // since we called list_values above which would return sliced values for
+ // a sliced list, we must adjust the offsets here as otherwise they would be invalid
+ let adjusted_offsets = adjust_offsets_for_slice(list);
+
+ Arc::new(ListArray::new(
+ field,
+ adjusted_offsets,
+ transformed_values,
+ list.nulls().cloned(),
+ )) as ArrayRef
+ }
+ DataType::LargeList(_) => {
+ let large_list = list_array.as_list();
+
+ // since we called list_values above which would return sliced values for
+ // a sliced list, we must adjust the offsets here as otherwise they would be invalid
+ let adjusted_offsets = adjust_offsets_for_slice(large_list);
+
+ Arc::new(LargeListArray::new(
+ field,
+ adjusted_offsets,
+ transformed_values,
+ large_list.nulls().cloned(),
+ ))
+ }
+ other => exec_err!("expected list, got {other}")?,
+ };
+
+ Ok(ColumnarValue::Array(transformed_list))
+ }
+
+ fn documentation(&self) -> Option<&Documentation> {
+ self.doc()
+ }
+}
+
+fn value_lambda_pair<'a, V: Debug, L: Debug>(
+ name: &str,
+ args: &'a [ValueOrLambda<V, L>],
+) -> Result<(&'a V, &'a L)> {
+ let [value, lambda] = take_function_args(name, args)?;
+
+ let (ValueOrLambda::Value(value), ValueOrLambda::Lambda(lambda)) = (value, lambda)
+ else {
+ return plan_err!(
+ "{name} expects a value followed by a lambda, got {value:?} and {lambda:?}"
+ );
+ };
+
+ Ok((value, lambda))
+}
+
+#[cfg(test)]
+mod tests {
+ use std::{collections::HashMap, sync::Arc};
+
+ use arrow::{
+ array::{Array, ArrayRef, AsArray, Int32Array, ListArray, RecordBatch},
+ buffer::{NullBuffer, OffsetBuffer},
+ datatypes::{DataType, Field},
+ };
+ use datafusion_common::{DFSchema, Result};
+ use datafusion_expr::{
+ Expr, col,
+ execution_props::ExecutionProps,
+ expr::{HigherOrderFunction, LambdaVariable},
+ lambda, lit,
+ };
+ use datafusion_physical_expr::create_physical_expr;
+
+ use crate::array_transform::array_transform_higher_order_function;
+
+ fn create_i32_list(
+ values: impl Into<Int32Array>,
+ offsets: OffsetBuffer<i32>,
+ nulls: Option<NullBuffer>,
+ ) -> ListArray {
+ let list_field = Arc::new(Field::new_list_field(DataType::Int32, true));
+
+ ListArray::new(list_field, offsets, Arc::new(values.into()), nulls)
+ }
+
+ fn divide_100_by(list: impl Array + Clone + 'static) -> Result<ArrayRef> {
+ let array_transform = array_transform_higher_order_function();
+
+ let schema = DFSchema::from_unqualified_fields(
+ vec![Field::new(
+ "list",
+ list.data_type().clone(),
+ list.is_nullable(),
+ )]
+ .into(),
+ HashMap::new(),
+ )?;
+
+ create_physical_expr(
+ &Expr::HigherOrderFunction(HigherOrderFunction::new(
+ array_transform,
+ vec![
+ col("list"),
+ lambda(
+ ["v"],
+ lit(100i32)
+ / Expr::LambdaVariable(LambdaVariable::new(
+ "v".to_string(),
+ Some(Arc::new(Field::new("v", DataType::Int32, true))),
+ )),
+ ),
+ ],
+ )),
+ &schema,
+ &ExecutionProps::new(),
+ )?
+ .evaluate(&RecordBatch::try_new(
+ Arc::clone(schema.inner()),
+ vec![Arc::new(list.clone())],
+ )?)?
+ .into_array(list.len())
+ }
+
+ #[test]
+ fn transform_on_sliced_list_should_not_evaluate_on_unreachable_values() {
+ let list = create_i32_list(
+ vec![
+ // Have 0 here so if the expression is called on data that it will fail
+ 0, 4, 100, 25, 20, 5, 2, 1, 10,
+ ],
+ OffsetBuffer::<i32>::from_lengths(vec![1, 3, 4, 1]),
+ None,
+ )
+ .slice(1, 3);
+
+ let res = divide_100_by(list).unwrap();
+
+ let actual_list = res.as_list::<i32>();
+
+ let expected_list = create_i32_list(
+ vec![25, 1, 4, 5, 20, 50, 100, 10],
+ OffsetBuffer::<i32>::from_lengths(vec![3, 4, 1]),
+ None,
+ );
+
+ assert_eq!(actual_list, &expected_list);
+ }
+
+ #[test]
+ fn transform_function_should_not_be_evaluated_on_values_underlying_null() {
+ let list = create_i32_list(
+ // 0 here for one of the values behind null, so if it will be evaluated
+ // it will fail due to divide by 0
+ vec![100, 20, 10, 0, 1, 2, 0, 1, 50],
+ OffsetBuffer::<i32>::from_lengths(vec![3, 4, 2]),
+ Some(NullBuffer::from(vec![true, false, true])),
+ );
+
+ let res = divide_100_by(list).unwrap();
+
+ let actual_list = res.as_list::<i32>();
+
+ let expected_list = create_i32_list(
+ vec![1, 5, 10, 100, 2],
+ OffsetBuffer::<i32>::from_lengths(vec![3, 0, 2]),
+ Some(NullBuffer::from(vec![true, false, true])),
+ );
+
+ assert_eq!(actual_list.data_type(), expected_list.data_type());
+ assert_eq!(actual_list, &expected_list);
+ }
+}
diff --git a/datafusion/functions-nested/src/lib.rs b/datafusion/functions-nested/src/lib.rs
index b1aa52d68c..96b0ef1557 100644
--- a/datafusion/functions-nested/src/lib.rs
+++ b/datafusion/functions-nested/src/lib.rs
@@ -37,8 +37,12 @@
#[macro_use]
pub mod macros;
+#[macro_use]
+pub mod macros_lambda;
+
pub mod array_compact;
pub mod array_has;
+pub mod array_transform;
pub mod arrays_zip;
pub mod cardinality;
pub mod concat;
@@ -73,7 +77,7 @@ pub mod utils;
use datafusion_common::Result;
use datafusion_execution::FunctionRegistry;
-use datafusion_expr::ScalarUDF;
+use datafusion_expr::{HigherOrderUDF, ScalarUDF};
use log::debug;
use std::sync::Arc;
@@ -83,6 +87,7 @@ pub mod expr_fn {
pub use super::array_has::array_has;
pub use super::array_has::array_has_all;
pub use super::array_has::array_has_any;
+ pub use super::array_transform::array_transform;
pub use super::arrays_zip::arrays_zip;
pub use super::cardinality::cardinality;
pub use super::concat::array_append;
@@ -184,6 +189,10 @@ pub fn all_default_nested_functions() -> Vec<Arc<ScalarUDF>> {
]
}
+pub fn all_default_higher_order_functions() -> Vec<Arc<dyn HigherOrderUDF>> {
+ vec![array_transform::array_transform_higher_order_function()]
+}
+
/// Registers all enabled packages with a [`FunctionRegistry`]
pub fn register_all(registry: &mut dyn FunctionRegistry) -> Result<()> {
let functions: Vec<Arc<ScalarUDF>> = all_default_nested_functions();
@@ -195,25 +204,43 @@ pub fn register_all(registry: &mut dyn FunctionRegistry) -> Result<()> {
Ok(()) as Result<()>
})?;
+ let functions: Vec<Arc<dyn HigherOrderUDF>> = all_default_higher_order_functions();
+ functions.into_iter().try_for_each(|function| {
+ let existing_function = registry.register_higher_order_function(function)?;
+ if let Some(existing_function) = existing_function {
+ debug!(
+ "Overwrite existing higher-order function: {}",
+ existing_function.name()
+ );
+ }
+ Ok(()) as Result<()>
+ })?;
+
Ok(())
}
#[cfg(test)]
mod tests {
- use crate::all_default_nested_functions;
+ use crate::{all_default_higher_order_functions, all_default_nested_functions};
use datafusion_common::Result;
use std::collections::HashSet;
#[test]
fn test_no_duplicate_name() -> Result<()> {
+ let scalars = all_default_nested_functions();
+ let scalars = scalars.iter().map(|s| (s.name(), s.aliases()));
+
+ let lambdas = all_default_higher_order_functions();
+ let lambdas = lambdas.iter().map(|l| (l.name(), l.aliases()));
+
let mut names = HashSet::new();
- for func in all_default_nested_functions() {
+
+ for (name, aliases) in scalars.chain(lambdas) {
assert!(
- names.insert(func.name().to_string().to_lowercase()),
- "duplicate function name: {}",
- func.name()
+ names.insert(name.to_string().to_lowercase()),
+ "duplicate function name: {name}",
);
- for alias in func.aliases() {
+ for alias in aliases {
assert!(
names.insert(alias.to_string().to_lowercase()),
"duplicate function name: {alias}"
diff --git a/datafusion/functions-nested/src/macros_lambda.rs b/datafusion/functions-nested/src/macros_lambda.rs
new file mode 100644
index 0000000000..8c15d8aed1
--- /dev/null
+++ b/datafusion/functions-nested/src/macros_lambda.rs
@@ -0,0 +1,107 @@
+// 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.
+
+/// Creates external API functions for an array UDF. Specifically, creates
+///
+/// 1. Single `HigherOrderUDF` instance
+///
+/// Creates a singleton `HigherOrderUDF` of the `$UDF` function named `STATIC_$(UDF)` and a
+/// function named `$HIGHER_ORDER_UDF_FUNC` which returns that function named `STATIC_$(UDF)`.
+///
+/// This is used to ensure creating the list of `HigherOrderUDF` only happens once.
+///
+/// # 2. `expr_fn` style function
+///
+/// These are functions that create an `Expr` that invokes the UDF, used
+/// primarily to programmatically create expressions.
+///
+/// For example:
+/// ```text
+/// pub fn array_to_string(delimiter: Expr) -> Expr {
+/// ...
+/// }
+/// ```
+/// # Arguments
+/// * `UDF`: name of the [`HigherOrderUDF`]
+/// * `EXPR_FN`: name of the expr_fn function to be created
+/// * `arg`: 0 or more named arguments for the function
+/// * `DOC`: documentation string for the function
+/// * `HIGHER_ORDER_UDF_FUNC`: name of the function to create (just) the `HigherOrderUDF`
+/// * (optional) `$CTOR`: Pass a custom constructor. When omitted it
+/// automatically resolves to `$UDF::new()`.
+///
+/// [`HigherOrderUDF`]: datafusion_expr::HigherOrderUDF
+macro_rules! make_higher_order_function_expr_and_func {
+ ($UDF:ident, $EXPR_FN:ident, $($arg:ident)*, $DOC:expr, $HIGHER_ORDER_UDF_FN:ident) => {
+ make_higher_order_function_expr_and_func!($UDF, $EXPR_FN, $($arg)*, $DOC, $HIGHER_ORDER_UDF_FN, $UDF::new);
+ };
+ ($UDF:ident, $EXPR_FN:ident, $($arg:ident)*, $DOC:expr, $HIGHER_ORDER_UDF_FN:ident, $CTOR:path) => {
+ // "fluent expr_fn" style function
+ #[doc = $DOC]
+ pub fn $EXPR_FN($($arg: datafusion_expr::Expr),*) -> datafusion_expr::Expr {
+ datafusion_expr::Expr::HigherOrderFunction(datafusion_expr::expr::HigherOrderFunction::new(
+ $HIGHER_ORDER_UDF_FN(),
+ vec![$($arg),*],
+ ))
+ }
+ create_higher_order!($UDF, $HIGHER_ORDER_UDF_FN, $CTOR);
+ };
+ ($UDF:ident, $EXPR_FN:ident, $DOC:expr, $HIGHER_ORDER_UDF_FN:ident) => {
+ make_higher_order_function_expr_and_func!($UDF, $EXPR_FN, $DOC, $HIGHER_ORDER_UDF_FN, $UDF::new);
+ };
+ ($UDF:ident, $EXPR_FN:ident, $DOC:expr, $HIGHER_ORDER_UDF_FN:ident, $CTOR:path) => {
+ // "fluent expr_fn" style function
+ #[doc = $DOC]
+ pub fn $EXPR_FN(arg: Vec<datafusion_expr::Expr>) -> datafusion_expr::Expr {
+ datafusion_expr::Expr::HigherOrderFunction(datafusion_expr::expr::HigherOrderFunction::new(
+ $HIGHER_ORDER_UDF_FN(),
+ arg,
+ ))
+ }
+ create_higher_order!($UDF, $HIGHER_ORDER_UDF_FN, $CTOR);
+ };
+}
+
+/// Creates a singleton `HigherOrderUDF` of the `$UDF` function named `STATIC_$(UDF)` and a
+/// function named `$HIGHER_ORDER_UDF_FUNC` which returns that function named `STATIC_$(UDF)`.
+///
+/// This is used to ensure creating the list of `HigherOrderUDF` only happens once.
+///
+/// # Arguments
+/// * `UDF`: name of the [`HigherOrderUDF`]
+/// * `HIGHER_ORDER_UDF_FUNC`: name of the function to create (just) the `HigherOrderUDF`
+/// * (optional) `$CTOR`: Pass a custom constructor. When omitted it
+/// automatically resolves to `$UDF::new()`.
+///
+/// [`HigherOrderUDF`]: datafusion_expr::HigherOrderUDF
+macro_rules! create_higher_order {
+ ($UDF:ident, $HIGHER_ORDER_UDF_FN:ident) => {
+ create_higher_order!($UDF, $HIGHER_ORDER_UDF_FN, $UDF::new);
+ };
+ ($UDF:ident, $HIGHER_ORDER_UDF_FN:ident, $CTOR:path) => {
+ #[doc = concat!("HigherOrderFunction that returns a [`HigherOrderUDF`](datafusion_expr::HigherOrderUDF) for ")]
+ #[doc = stringify!($UDF)]
+ pub fn $HIGHER_ORDER_UDF_FN() -> std::sync::Arc<dyn datafusion_expr::HigherOrderUDF> {
+ // Singleton instance of [`$UDF`], ensures the UDF is only created once
+ static INSTANCE: std::sync::LazyLock<std::sync::Arc<dyn datafusion_expr::HigherOrderUDF>> =
+ std::sync::LazyLock::new(|| {
+ std::sync::Arc::new($CTOR())
+ });
+ std::sync::Arc::clone(&INSTANCE)
+ }
+ };
+}
diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs
index bb17d13f5c..7b81feab47 100644
--- a/datafusion/optimizer/src/analyzer/type_coercion.rs
+++ b/datafusion/optimizer/src/analyzer/type_coercion.rs
@@ -35,8 +35,9 @@ use datafusion_common::{
plan_err,
};
use datafusion_expr::expr::{
- self, AggregateFunctionParams, Alias, Between, BinaryExpr, Case, Exists, InList,
- InSubquery, Like, ScalarFunction, SetComparison, Sort, WindowFunction,
+ self, AggregateFunctionParams, Alias, Between, BinaryExpr, Case, Exists,
+ HigherOrderFunction, InList, InSubquery, Like, ScalarFunction, SetComparison, Sort,
+ WindowFunction,
};
use datafusion_expr::expr_rewriter::coerce_plan_expr_for_schema;
use datafusion_expr::expr_schema::cast_subquery;
@@ -44,7 +45,9 @@ use datafusion_expr::logical_plan::Subquery;
use datafusion_expr::type_coercion::binary::{
comparison_coercion, like_coercion, type_union_coercion,
};
-use datafusion_expr::type_coercion::functions::{UDFCoercionExt, fields_with_udf};
+use datafusion_expr::type_coercion::functions::{
+ UDFCoercionExt, fields_with_udf, value_fields_with_higher_order_udf_and_lambdas,
+};
use datafusion_expr::type_coercion::other::{
get_coerce_type_for_case_expression, get_coerce_type_for_case_when,
get_coerce_type_for_list,
@@ -55,8 +58,8 @@ use datafusion_expr::type_coercion::{
use datafusion_expr::utils::merge_schema;
use datafusion_expr::{
Cast, Expr, ExprSchemable, Join, Limit, LogicalPlan, Operator, Projection, Union,
- WindowFrame, WindowFrameBound, WindowFrameUnits, is_false, is_not_false, is_not_true,
- is_not_unknown, is_true, is_unknown, lit, not,
+ ValueOrLambda, WindowFrame, WindowFrameBound, WindowFrameUnits, is_false,
+ is_not_false, is_not_true, is_not_unknown, is_true, is_unknown, lit, not,
};
/// Performs type coercion by determining the schema
@@ -763,6 +766,35 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> {
});
Ok(Transformed::yes(new_expr))
}
+ Expr::HigherOrderFunction(HigherOrderFunction { func, args }) => {
+ let current_fields = args
+ .iter()
+ .map(|arg| match arg {
+ Expr::Lambda(lambda) => Ok(ValueOrLambda::Lambda(
+ lambda.body.to_field(self.schema)?.1,
+ )),
+ _ => Ok(ValueOrLambda::Value(arg.to_field(self.schema)?.1)),
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ let new_fields = value_fields_with_higher_order_udf_and_lambdas(
+ ¤t_fields,
+ func.as_ref(),
+ )?;
+
+ let new_args = std::iter::zip(args, new_fields)
+ .map(|(arg, new_field)| match (&arg, new_field) {
+ (Expr::Lambda(_lambda), ValueOrLambda::Lambda(_)) => Ok(arg),
+ (Expr::Lambda(_lambda), ValueOrLambda::Value(_)) => internal_err!("value_fields_with_higher_order_udf returned a value for a lambda argument"),
+ (_, ValueOrLambda::Value(new_field)) => arg.cast_to(new_field.data_type(), self.schema),
+ (_, ValueOrLambda::Lambda(_)) => internal_err!("value_fields_with_higher_order_udf returned a lambda for a value argument"),
+ })
+ .collect::<Result<_>>()?;
+
+ Ok(Transformed::yes(Expr::HigherOrderFunction(
+ HigherOrderFunction::new(func, new_args),
+ )))
+ }
// TODO: remove the next line after `Expr::Wildcard` is removed
#[expect(deprecated)]
Expr::Alias(_)
@@ -777,7 +809,9 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> {
| Expr::Wildcard { .. }
| Expr::GroupingSet(_)
| Expr::Placeholder(_)
- | Expr::OuterReferenceColumn(_, _) => Ok(Transformed::no(expr)),
+ | Expr::OuterReferenceColumn(_, _)
+ | Expr::Lambda(_)
+ | Expr::LambdaVariable(_) => Ok(Transformed::no(expr)),
}
}
}
diff --git a/datafusion/optimizer/src/common_subexpr_eliminate.rs b/datafusion/optimizer/src/common_subexpr_eliminate.rs
index c02ba60247..2775d62144 100644
--- a/datafusion/optimizer/src/common_subexpr_eliminate.rs
+++ b/datafusion/optimizer/src/common_subexpr_eliminate.rs
@@ -30,7 +30,7 @@ use datafusion_common::alias::AliasGenerator;
use datafusion_common::cse::{CSE, CSEController, FoundCommonNodes};
use datafusion_common::tree_node::{Transformed, TreeNode};
use datafusion_common::{Column, DFSchema, DFSchemaRef, Result, qualified_name};
-use datafusion_expr::expr::{Alias, ScalarFunction};
+use datafusion_expr::expr::{Alias, HigherOrderFunction, ScalarFunction};
use datafusion_expr::logical_plan::{
Aggregate, Filter, LogicalPlan, Projection, Sort, Window,
};
@@ -651,12 +651,15 @@ impl CSEController for ExprCSEController<'_> {
fn conditional_children(node: &Expr) -> Option<(Vec<&Expr>, Vec<&Expr>)> {
match node {
- // In case of `ScalarFunction`s we don't know which children are surely
+ // In case of `ScalarFunction`s and `HigherOrderFunction`s we don't know which children are surely
// executed so start visiting all children conditionally and stop the
// recursion with `TreeNodeRecursion::Jump`.
Expr::ScalarFunction(ScalarFunction { func, args }) => {
func.conditional_arguments(args)
}
+ Expr::HigherOrderFunction(HigherOrderFunction { func, args }) => {
+ func.conditional_arguments(args)
+ }
// In case of `And` and `Or` the first child is surely executed, but we
// account subexpressions as conditional in the second.
@@ -697,6 +700,7 @@ impl CSEController for ExprCSEController<'_> {
fn is_valid(node: &Expr) -> bool {
!node.is_volatile_node()
+ && !matches!(node, Expr::Lambda(_) | Expr::LambdaVariable(_))
}
fn is_ignored(&self, node: &Expr) -> bool {
@@ -726,6 +730,8 @@ impl CSEController for ExprCSEController<'_> {
| Expr::ScalarVariable(..)
| Expr::Alias(..)
| Expr::Wildcard { .. }
+ | Expr::Lambda(_)
+ | Expr::LambdaVariable(_)
);
let is_aggr = matches!(node, Expr::AggregateFunction(..));
diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs
index 35b2e3e3b6..9c69276fa1 100644
--- a/datafusion/optimizer/src/push_down_filter.rs
+++ b/datafusion/optimizer/src/push_down_filter.rs
@@ -262,7 +262,10 @@ fn can_evaluate_as_join_condition(predicate: &Expr) -> Result<bool> {
| Expr::Cast(_)
| Expr::TryCast(_)
| Expr::InList { .. }
- | Expr::ScalarFunction(_) => Ok(TreeNodeRecursion::Continue),
+ | Expr::ScalarFunction(_)
+ | Expr::HigherOrderFunction(_)
+ | Expr::Lambda(_)
+ | Expr::LambdaVariable(_) => Ok(TreeNodeRecursion::Continue),
// TODO: remove the next line after `Expr::Wildcard` is removed
#[expect(deprecated)]
Expr::AggregateFunction(_)
diff --git a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs
index 1af08c91c1..143d8eae69 100644
--- a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs
+++ b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs
@@ -39,6 +39,7 @@ use datafusion_common::{
metadata::FieldMetadata,
tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRewriter},
};
+use datafusion_expr::expr::HigherOrderFunction;
use datafusion_expr::{
BinaryExpr, Case, ColumnarValue, Expr, ExprSchemable, Like, Operator, Volatility,
and, binary::BinaryTypeCoercer, lit, or, preimage::PreimageResult,
@@ -645,6 +646,9 @@ impl ConstEvaluator {
Expr::ScalarFunction(ScalarFunction { func, .. }) => {
Self::volatility_ok(func.signature().volatility)
}
+ Expr::HigherOrderFunction(HigherOrderFunction { func, .. }) => {
+ Self::volatility_ok(func.signature().volatility)
+ }
Expr::Cast(Cast { expr, field }) | Expr::TryCast(TryCast { expr, field }) => {
if let (
Ok(DataType::Struct(source_fields)),
@@ -691,7 +695,9 @@ impl ConstEvaluator {
| Expr::Like { .. }
| Expr::SimilarTo { .. }
| Expr::Case(_)
- | Expr::InList { .. } => true,
+ | Expr::InList { .. }
+ | Expr::Lambda(_)
+ | Expr::LambdaVariable(_) => true,
}
}
diff --git a/datafusion/optimizer/tests/optimizer_integration.rs b/datafusion/optimizer/tests/optimizer_integration.rs
index b1ba313e1d..14e02c8693 100644
--- a/datafusion/optimizer/tests/optimizer_integration.rs
+++ b/datafusion/optimizer/tests/optimizer_integration.rs
@@ -879,6 +879,13 @@ impl ContextProvider for MyContextProvider {
None
}
+ fn get_higher_order_meta(
+ &self,
+ _name: &str,
+ ) -> Option<Arc<dyn datafusion_expr::HigherOrderUDF>> {
+ None
+ }
+
fn get_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>> {
self.udafs.get(name).cloned()
}
@@ -907,6 +914,10 @@ impl ContextProvider for MyContextProvider {
Vec::new()
}
+ fn higher_order_function_names(&self) -> Vec<String> {
+ Vec::new()
+ }
+
fn udaf_names(&self) -> Vec<String> {
Vec::new()
}
diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs
index 25aedecc2c..bf95f8e6ac 100644
--- a/datafusion/physical-expr/src/expressions/case.rs
+++ b/datafusion/physical-expr/src/expressions/case.rs
@@ -19,7 +19,7 @@ mod literal_lookup_table;
use super::{Column, Literal};
use crate::PhysicalExpr;
-use crate::expressions::{lit, try_cast};
+use crate::expressions::{LambdaExpr, LambdaVariable, lit, try_cast};
use arrow::array::*;
use arrow::compute::kernels::zip::zip;
use arrow::compute::{
@@ -133,6 +133,13 @@ impl CaseBody {
expr.apply(|expr| {
if let Some(column) = expr.downcast_ref::<Column>() {
used_column_indices.insert(column.index());
+ } else if let Some(lambda_variable) =
+ expr.downcast_ref::<LambdaVariable>()
+ {
+ used_column_indices.insert(lambda_variable.index());
+ } else if expr.is::<LambdaExpr>() {
+ //todo: remove this branch when lambda supports column capture
+ return Ok(TreeNodeRecursion::Jump);
}
Ok(TreeNodeRecursion::Continue)
})
@@ -171,6 +178,20 @@ impl CaseBody {
projected,
))));
}
+ } else if let Some(lambda_variable) =
+ e.downcast_ref::<LambdaVariable>()
+ {
+ let original = lambda_variable.index();
+ let projected = *column_index_map.get(&original).unwrap();
+ if projected != original {
+ return Ok(Transformed::yes(Arc::new(LambdaVariable::new(
+ projected,
+ Arc::clone(lambda_variable.field()),
+ ))));
+ }
+ } else if e.is::<LambdaExpr>() {
+ //todo: remove this branch when lambda supports column capture
+ return Ok(Transformed::new(e, false, TreeNodeRecursion::Jump));
}
Ok(Transformed::no(e))
})
diff --git a/datafusion/physical-expr/src/expressions/lambda.rs b/datafusion/physical-expr/src/expressions/lambda.rs
new file mode 100644
index 0000000000..267f5f605f
--- /dev/null
+++ b/datafusion/physical-expr/src/expressions/lambda.rs
@@ -0,0 +1,160 @@
+// 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.
+
+//! Physical lambda expression: [`LambdaExpr`]
+
+use std::hash::Hash;
+use std::sync::Arc;
+
+use crate::physical_expr::PhysicalExpr;
+use arrow::{
+ datatypes::{DataType, Schema},
+ record_batch::RecordBatch,
+};
+use datafusion_common::plan_err;
+use datafusion_common::{HashSet, Result, internal_err};
+use datafusion_expr::ColumnarValue;
+
+/// Represents a lambda with the given parameters names and body
+#[derive(Debug, Eq, Clone)]
+pub struct LambdaExpr {
+ params: Vec<String>,
+ body: Arc<dyn PhysicalExpr>,
+}
+
+// Manually derive PartialEq and Hash to work around https://github.com/rust-lang/rust/issues/78808 [https://github.com/apache/datafusion/issues/13196]
+impl PartialEq for LambdaExpr {
+ fn eq(&self, other: &Self) -> bool {
+ self.params.eq(&other.params) && self.body.eq(&other.body)
+ }
+}
+
+impl Hash for LambdaExpr {
+ fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+ self.params.hash(state);
+ self.body.hash(state);
+ }
+}
+
+impl LambdaExpr {
+ /// Create a new lambda expression with the given parameters and body
+ pub fn try_new(params: Vec<String>, body: Arc<dyn PhysicalExpr>) -> Result<Self> {
+ if all_unique(¶ms) {
+ Ok(Self::new(params, body))
+ } else {
+ plan_err!("lambda params must be unique, got ({})", params.join(", "))
+ }
+ }
+
+ fn new(params: Vec<String>, body: Arc<dyn PhysicalExpr>) -> Self {
+ Self { params, body }
+ }
+
+ /// Get the lambda's params names
+ pub fn params(&self) -> &[String] {
+ &self.params
+ }
+
+ /// Get the lambda's body
+ pub fn body(&self) -> &Arc<dyn PhysicalExpr> {
+ &self.body
+ }
+}
+
+impl std::fmt::Display for LambdaExpr {
+ fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+ write!(f, "({}) -> {}", self.params.join(", "), self.body)
+ }
+}
+
+impl PhysicalExpr for LambdaExpr {
+ fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
+ Ok(DataType::Null)
+ }
+
+ fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
+ Ok(true)
+ }
+
+ fn evaluate(&self, _batch: &RecordBatch) -> Result<ColumnarValue> {
+ internal_err!("LambdaExpr::evaluate() should not be called")
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+ vec![&self.body]
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn PhysicalExpr>>,
+ ) -> Result<Arc<dyn PhysicalExpr>> {
+ let [body] = children.as_slice() else {
+ return internal_err!(
+ "LambdaExpr expects exactly 1 child, got {}",
+ children.len()
+ );
+ };
+
+ Ok(Arc::new(Self::new(self.params.clone(), Arc::clone(body))))
+ }
+
+ fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ write!(f, "({}) -> {}", self.params.join(", "), self.body)
+ }
+}
+
+/// Create a lambda expression
+pub fn lambda(
+ params: impl IntoIterator<Item = impl Into<String>>,
+ body: Arc<dyn PhysicalExpr>,
+) -> Result<Arc<dyn PhysicalExpr>> {
+ Ok(Arc::new(LambdaExpr::try_new(
+ params.into_iter().map(Into::into).collect(),
+ body,
+ )?))
+}
+
+fn all_unique(params: &[String]) -> bool {
+ match params.len() {
+ 0 | 1 => true,
+ 2 => params[0] != params[1],
+ _ => {
+ let mut set = HashSet::with_capacity(params.len());
+
+ params.iter().all(|p| set.insert(p.as_str()))
+ }
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use crate::expressions::{NoOp, lambda::lambda};
+ use arrow::{array::RecordBatch, datatypes::Schema};
+ use std::sync::Arc;
+
+ #[test]
+ fn test_lambda_evaluate() {
+ let lambda = lambda(["a"], Arc::new(NoOp::new())).unwrap();
+ let batch = RecordBatch::new_empty(Arc::new(Schema::empty()));
+ assert!(lambda.evaluate(&batch).is_err());
+ }
+
+ #[test]
+ fn test_lambda_duplicate_name() {
+ assert!(lambda(["a", "a"], Arc::new(NoOp::new())).is_err());
+ }
+}
diff --git a/datafusion/physical-expr/src/expressions/lambda_variable.rs b/datafusion/physical-expr/src/expressions/lambda_variable.rs
new file mode 100644
index 0000000000..fb9657e897
--- /dev/null
+++ b/datafusion/physical-expr/src/expressions/lambda_variable.rs
@@ -0,0 +1,145 @@
+// 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.
+
+//! Physical lambda variable reference: [`LambdaVariable`]
+
+use std::hash::Hash;
+use std::sync::Arc;
+
+use crate::physical_expr::PhysicalExpr;
+use arrow::datatypes::FieldRef;
+use arrow::{
+ datatypes::{DataType, Schema},
+ record_batch::RecordBatch,
+};
+
+use datafusion_common::{Result, exec_err, internal_err};
+use datafusion_expr::ColumnarValue;
+
+/// Represents the lambda variable with a given index and field
+#[derive(Debug, Clone)]
+pub struct LambdaVariable {
+ index: usize,
+ field: FieldRef,
+}
+
+impl Eq for LambdaVariable {}
+
+impl PartialEq for LambdaVariable {
+ fn eq(&self, other: &Self) -> bool {
+ self.index == other.index && self.field == other.field
+ }
+}
+
+impl Hash for LambdaVariable {
+ fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+ self.index.hash(state);
+ self.field.hash(state);
+ }
+}
+
+impl LambdaVariable {
+ /// Create a new lambda variable expression
+ pub fn new(index: usize, field: FieldRef) -> Self {
+ Self { index, field }
+ }
+
+ /// Get the variable's name
+ pub fn name(&self) -> &str {
+ self.field.name()
+ }
+
+ /// Get the variable's index
+ pub fn index(&self) -> usize {
+ self.index
+ }
+
+ /// Get the variable's field
+ pub fn field(&self) -> &FieldRef {
+ &self.field
+ }
+}
+
+impl std::fmt::Display for LambdaVariable {
+ fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+ write!(f, "{}@{}", self.name(), self.index)
+ }
+}
+
+impl PhysicalExpr for LambdaVariable {
+ fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
+ Ok(self.field.data_type().clone())
+ }
+
+ fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
+ Ok(self.field.is_nullable())
+ }
+
+ fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
+ if self.index >= batch.num_columns() {
+ return internal_err!(
+ "PhysicalExpr LambdaVariable references column '{}' at index {} (zero-based) but batch only has {} columns: {:?}",
+ self.name(),
+ self.index,
+ batch.num_columns(),
+ batch
+ .schema_ref()
+ .fields()
+ .iter()
+ .map(|f| f.name())
+ .collect::<Vec<_>>()
+ );
+ }
+
+ if self.field.as_ref() != batch.schema_ref().field(self.index) {
+ return exec_err!(
+ "Physical LambdaVariable field doesn't match batch field during evaluation {} != {}",
+ self.field,
+ batch.schema_ref().field(self.index)
+ );
+ }
+
+ Ok(ColumnarValue::Array(Arc::clone(batch.column(self.index))))
+ }
+
+ fn return_field(&self, _input_schema: &Schema) -> Result<FieldRef> {
+ Ok(Arc::clone(&self.field))
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+ vec![]
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ _children: Vec<Arc<dyn PhysicalExpr>>,
+ ) -> Result<Arc<dyn PhysicalExpr>> {
+ Ok(self)
+ }
+
+ fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ write!(f, "{}@{}", self.name(), self.index)
+ }
+}
+
+/// Create a lambda variable expression
+pub fn lambda_variable(name: &str, schema: &Schema) -> Result<Arc<dyn PhysicalExpr>> {
+ let index = schema.index_of(name)?;
+ let field = Arc::clone(&schema.fields()[index]);
+
+ Ok(Arc::new(LambdaVariable::new(index, field)))
+}
diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs
index 1e08235548..bb6d9ca9c9 100644
--- a/datafusion/physical-expr/src/expressions/mod.rs
+++ b/datafusion/physical-expr/src/expressions/mod.rs
@@ -26,6 +26,8 @@ mod dynamic_filters;
mod in_list;
mod is_not_null;
mod is_null;
+mod lambda;
+mod lambda_variable;
mod like;
mod literal;
mod negative;
@@ -47,6 +49,8 @@ pub use dynamic_filters::DynamicFilterPhysicalExpr;
pub use in_list::{InListExpr, in_list};
pub use is_not_null::{IsNotNullExpr, is_not_null};
pub use is_null::{IsNullExpr, is_null};
+pub use lambda::{LambdaExpr, lambda};
+pub use lambda_variable::{LambdaVariable, lambda_variable};
pub use like::{LikeExpr, like};
pub use literal::{Literal, lit};
pub use negative::{NegativeExpr, negative};
diff --git a/datafusion/physical-expr/src/higher_order_function.rs b/datafusion/physical-expr/src/higher_order_function.rs
new file mode 100644
index 0000000000..7cacf42d4c
--- /dev/null
+++ b/datafusion/physical-expr/src/higher_order_function.rs
@@ -0,0 +1,686 @@
+// 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.
+
+//! Declaration of built-in (higher order) functions.
+//! This module contains built-in functions' enumeration and metadata.
+//!
+//! Generally, a function has:
+//! * a signature
+//! * a return type, that is a function of the incoming argument's types
+//! * the computation, that must accept each valid signature
+//!
+//! * Signature: see `Signature`
+//! * Return type: a function `(arg_types) -> return_type`. E.g. for array_transform, ([[f32]], v -> v*2) -> [f32], ([[f32]], v -> v > 3.0) -> [bool].
+//!
+//! This module also has a set of coercion rules to improve user experience: if an argument i32 is passed
+//! to a function that supports f64, it is coerced to f64.
+
+use std::fmt::{self, Debug, Formatter};
+use std::hash::{Hash, Hasher};
+use std::sync::Arc;
+
+use crate::PhysicalExpr;
+use crate::expressions::{LambdaExpr, Literal};
+
+use arrow::array::{Array, RecordBatch};
+use arrow::datatypes::{DataType, FieldRef, Schema};
+use datafusion_common::config::{ConfigEntry, ConfigOptions};
+use datafusion_common::datatype::FieldExt;
+use datafusion_common::utils::remove_list_null_values;
+use datafusion_common::{
+ Result, ScalarValue, exec_datafusion_err, exec_err, internal_datafusion_err,
+ internal_err, plan_err,
+};
+use datafusion_expr::type_coercion::functions::value_fields_with_higher_order_udf;
+use datafusion_expr::{
+ ColumnarValue, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, HigherOrderUDF,
+ LambdaArgument, LambdaParametersProgress, ValueOrLambda, Volatility, expr_vec_fmt,
+};
+
+/// Physical expression of a higher order function
+pub struct HigherOrderFunctionExpr {
+ /// A shared instance of the higher-order function
+ fun: Arc<dyn HigherOrderUDF>,
+ /// The name of the higher-order function
+ name: String,
+ /// List of expressions to feed to the function as arguments
+ ///
+ /// For example, for `array_transform([2, 3], v -> v != 2)`, this will be:
+ ///
+ /// ```text
+ /// ListExpression [2,3]
+ /// LambdaExpression
+ /// parameters: ["v"]
+ /// body:
+ /// BinaryExpression (!=)
+ /// left:
+ /// LambdaVariableExpression("v", Field::new("", Int32, false))
+ /// right:
+ /// LiteralExpression(2)
+ /// ```
+ args: Vec<Arc<dyn PhysicalExpr>>,
+ /// Positions in `args` where lambdas were top level arguments during try_new_with_schema
+ /// but may have been wrapped by tree node rewrites via with_new_children, like:
+ ///
+ /// ```ignore
+ /// expr.transform(|expr| Ok(Transformed::yes(Arc::new(DebugExpr::new(expr)))))
+ /// ```
+ ///
+ /// For example, for `array_transform([2, 3], v -> v != 2)`, this will be `vec![1]`
+ lambda_positions: Vec<usize>,
+ /// The output field associated this expression
+ ///
+ /// For example, for `array_transform([2, 3], v -> v != 2)`, this will be
+ /// `Field::new("", DataType::new_list(DataType::Boolean, true), true)`
+ return_field: FieldRef,
+ /// The config options at execution time
+ config_options: Arc<ConfigOptions>,
+}
+
+impl Debug for HigherOrderFunctionExpr {
+ fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+ f.debug_struct("HigherOrderFunctionExpr")
+ .field("fun", &"<FUNC>")
+ .field("name", &self.name)
+ .field("args", &self.args)
+ .field("lambda_positions", &self.lambda_positions)
+ .field("return_field", &self.return_field)
+ .finish()
+ }
+}
+
+impl HigherOrderFunctionExpr {
+ /// Create a new Higher Order function
+ ///
+ /// Note that lambda arguments must be present directly in args as [LambdaExpr],
+ /// and not as a wrapped child of any arg
+ pub fn try_new_with_schema(
+ fun: Arc<dyn HigherOrderUDF>,
+ args: Vec<Arc<dyn PhysicalExpr>>,
+ schema: &Schema,
+ config_options: Arc<ConfigOptions>,
+ ) -> Result<Self> {
+ let name = fun.name().to_string();
+ let mut lambda_positions = vec![];
+ let arg_fields = args
+ .iter()
+ .enumerate()
+ .map(|(i, e)| match e.downcast_ref::<LambdaExpr>() {
+ Some(lambda) => {
+ lambda_positions.push(i);
+
+ Ok(ValueOrLambda::Lambda(lambda.body().return_field(schema)?))
+ }
+ None => Ok(ValueOrLambda::Value(e.return_field(schema)?)),
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ // verify that input data types is consistent with function's `HigherOrderTypeSignature`
+ value_fields_with_higher_order_udf(&arg_fields, fun.as_ref())?;
+
+ let arguments = args
+ .iter()
+ .map(|e| e.downcast_ref::<Literal>().map(|literal| literal.value()))
+ .collect::<Vec<_>>();
+
+ let ret_args = HigherOrderReturnFieldArgs {
+ arg_fields: &arg_fields,
+ scalar_arguments: &arguments,
+ };
+
+ let return_field = fun.return_field_from_args(ret_args)?;
+
+ Ok(Self {
+ fun,
+ name,
+ args,
+ lambda_positions,
+ return_field,
+ config_options,
+ })
+ }
+
+ /// Get the higher order function implementation
+ pub fn fun(&self) -> &dyn HigherOrderUDF {
+ self.fun.as_ref()
+ }
+
+ /// The name for this expression
+ pub fn name(&self) -> &str {
+ &self.name
+ }
+
+ /// Input arguments
+ pub fn args(&self) -> &[Arc<dyn PhysicalExpr>] {
+ &self.args
+ }
+
+ /// Data type produced by this expression
+ pub fn return_type(&self) -> &DataType {
+ self.return_field.data_type()
+ }
+
+ pub fn nullable(&self) -> bool {
+ self.return_field.is_nullable()
+ }
+
+ pub fn config_options(&self) -> &ConfigOptions {
+ &self.config_options
+ }
+
+ pub fn lambda_positions(&self) -> &[usize] {
+ &self.lambda_positions
+ }
+}
+
+impl fmt::Display for HigherOrderFunctionExpr {
+ fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+ write!(f, "{}({})", self.name, expr_vec_fmt!(self.args))
+ }
+}
+
+impl PartialEq for HigherOrderFunctionExpr {
+ fn eq(&self, o: &Self) -> bool {
+ if std::ptr::eq(self, o) {
+ // The equality implementation is somewhat expensive, so let's short-circuit when possible.
+ return true;
+ }
+ let Self {
+ fun,
+ name,
+ args,
+ lambda_positions,
+ return_field,
+ config_options,
+ } = self;
+ fun.eq(&o.fun)
+ && name.eq(&o.name)
+ && args.eq(&o.args)
+ && lambda_positions.eq(&o.lambda_positions)
+ && return_field.eq(&o.return_field)
+ && (Arc::ptr_eq(config_options, &o.config_options)
+ || sorted_config_entries(config_options)
+ == sorted_config_entries(&o.config_options))
+ }
+}
+impl Eq for HigherOrderFunctionExpr {}
+impl Hash for HigherOrderFunctionExpr {
+ fn hash<H: Hasher>(&self, state: &mut H) {
+ let Self {
+ fun,
+ name,
+ args,
+ lambda_positions,
+ return_field,
+ config_options: _, // expensive to hash, and often equal
+ } = self;
+ fun.hash(state);
+ name.hash(state);
+ args.hash(state);
+ lambda_positions.hash(state);
+ return_field.hash(state);
+ }
+}
+
+fn sorted_config_entries(config_options: &ConfigOptions) -> Vec<ConfigEntry> {
+ let mut entries = config_options.entries();
+ entries.sort_by(|l, r| l.key.cmp(&r.key));
+ entries
+}
+
+impl PhysicalExpr for HigherOrderFunctionExpr {
+ fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
+ let arg_fields = self
+ .args
+ .iter()
+ .enumerate()
+ .map(|(i, e)| {
+ if self.lambda_positions.contains(&i) {
+ let lambda = wrapped_lambda(e).ok_or_else(|| {
+ exec_datafusion_err!(
+ "{} unable to unwrap lambda from {e} at position {i}",
+ self.name()
+ )
+ })?;
+
+ Ok(ValueOrLambda::Lambda(
+ lambda.body().return_field(batch.schema_ref())?,
+ ))
+ } else {
+ Ok(ValueOrLambda::Value(e.return_field(batch.schema_ref())?))
+ }
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ let fields = arg_fields
+ .iter()
+ .map(|field| match field {
+ ValueOrLambda::Lambda(field) => {
+ ValueOrLambda::Lambda(Some(Arc::clone(field)))
+ }
+ ValueOrLambda::Value(field) => ValueOrLambda::Value(Arc::clone(field)),
+ })
+ .collect::<Vec<_>>();
+
+ // lambda_parameters refers only to lambdas and not to values, so instead
+ // of zipping it with self.args, we iterate over self.args and only
+ // consume from lambda_parameters when a given argument is a lambda
+ // to reconstruct the arguments list with the correct order
+ // this supports any value and lambda positioning including
+ // multiple lambdas interleaved with values
+ let mut lambda_parameters = match self.fun().lambda_parameters(0, &fields)? {
+ LambdaParametersProgress::Partial(_) => {
+ return plan_err!(
+ "{} lambda_parameters returned a partial result when the return type of all it's lambdas were provided",
+ self.name()
+ );
+ }
+ LambdaParametersProgress::Complete(items) => items.into_iter(),
+ };
+
+ let num_lambdas = self.args.len() - fields.len();
+
+ // functions can support multiple lambdas where some trailing ones are optional,
+ // but to simplify the implementor, lambda_parameters returns the parameters of all of them,
+ // so we can't do equality check. one example is spark reduce:
+ // https://spark.apache.org/docs/latest/api/sql/index.html#reduce
+ if lambda_parameters.len() < num_lambdas {
+ return exec_err!(
+ "{} invocation defined {num_lambdas} but lambda_parameters returned only {}",
+ self.name(),
+ lambda_parameters.len()
+ );
+ }
+
+ let args = self
+ .args
+ .iter()
+ .enumerate()
+ .map(|(i, arg)| {
+ if self.lambda_positions.contains(&i) {
+ let lambda = wrapped_lambda(arg).ok_or_else(|| {
+ exec_datafusion_err!(
+ "{} unable to unwrap lambda from {arg} at position {i}",
+ self.name()
+ )
+ })?;
+
+ let lambda_params = lambda_parameters.next().ok_or_else(|| {
+ internal_datafusion_err!(
+ "params len should have been checked above"
+ )
+ })?;
+
+ if lambda.params().len() > lambda_params.len() {
+ return exec_err!(
+ "lambda defined {} params but higher-order function support only {}",
+ lambda.params().len(),
+ lambda_params.len()
+ );
+ }
+
+ let params = std::iter::zip(lambda.params(), lambda_params)
+ .map(|(name, param)| param.renamed(name.as_str()))
+ .collect();
+
+ Ok(ValueOrLambda::Lambda(LambdaArgument::new(
+ params,
+ Arc::clone(lambda.body()),
+ )))
+ } else {
+ let value = arg.evaluate(batch)?;
+
+ let value =
+ if self.fun.clear_null_values() && matches!(value.data_type(), DataType::List(_) | DataType::LargeList(_)) {
+ ColumnarValue::Array(remove_list_null_values(
+ &value.into_array(batch.num_rows())?,
+ )?)
+ } else {
+ value
+ };
+
+ Ok(ValueOrLambda::Value(value))
+ }
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ let input_empty = args.is_empty();
+ let input_all_scalar = args
+ .iter()
+ .all(|arg| matches!(arg, ValueOrLambda::Value(ColumnarValue::Scalar(_))));
+
+ // evaluate the function
+ let output = self.fun.invoke_with_args(HigherOrderFunctionArgs {
+ args,
+ arg_fields,
+ number_rows: batch.num_rows(),
+ return_field: Arc::clone(&self.return_field),
+ config_options: Arc::clone(&self.config_options),
+ })?;
+
+ if let ColumnarValue::Array(array) = &output
+ && array.len() != batch.num_rows()
+ {
+ // If the arguments are a non-empty slice of scalar values, we can assume that
+ // returning a one-element array is equivalent to returning a scalar.
+ let preserve_scalar = array.len() == 1 && !input_empty && input_all_scalar;
+ return if preserve_scalar {
+ ScalarValue::try_from_array(array, 0).map(ColumnarValue::Scalar)
+ } else {
+ internal_err!(
+ "higher-order function {} returned a different number of rows than expected. Expected: {}, Got: {}",
+ self.name,
+ batch.num_rows(),
+ array.len()
+ )
+ };
+ }
+ Ok(output)
+ }
+
+ fn return_field(&self, _input_schema: &Schema) -> Result<FieldRef> {
+ Ok(Arc::clone(&self.return_field))
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+ self.args.iter().collect()
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn PhysicalExpr>>,
+ ) -> Result<Arc<dyn PhysicalExpr>> {
+ if children.len() != self.args.len() {
+ return internal_err!(
+ "HigherOrderFunctionExpr expects exactly {} child, got {}",
+ self.args.len(),
+ children.len()
+ );
+ }
+
+ for (i, child) in children.iter().enumerate() {
+ if self.lambda_positions.contains(&i) {
+ if wrapped_lambda(child).is_none() {
+ return plan_err!(
+ "{} unable to unwrap lambda from {} at position {i}",
+ &children[i],
+ self.name()
+ );
+ }
+ } else if child.is::<LambdaExpr>() {
+ return plan_err!(
+ "{} received a lambda via with_new_children at position {i} that wasn't a lambda before",
+ self.name()
+ );
+ }
+ }
+
+ Ok(Arc::new(HigherOrderFunctionExpr {
+ name: self.name.clone(),
+ fun: Arc::clone(&self.fun),
+ args: children,
+ lambda_positions: self.lambda_positions.clone(),
+ return_field: Arc::clone(&self.return_field),
+ config_options: Arc::clone(&self.config_options),
+ }))
+ }
+
+ fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result {
+ write!(f, "{}(", self.name)?;
+ for (i, expr) in self.args.iter().enumerate() {
+ if i > 0 {
+ write!(f, ", ")?;
+ }
+ expr.fmt_sql(f)?;
+ }
+ write!(f, ")")
+ }
+
+ fn is_volatile_node(&self) -> bool {
+ self.fun.signature().volatility == Volatility::Volatile
+ }
+}
+
+fn wrapped_lambda(expr: &Arc<dyn PhysicalExpr>) -> Option<&LambdaExpr> {
+ let mut current = expr;
+
+ loop {
+ if let Some(lambda) = current.downcast_ref::<LambdaExpr>() {
+ return Some(lambda);
+ } else if current.is::<HigherOrderFunctionExpr>() {
+ return None;
+ }
+
+ match current.children().as_slice() {
+ [single_child] => current = *single_child,
+ _ => return None,
+ }
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use super::*;
+ use crate::HigherOrderFunctionExpr;
+ use crate::expressions::Column;
+ use crate::expressions::NoOp;
+ use crate::expressions::lambda;
+ use crate::expressions::not;
+ use arrow::array::NullArray;
+ use arrow::array::RecordBatchOptions;
+ use arrow::datatypes::{DataType, Field, Schema};
+ use datafusion_common::Result;
+ use datafusion_common::assert_contains;
+ use datafusion_expr::{
+ HigherOrderFunctionArgs, HigherOrderSignature, HigherOrderUDF,
+ };
+ use datafusion_expr_common::columnar_value::ColumnarValue;
+ use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
+ use datafusion_physical_expr_common::physical_expr::is_volatile;
+
+ /// Test helper to create a mock UDF with a specific volatility
+ #[derive(Debug, PartialEq, Eq, Hash)]
+ struct MockHigherOrderUDF {
+ signature: HigherOrderSignature,
+ }
+
+ impl HigherOrderUDF for MockHigherOrderUDF {
+ fn name(&self) -> &str {
+ "mock_function"
+ }
+
+ fn signature(&self) -> &HigherOrderSignature {
+ &self.signature
+ }
+
+ fn lambda_parameters(
+ &self,
+ _step: usize,
+ _fields: &[ValueOrLambda<FieldRef, Option<FieldRef>>],
+ ) -> Result<LambdaParametersProgress> {
+ Ok(LambdaParametersProgress::Complete(vec![vec![Arc::new(
+ Field::new("", DataType::Null, true),
+ )]]))
+ }
+
+ fn return_field_from_args(
+ &self,
+ args: HigherOrderReturnFieldArgs,
+ ) -> Result<FieldRef> {
+ match &args.arg_fields[0] {
+ ValueOrLambda::Lambda(field) | ValueOrLambda::Value(field) => {
+ Ok(Arc::clone(field))
+ }
+ }
+ }
+
+ fn invoke_with_args(
+ &self,
+ args: HigherOrderFunctionArgs,
+ ) -> Result<ColumnarValue> {
+ match &args.args[0] {
+ ValueOrLambda::Lambda(lambda) => {
+ lambda.evaluate(&[&|| Ok(Arc::new(NullArray::new(args.number_rows)))])
+ }
+ ValueOrLambda::Value(value) => Ok(value.clone()),
+ }
+ }
+ }
+
+ #[test]
+ fn test_higher_order_function_volatile_node() {
+ // Create a volatile UDF
+ let volatile_udf = Arc::new(MockHigherOrderUDF {
+ signature: HigherOrderSignature::variadic_any(Volatility::Volatile),
+ });
+
+ // Create a non-volatile UDF
+ let stable_udf = Arc::new(MockHigherOrderUDF {
+ signature: HigherOrderSignature::variadic_any(Volatility::Stable),
+ });
+
+ let schema = Schema::new(vec![Field::new("a", DataType::Float32, false)]);
+ let args = vec![Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>];
+ let config_options = Arc::new(ConfigOptions::new());
+
+ // Test volatile function
+ let volatile_expr = HigherOrderFunctionExpr::try_new_with_schema(
+ volatile_udf,
+ args.clone(),
+ &schema,
+ Arc::clone(&config_options),
+ )
+ .unwrap();
+
+ assert!(volatile_expr.is_volatile_node());
+ let volatile_arc: Arc<dyn PhysicalExpr> = Arc::new(volatile_expr);
+ assert!(is_volatile(&volatile_arc));
+
+ // Test non-volatile function
+ let stable_expr = HigherOrderFunctionExpr::try_new_with_schema(
+ stable_udf,
+ args,
+ &schema,
+ config_options,
+ )
+ .unwrap();
+
+ assert!(!stable_expr.is_volatile_node());
+ let stable_arc: Arc<dyn PhysicalExpr> = Arc::new(stable_expr);
+ assert!(!is_volatile(&stable_arc));
+ }
+
+ #[test]
+ fn test_higher_order_function_wrapped_lambda() {
+ let fun = Arc::new(MockHigherOrderUDF {
+ signature: HigherOrderSignature::variadic_any(Volatility::Stable),
+ });
+
+ let expected = ScalarValue::Int32(Some(42));
+
+ let hof = HigherOrderFunctionExpr::try_new_with_schema(
+ fun,
+ vec![lambda(["a"], Arc::new(Literal::new(expected.clone()))).unwrap()],
+ &Schema::empty(),
+ Arc::new(ConfigOptions::new()),
+ )
+ .unwrap();
+
+ let new_children = vec![not(Arc::clone(&hof.args[0])).unwrap()];
+ let wrapped = Arc::new(hof).with_new_children(new_children).unwrap();
+
+ let result = wrapped
+ .evaluate(
+ &RecordBatch::try_new_with_options(
+ Arc::new(Schema::empty()),
+ vec![],
+ &RecordBatchOptions::new().with_row_count(Some(0)),
+ )
+ .unwrap(),
+ )
+ .unwrap();
+
+ let ColumnarValue::Scalar(result) = result else {
+ unreachable!()
+ };
+
+ assert_eq!(result, expected);
+ }
+
+ #[test]
+ fn test_higher_order_function_badly_wrapped_lambda() {
+ let fun = Arc::new(MockHigherOrderUDF {
+ signature: HigherOrderSignature::variadic_any(Volatility::Stable),
+ });
+
+ let hof = HigherOrderFunctionExpr::try_new_with_schema(
+ fun,
+ vec![
+ not(
+ lambda(["a"], Arc::new(Literal::new(ScalarValue::Int32(Some(42)))))
+ .unwrap(),
+ )
+ .unwrap(),
+ ],
+ &Schema::empty(),
+ Arc::new(ConfigOptions::new()),
+ )
+ .unwrap();
+
+ let result = hof
+ .evaluate(
+ &RecordBatch::try_new_with_options(
+ Arc::new(Schema::empty()),
+ vec![],
+ &RecordBatchOptions::new().with_row_count(Some(0)),
+ )
+ .unwrap(),
+ )
+ .unwrap_err();
+
+ assert_contains!(
+ result.to_string(),
+ "LambdaExpr::evaluate() should not be called"
+ );
+ }
+
+ #[test]
+ fn test_higher_order_function_unexpected_lambda() {
+ let fun = Arc::new(MockHigherOrderUDF {
+ signature: HigherOrderSignature::variadic_any(Volatility::Stable),
+ });
+
+ let hof = HigherOrderFunctionExpr::try_new_with_schema(
+ fun,
+ vec![Arc::new(NoOp::new())],
+ &Schema::empty(),
+ Arc::new(ConfigOptions::new()),
+ )
+ .unwrap();
+
+ let result = Arc::new(hof)
+ .with_new_children(vec![lambda(["a"], Arc::new(NoOp::new())).unwrap()])
+ .unwrap_err();
+
+ assert_contains!(
+ result.to_string(),
+ "mock_function received a lambda via with_new_children at position 0 that wasn't a lambda before"
+ );
+ }
+}
diff --git a/datafusion/physical-expr/src/lib.rs b/datafusion/physical-expr/src/lib.rs
index 9c567ce862..848bf81d15 100644
--- a/datafusion/physical-expr/src/lib.rs
+++ b/datafusion/physical-expr/src/lib.rs
@@ -34,6 +34,7 @@ pub mod binary_map {
pub mod async_scalar_function;
pub mod equivalence;
pub mod expressions;
+pub mod higher_order_function;
pub mod intervals;
mod partitioning;
mod physical_expr;
@@ -70,6 +71,7 @@ pub use datafusion_physical_expr_common::sort_expr::{
PhysicalSortRequirement,
};
+pub use higher_order_function::HigherOrderFunctionExpr;
pub use planner::{create_physical_expr, create_physical_exprs};
pub use scalar_function::ScalarFunctionExpr;
pub use simplifier::PhysicalExprSimplifier;
diff --git a/datafusion/physical-expr/src/planner.rs b/datafusion/physical-expr/src/planner.rs
index 5d4c5d4ce1..9cb20de252 100644
--- a/datafusion/physical-expr/src/planner.rs
+++ b/datafusion/physical-expr/src/planner.rs
@@ -15,10 +15,11 @@
// specific language governing permissions and limitations
// under the License.
+use std::collections::HashMap;
use std::sync::Arc;
-use crate::ScalarFunctionExpr;
use crate::scalar_subquery::ScalarSubqueryExpr;
+use crate::{HigherOrderFunctionExpr, ScalarFunctionExpr};
use crate::{
PhysicalExpr,
expressions::{self, Column, Literal, binary, like, similar_to},
@@ -26,12 +27,17 @@ use crate::{
use arrow::datatypes::Schema;
use datafusion_common::config::ConfigOptions;
+use datafusion_common::datatype::FieldExt;
use datafusion_common::metadata::{FieldMetadata, format_type_and_metadata};
use datafusion_common::{
- DFSchema, Result, ScalarValue, ToDFSchema, exec_err, not_impl_err, plan_err,
+ DFSchema, Result, ScalarValue, ToDFSchema, exec_err, internal_datafusion_err,
+ not_impl_err, plan_datafusion_err, plan_err,
};
use datafusion_expr::execution_props::ExecutionProps;
-use datafusion_expr::expr::{Alias, Cast, InList, Placeholder, ScalarFunction};
+use datafusion_expr::expr::{
+ Alias, Cast, HigherOrderFunction, InList, Lambda, LambdaVariable, Placeholder,
+ ScalarFunction,
+};
use datafusion_expr::var_provider::VarType;
use datafusion_expr::var_provider::is_system_variables;
use datafusion_expr::{
@@ -429,6 +435,107 @@ pub fn create_physical_expr(
Expr::Placeholder(Placeholder { id, .. }) => {
exec_err!("Placeholder '{id}' was not provided a value for execution.")
}
+ Expr::HigherOrderFunction(invocation @ HigherOrderFunction { func, args }) => {
+ let num_lambdas = args
+ .iter()
+ .filter(|arg| matches!(arg, Expr::Lambda(_)))
+ .count();
+
+ let mut lambda_parameters =
+ invocation.lambda_parameters(input_dfschema)?.into_iter();
+
+ if num_lambdas > lambda_parameters.len() {
+ return plan_err!(
+ "{} lambda_parameters returned only {} values for {num_lambdas} lambdas",
+ func.name(),
+ lambda_parameters.len()
+ );
+ }
+
+ let physical_args = args
+ .iter()
+ .map(|arg| match arg {
+ Expr::Lambda(lambda) => {
+ let lambda_parameters = lambda_parameters
+ .next()
+ .ok_or_else(|| {
+ internal_datafusion_err!(
+ "lambda_parameters len should have been checked above"
+ )
+ })?
+ .into_iter()
+ .zip(&lambda.params)
+ .map(|(field, name)| field.renamed(name.as_str()))
+ .collect();
+
+ let lambda_schema = DFSchema::from_unqualified_fields(
+ lambda_parameters,
+ HashMap::new(),
+ )?;
+
+ create_physical_expr(arg, &lambda_schema, execution_props)
+ }
+ _ => create_physical_expr(arg, input_dfschema, execution_props),
+ })
+ .collect::<Result<_>>()?;
+
+ let config_options = match execution_props.config_options.as_ref() {
+ Some(config_options) => Arc::clone(config_options),
+ None => Arc::new(ConfigOptions::default()),
+ };
+
+ Ok(Arc::new(HigherOrderFunctionExpr::try_new_with_schema(
+ Arc::clone(func),
+ physical_args,
+ input_schema,
+ config_options,
+ )?))
+ }
+ Expr::Lambda(Lambda { params, body }) => {
+ // tracked at https://github.com/apache/datafusion/issues/21172
+ if body.any_column_refs() {
+ return plan_err!("lambda doesn't support column capture");
+ }
+
+ expressions::lambda(
+ params,
+ create_physical_expr(body, input_dfschema, execution_props)?,
+ )
+ }
+ Expr::LambdaVariable(LambdaVariable {
+ name,
+ field,
+ spans: _,
+ }) => {
+ let field = field.as_ref().ok_or_else(|| {
+ plan_datafusion_err!("unresolved LambdaVariable {name}")
+ })?;
+
+ let index = input_dfschema.inner().index_of(name)?;
+ let schema_field = input_dfschema.field(index);
+
+ // LambdaVariable.field will be made optional as in Expr::Placeholder
+ // and only LambdaVariable.name used, and field.name ignored,
+ // so they're not enforced to match for logical expressions
+ // Rename the field to match the schema one and use it's PartialEq impl instead
+ // of checking property by property and fail if new properties get's added to it.
+ // While not necessary, the sql planner does create lambda vars with matching names,
+ // so this shouldn't allocate with a lambda var from it
+ let renamed_field = Arc::clone(field).renamed(name);
+
+ if &renamed_field != schema_field {
+ return plan_err!(
+ "LambdaVariable field and schema field mismatch {} != {}",
+ renamed_field,
+ schema_field
+ );
+ }
+
+ Ok(Arc::new(expressions::LambdaVariable::new(
+ index,
+ Arc::clone(schema_field),
+ )))
+ }
other => {
not_impl_err!("Physical plan does not support logical expression {other:?}")
}
diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs
index 1417ba850b..7ae5cbeed3 100644
--- a/datafusion/proto/src/logical_plan/mod.rs
+++ b/datafusion/proto/src/logical_plan/mod.rs
@@ -55,7 +55,8 @@ use datafusion_datasource_json::file_format::{
#[cfg(feature = "parquet")]
use datafusion_datasource_parquet::file_format::{ParquetFormat, ParquetFormatFactory};
use datafusion_expr::{
- AggregateUDF, DmlStatement, FetchType, RecursiveQuery, SkipType, TableSource, Unnest,
+ AggregateUDF, DmlStatement, FetchType, HigherOrderUDF, RecursiveQuery, SkipType,
+ TableSource, Unnest,
};
use datafusion_expr::{
DistinctOn, DropView, Expr, LogicalPlan, LogicalPlanBuilder, ScalarUDF, SortExpr,
@@ -156,6 +157,24 @@ pub trait LogicalExtensionCodec: Debug + Send + Sync + std::any::Any {
Ok(())
}
+ fn try_decode_higher_order_function(
+ &self,
+ name: &str,
+ _buf: &[u8],
+ ) -> Result<Arc<dyn HigherOrderUDF>> {
+ not_impl_err!(
+ "LogicalExtensionCodec is not provided for higher order function {name}"
+ )
+ }
+
+ fn try_encode_higher_order_function(
+ &self,
+ _node: &dyn HigherOrderUDF,
+ _buf: &mut Vec<u8>,
+ ) -> Result<()> {
+ Ok(())
+ }
+
fn try_decode_udaf(&self, name: &str, _buf: &[u8]) -> Result<Arc<AggregateUDF>> {
not_impl_err!(
"LogicalExtensionCodec is not provided for aggregate function {name}"
diff --git a/datafusion/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs
index bd5c4b585c..d79107d1d0 100644
--- a/datafusion/proto/src/logical_plan/to_proto.rs
+++ b/datafusion/proto/src/logical_plan/to_proto.rs
@@ -634,6 +634,11 @@ pub fn serialize_expr(
.unwrap_or(HashMap::new()),
})),
},
+ Expr::HigherOrderFunction(_) | Expr::Lambda(_) | Expr::LambdaVariable(_) => {
+ return Err(Error::General(
+ "Proto serialization error: Lambda not implemented".to_string(),
+ ));
+ }
};
Ok(expr_node)
diff --git a/datafusion/session/src/session.rs b/datafusion/session/src/session.rs
index 34506a0a17..82dda6655f 100644
--- a/datafusion/session/src/session.rs
+++ b/datafusion/session/src/session.rs
@@ -23,7 +23,9 @@ use datafusion_execution::config::SessionConfig;
use datafusion_execution::runtime_env::RuntimeEnv;
use datafusion_expr::execution_props::ExecutionProps;
use datafusion_expr::registry::ExtensionTypeRegistryRef;
-use datafusion_expr::{AggregateUDF, Expr, LogicalPlan, ScalarUDF, WindowUDF};
+use datafusion_expr::{
+ AggregateUDF, Expr, HigherOrderUDF, LogicalPlan, ScalarUDF, WindowUDF,
+};
use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr};
use parking_lot::{Mutex, RwLock};
use std::any::Any;
@@ -111,6 +113,9 @@ pub trait Session: Send + Sync {
/// Return reference to scalar_functions
fn scalar_functions(&self) -> &HashMap<String, Arc<ScalarUDF>>;
+ /// Return reference to higher_order_functions
+ fn higher_order_functions(&self) -> &HashMap<String, Arc<dyn HigherOrderUDF>>;
+
/// Return reference to aggregate_functions
fn aggregate_functions(&self) -> &HashMap<String, Arc<AggregateUDF>>;
@@ -153,6 +158,7 @@ impl From<&dyn Session> for TaskContext {
state.session_id().to_string(),
state.config().clone(),
state.scalar_functions().clone(),
+ state.higher_order_functions().clone(),
state.aggregate_functions().clone(),
state.window_functions().clone(),
Arc::clone(state.runtime_env()),
diff --git a/datafusion/spark/src/lib.rs b/datafusion/spark/src/lib.rs
index 9575f560b8..2eee94c52e 100644
--- a/datafusion/spark/src/lib.rs
+++ b/datafusion/spark/src/lib.rs
@@ -43,7 +43,7 @@
//!
//! ```
//! # use datafusion_execution::FunctionRegistry;
-//! # use datafusion_expr::{ScalarUDF, AggregateUDF, WindowUDF};
+//! # use datafusion_expr::{ScalarUDF, AggregateUDF, WindowUDF, HigherOrderUDF};
//! # use datafusion_expr::planner::ExprPlanner;
//! # use datafusion_common::Result;
//! # use std::collections::HashSet;
@@ -55,9 +55,11 @@
//! # impl FunctionRegistry for SessionContext {
//! # fn register_udf(&mut self, _udf: Arc<ScalarUDF>) -> Result<Option<Arc<ScalarUDF>>> { Ok (None) }
//! # fn udfs(&self) -> HashSet<String> { unimplemented!() }
+//! # fn higher_order_function_names(&self) -> HashSet<String> { unimplemented!() }
//! # fn udafs(&self) -> HashSet<String> { unimplemented!() }
//! # fn udwfs(&self) -> HashSet<String> { unimplemented!() }
//! # fn udf(&self, _name: &str) -> Result<Arc<ScalarUDF>> { unimplemented!() }
+//! # fn higher_order_function(&self, name: &str) -> Result<Arc<dyn HigherOrderUDF>> { unimplemented!() }
//! # fn udaf(&self, name: &str) -> Result<Arc<AggregateUDF>> {unimplemented!() }
//! # fn udwf(&self, name: &str) -> Result<Arc<WindowUDF>> { unimplemented!() }
//! # fn expr_planners(&self) -> Vec<Arc<dyn ExprPlanner>> { unimplemented!() }
diff --git a/datafusion/sql/examples/sql.rs b/datafusion/sql/examples/sql.rs
index dbedaf3f15..dc49b4460f 100644
--- a/datafusion/sql/examples/sql.rs
+++ b/datafusion/sql/examples/sql.rs
@@ -21,11 +21,11 @@ use arrow::datatypes::{DataType, Field, Schema};
use datafusion_common::config::ConfigOptions;
use datafusion_common::{Result, TableReference, plan_err};
-use datafusion_expr::WindowUDF;
use datafusion_expr::planner::ExprPlanner;
use datafusion_expr::{
AggregateUDF, ScalarUDF, TableSource, logical_plan::builder::LogicalTableSource,
};
+use datafusion_expr::{HigherOrderUDF, WindowUDF};
use datafusion_functions::core::planner::CoreFunctionPlanner;
use datafusion_functions_aggregate::count::count_udaf;
use datafusion_functions_aggregate::sum::sum_udaf;
@@ -138,6 +138,10 @@ impl ContextProvider for MyContextProvider {
None
}
+ fn get_higher_order_meta(&self, _name: &str) -> Option<Arc<dyn HigherOrderUDF>> {
+ None
+ }
+
fn get_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>> {
self.udafs.get(name).cloned()
}
@@ -158,6 +162,10 @@ impl ContextProvider for MyContextProvider {
Vec::new()
}
+ fn higher_order_function_names(&self) -> Vec<String> {
+ Vec::new()
+ }
+
fn udaf_names(&self) -> Vec<String> {
Vec::new()
}
diff --git a/datafusion/sql/src/expr/function.rs b/datafusion/sql/src/expr/function.rs
index b7f4ebe9a6..1790e66a02 100644
--- a/datafusion/sql/src/expr/function.rs
+++ b/datafusion/sql/src/expr/function.rs
@@ -17,22 +17,26 @@
use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
-use arrow::datatypes::DataType;
+use arrow::datatypes::{DataType, FieldRef};
use datafusion_common::{
- DFSchema, Dependency, Diagnostic, Result, Span, internal_datafusion_err,
- internal_err, not_impl_err, plan_datafusion_err, plan_err,
+ DFSchema, Dependency, Diagnostic, HashSet, Result, Span, datatype::FieldExt,
+ internal_datafusion_err, internal_err, not_impl_err, plan_datafusion_err, plan_err,
};
use datafusion_expr::{
- Expr, ExprSchemable, SortExpr, WindowFrame, WindowFunctionDefinition,
+ Expr, ExprSchemable, LambdaParametersProgress, SortExpr, ValueOrLambda, WindowFrame,
+ WindowFunctionDefinition,
arguments::ArgumentName,
- expr,
- expr::{NullTreatment, ScalarFunction, Unnest, WildcardOptions, WindowFunction},
+ expr::{
+ self, HigherOrderFunction, Lambda, NullTreatment, ScalarFunction, Unnest,
+ WildcardOptions, WindowFunction,
+ },
planner::{PlannerResult, RawAggregateExpr, RawWindowExpr},
+ type_coercion::functions::value_fields_with_higher_order_udf,
};
use sqlparser::ast::{
DuplicateTreatment, Expr as SQLExpr, Function as SQLFunction, FunctionArg,
FunctionArgExpr, FunctionArgumentClause, FunctionArgumentList, FunctionArguments,
- ObjectName, OrderByExpr, Spanned, WindowType,
+ LambdaFunction, ObjectName, OrderByExpr, Spanned, WindowType,
};
/// Suggest a valid function based on an invalid input function name
@@ -57,6 +61,7 @@ pub fn suggest_valid_function(
let mut funcs = Vec::new();
funcs.extend(ctx.udf_names());
+ funcs.extend(ctx.higher_order_function_names());
funcs.extend(ctx.udaf_names());
funcs
@@ -363,6 +368,184 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
}
}
+ if let Some(fm) = self.context_provider.get_higher_order_meta(&name) {
+ // plan non-lambda arguments first so we can get theirs datatype and call
+ // HigherOrderUDF::lambda_parameters to then plan the lambda arguments with
+ // resolved lambda variables
+ enum ExprOrLambda {
+ Expr(Expr),
+ Lambda(LambdaFunction),
+ }
+
+ let partially_planned = args
+ .into_iter()
+ .map(|a| match a {
+ FunctionArg::Unnamed(FunctionArgExpr::Expr(SQLExpr::Lambda(
+ lambda,
+ ))) => {
+ if !all_unique(&lambda.params) {
+ return plan_err!(
+ "lambda parameters names must be unique, got {}",
+ lambda.params
+ );
+ }
+
+ Ok(ExprOrLambda::Lambda(lambda))
+ }
+ _ => Ok(ExprOrLambda::Expr(self.sql_fn_arg_to_logical_expr(
+ a,
+ schema,
+ planner_context,
+ )?)),
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ let current_fields = partially_planned
+ .iter()
+ .map(|e| match e {
+ ExprOrLambda::Expr(expr) => {
+ Ok(ValueOrLambda::Value(expr.to_field(schema)?.1))
+ }
+ ExprOrLambda::Lambda(_lambda_function) => {
+ Ok(ValueOrLambda::Lambda(()))
+ }
+ })
+ .collect::<Result<Vec<_>>>()?;
+
+ // coerce fields because coercion may alter the lambda parameters
+ let mut fields =
+ value_fields_with_higher_order_udf(¤t_fields, fm.as_ref())?
+ .into_iter()
+ .map(|arg| match arg {
+ ValueOrLambda::Value(value) => ValueOrLambda::Value(value),
+ ValueOrLambda::Lambda(_lambda) => ValueOrLambda::Lambda(None),
+ })
+ .collect::<Vec<_>>();
+
+ let lambda_count = partially_planned
+ .iter()
+ .filter(|a| matches!(a, ExprOrLambda::Lambda(_)))
+ .count();
+
+ let mut step = 0;
+
+ let args = loop {
+ match fm.lambda_parameters(step, &fields)? {
+ LambdaParametersProgress::Partial(params) => {
+ let mut params = params.into_iter();
+
+ if params.len() != lambda_count {
+ return plan_err!(
+ "{} lambda_parameters returned {} lambdas but {lambda_count} expected",
+ fm.name(),
+ params.len()
+ );
+ }
+
+ for (arg, field) in
+ std::iter::zip(&partially_planned, &mut fields)
+ {
+ match (arg, field) {
+ (
+ ExprOrLambda::Lambda(lambda),
+ ValueOrLambda::Lambda(field),
+ ) => {
+ let params = params.next().ok_or_else(|| {
+ internal_datafusion_err!(
+ "params len should have been checked above"
+ )
+ })?;
+ if let Some(params) = params {
+ let lambda = self.sql_lambda_to_logical_lambda(
+ planner_context,
+ schema,
+ lambda.clone(),
+ params,
+ )?;
+
+ *field = Some(lambda.body.to_field(schema)?.1);
+ }
+ }
+ (ExprOrLambda::Expr(_), ValueOrLambda::Value(_)) => {} // nothing to do
+ (ExprOrLambda::Expr(_), ValueOrLambda::Lambda(_)) => {
+ return internal_err!(
+ "value_fields_with_higher_order_udf returned a value for a lambda argument"
+ );
+ }
+ (ExprOrLambda::Lambda(_), ValueOrLambda::Value(_)) => {
+ return internal_err!(
+ "value_fields_with_higher_order_udf returned a lambda for a value argument"
+ );
+ }
+ }
+ }
+ }
+ LambdaParametersProgress::Complete(params) => {
+ let mut params = params.into_iter();
+
+ if params.len() != lambda_count {
+ return plan_err!(
+ "{} lambda_parameters returned {} lambdas but {lambda_count} expected",
+ fm.name(),
+ params.len()
+ );
+ }
+
+ break partially_planned
+ .into_iter()
+ .map(|arg| match arg {
+ ExprOrLambda::Expr(expr) => Ok(expr),
+ ExprOrLambda::Lambda(lambda) => {
+ let params = params.next().ok_or_else(|| {
+ internal_datafusion_err!(
+ "params len should have been checked above"
+ )
+ })?;
+
+ Ok(Expr::Lambda(self.sql_lambda_to_logical_lambda(
+ planner_context,
+ schema,
+ lambda,
+ params,
+ )?))
+ }
+ })
+ .collect::<Result<Vec<_>>>()?;
+ }
+ }
+
+ let limit = fm.signature().lambda_parameters_max_iterations;
+
+ step += 1;
+
+ if step > limit {
+ return plan_err!(
+ "{} lambda_parameters called {limit} times without completion",
+ fm.name()
+ );
+ }
+ };
+
+ let inner = HigherOrderFunction::new(fm, args);
+
+ if name.eq_ignore_ascii_case(inner.name()) {
+ return Ok(Expr::HigherOrderFunction(inner));
+ } else {
+ // If the function is called by an alias, a verbose string representation is created
+ // (e.g., "my_alias(arg1, arg2)") and the expression is wrapped in an `Alias`
+ // to ensure the output column name matches the user's query.
+ let arg_names = inner
+ .args
+ .iter()
+ .map(|arg| arg.to_string())
+ .collect::<Vec<_>>()
+ .join(",");
+ let verbose_alias = format!("{name}({arg_names})");
+
+ return Ok(Expr::HigherOrderFunction(inner).alias(verbose_alias));
+ }
+ }
+
// Build Unnest expression
if name.eq("unnest") {
let mut exprs = self.function_args_to_expr(args, schema, planner_context)?;
@@ -729,6 +912,44 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
}
}
+ fn sql_lambda_to_logical_lambda(
+ &self,
+ planner_context: &mut PlannerContext,
+ schema: &DFSchema,
+ lambda: LambdaFunction,
+ lambda_params: Vec<FieldRef>,
+ ) -> Result<Lambda> {
+ if lambda.params.len() > lambda_params.len() {
+ return plan_err!(
+ "lambda defined {} params but UDF support only {}",
+ lambda.params.len(),
+ lambda_params.len()
+ );
+ }
+
+ let params = lambda
+ .params
+ .iter()
+ .map(|p| crate::utils::normalize_ident(p.clone()))
+ .collect();
+
+ let lambda_parameters = std::iter::zip(lambda_params, ¶ms)
+ .map(|(f, n): (FieldRef, &String)| f.renamed(n.as_str()));
+
+ let mut planner_context = planner_context
+ .clone()
+ .with_lambda_parameters(lambda_parameters);
+
+ Ok(Lambda {
+ params,
+ body: Box::new(self.sql_expr_to_logical_expr(
+ *lambda.body,
+ schema,
+ &mut planner_context,
+ )?),
+ })
+ }
+
pub(super) fn sql_fn_name_to_expr(
&self,
expr: SQLExpr,
@@ -956,3 +1177,24 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
}
}
}
+
+/// After normalization with [normalize_ident], check whether all params are unique
+///
+/// [normalize_ident]: crate::utils::normalize_ident
+fn all_unique(params: &[sqlparser::ast::Ident]) -> bool {
+ match params.len() {
+ 0 | 1 => true,
+ 2 => {
+ crate::utils::normalize_ident(params[0].clone())
+ != crate::utils::normalize_ident(params[1].clone())
+ }
+ _ => {
+ let mut set = HashSet::with_capacity(params.len());
+
+ params
+ .iter()
+ .map(|p| crate::utils::normalize_ident(p.clone()))
+ .all(|p| set.insert(p))
+ }
+ }
+}
diff --git a/datafusion/sql/src/expr/identifier.rs b/datafusion/sql/src/expr/identifier.rs
index cca09df0db..775cb2cfe8 100644
--- a/datafusion/sql/src/expr/identifier.rs
+++ b/datafusion/sql/src/expr/identifier.rs
@@ -21,6 +21,7 @@ use datafusion_common::{
Column, DFSchema, Result, Span, TableReference, assert_or_internal_err,
exec_datafusion_err, internal_err, not_impl_err, plan_datafusion_err, plan_err,
};
+use datafusion_expr::expr::LambdaVariable;
use datafusion_expr::planner::PlannerResult;
use datafusion_expr::{Case, Expr};
use sqlparser::ast::{CaseWhen, Expr as SQLExpr, Ident};
@@ -59,6 +60,19 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
// identifier. (e.g. it is "foo.bar" not foo.bar)
let normalize_ident = self.ident_normalizer.normalize(id);
+ // lambdas parameters have higher precedence
+ if let Some(field) = planner_context.lambda_parameters().get(&normalize_ident)
+ {
+ let mut lambda_var =
+ LambdaVariable::new(normalize_ident, Some(Arc::clone(field)));
+ if self.options.collect_spans
+ && let Some(span) = Span::try_from_sqlparser_span(id_span)
+ {
+ lambda_var.spans_mut().add_span(span);
+ }
+ return Ok(Expr::LambdaVariable(lambda_var));
+ }
+
// Check for qualified field with unqualified name
if let Ok((qualifier, _)) =
schema.qualified_field_with_unqualified_name(normalize_ident.as_str())
diff --git a/datafusion/sql/src/expr/mod.rs b/datafusion/sql/src/expr/mod.rs
index ae52d8403e..5cbc1c84bd 100644
--- a/datafusion/sql/src/expr/mod.rs
+++ b/datafusion/sql/src/expr/mod.rs
@@ -1351,7 +1351,9 @@ mod tests {
use datafusion_common::TableReference;
use datafusion_common::config::ConfigOptions;
use datafusion_expr::logical_plan::builder::LogicalTableSource;
- use datafusion_expr::{AggregateUDF, ScalarUDF, TableSource, WindowUDF};
+ use datafusion_expr::{
+ AggregateUDF, HigherOrderUDF, ScalarUDF, TableSource, WindowUDF,
+ };
use super::*;
@@ -1391,6 +1393,10 @@ mod tests {
None
}
+ fn get_higher_order_meta(&self, _name: &str) -> Option<Arc<dyn HigherOrderUDF>> {
+ None
+ }
+
fn get_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>> {
match name {
"sum" => Some(datafusion_functions_aggregate::sum::sum_udaf()),
@@ -1414,6 +1420,10 @@ mod tests {
Vec::new()
}
+ fn higher_order_function_names(&self) -> Vec<String> {
+ Vec::new()
+ }
+
fn udaf_names(&self) -> Vec<String> {
vec!["sum".to_string()]
}
diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs
index 1bd257adab..df0e68eab6 100644
--- a/datafusion/sql/src/planner.rs
+++ b/datafusion/sql/src/planner.rs
@@ -274,6 +274,8 @@ pub struct PlannerContext {
/// (UNION/INTERSECT/EXCEPT), holds the schema of the left-most query.
/// Used to alias duplicate expressions to match the left side's field names.
set_expr_left_schema: Option<DFSchemaRef>,
+ /// The parameters of all lambdas seen so far
+ lambda_parameters: HashMap<String, FieldRef>,
}
impl Default for PlannerContext {
@@ -292,6 +294,7 @@ impl PlannerContext {
outer_from_schema: None,
create_table_schema: None,
set_expr_left_schema: None,
+ lambda_parameters: HashMap::new(),
}
}
@@ -401,6 +404,20 @@ impl PlannerContext {
self.ctes.get(cte_name).map(|cte| cte.as_ref())
}
+ pub fn lambda_parameters(&self) -> &HashMap<String, FieldRef> {
+ &self.lambda_parameters
+ }
+
+ pub fn with_lambda_parameters(
+ mut self,
+ parameters: impl IntoIterator<Item = FieldRef>,
+ ) -> Self {
+ self.lambda_parameters
+ .extend(parameters.into_iter().map(|f| (f.name().clone(), f)));
+
+ self
+ }
+
/// Remove the plan of CTE / Subquery for the specified name
pub(super) fn remove_cte(&mut self, cte_name: &str) {
self.ctes.remove(cte_name);
diff --git a/datafusion/sql/src/unparser/dialect.rs b/datafusion/sql/src/unparser/dialect.rs
index 027aec4ac5..d934462240 100644
--- a/datafusion/sql/src/unparser/dialect.rs
+++ b/datafusion/sql/src/unparser/dialect.rs
@@ -170,6 +170,18 @@ pub trait Dialect: Send + Sync {
Ok(None)
}
+ /// Allows the dialect to override higher order function unparsing if the dialect has specific rules.
+ /// Returns None if the default unparsing should be used, or Some(ast::Expr) if there is
+ /// a custom implementation for the function.
+ fn higher_order_function_to_sql_overrides(
+ &self,
+ _unparser: &Unparser,
+ _func_name: &str,
+ _args: &[Expr],
+ ) -> Result<Option<ast::Expr>> {
+ Ok(None)
+ }
+
/// Allows the dialect to choose to omit window frame in unparsing
/// based on function name and window frame bound
/// Returns false if specific function name / window frame bound indicates no window frame is needed in unparsing
diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs
index ca51af2769..e76aea5849 100644
--- a/datafusion/sql/src/unparser/expr.rs
+++ b/datafusion/sql/src/unparser/expr.rs
@@ -16,13 +16,16 @@
// under the License.
use datafusion_common::datatype::DataTypeExt;
-use datafusion_expr::expr::{AggregateFunctionParams, Unnest, WindowFunctionParams};
+use datafusion_expr::expr::{
+ AggregateFunctionParams, HigherOrderFunction, WindowFunctionParams,
+};
+use datafusion_expr::expr::{Lambda, Unnest};
use sqlparser::ast::Value::SingleQuotedString;
use sqlparser::ast::{
- self, Array, BinaryOperator, CaseWhen, DuplicateTreatment, Expr as AstExpr, Function,
- Ident, Interval, ObjectName, OrderByOptions, Subscript, TimezoneInfo, UnaryOperator,
- ValueWithSpan,
+ self, Array, BinaryOperator, Expr as AstExpr, Function, Ident, Interval, ObjectName,
+ Subscript, TimezoneInfo, UnaryOperator,
};
+use sqlparser::ast::{CaseWhen, DuplicateTreatment, OrderByOptions, ValueWithSpan};
use std::sync::Arc;
use std::vec;
@@ -552,6 +555,33 @@ impl Unparser<'_> {
}
Expr::OuterReferenceColumn(_, col) => self.col_to_sql(col),
Expr::Unnest(unnest) => self.unnest_to_sql(unnest),
+ Expr::HigherOrderFunction(HigherOrderFunction { func, args }) => {
+ let func_name = func.name();
+
+ if let Some(expr) = self
+ .dialect
+ .higher_order_function_to_sql_overrides(self, func_name, args)?
+ {
+ return Ok(expr);
+ }
+
+ self.function_to_sql_internal(func_name, args)
+ }
+ Expr::Lambda(Lambda { params, body }) => {
+ Ok(ast::Expr::Lambda(ast::LambdaFunction {
+ params: ast::OneOrManyWithParens::Many(
+ params
+ .iter()
+ .map(|param| self.new_ident_quoted_if_needs(param.clone()))
+ .collect(),
+ ),
+ body: Box::new(self.expr_to_sql_inner(body)?),
+ syntax: ast::LambdaSyntax::Arrow,
+ }))
+ }
+ Expr::LambdaVariable(l) => Ok(ast::Expr::Identifier(
+ self.new_ident_quoted_if_needs(l.name.clone()),
+ )),
}
}
@@ -567,11 +597,11 @@ impl Unparser<'_> {
"get_field" => self.get_field_to_sql(args),
"map" => self.map_to_sql(args),
// TODO: support for the construct and access functions of the `map` type
- _ => self.scalar_function_to_sql_internal(func_name, args),
+ _ => self.function_to_sql_internal(func_name, args),
}
}
- fn scalar_function_to_sql_internal(
+ fn function_to_sql_internal(
&self,
func_name: &str,
args: &[Expr],
@@ -1846,12 +1876,13 @@ mod tests {
use ast::ObjectName;
use datafusion_common::datatype::DataTypeExt;
use datafusion_common::{Spans, TableReference};
- use datafusion_expr::expr::WildcardOptions;
+ use datafusion_expr::expr::{LambdaVariable, WildcardOptions};
use datafusion_expr::{
- ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature,
- Volatility, WindowFrame, WindowFunctionDefinition, case, cast, col, cube, exists,
- grouping_set, interval_datetime_lit, interval_year_month_lit, lit, not,
- not_exists, out_ref_col, placeholder, rollup, table_scan, try_cast, when,
+ ColumnarValue, HigherOrderUDF, LambdaParametersProgress, ScalarFunctionArgs,
+ ScalarUDF, ScalarUDFImpl, Signature, ValueOrLambda, Volatility, WindowFrame,
+ WindowFunctionDefinition, case, cast, col, cube, exists, grouping_set,
+ interval_datetime_lit, interval_year_month_lit, lambda, lit, not, not_exists,
+ out_ref_col, placeholder, rollup, table_scan, try_cast, when,
};
use datafusion_expr::{ExprFunctionExt, interval_month_day_nano_lit};
use datafusion_functions::datetime::from_unixtime::FromUnixtimeFunc;
@@ -1905,6 +1936,41 @@ mod tests {
}
// See sql::tests for E2E tests.
+ #[derive(Debug, Hash, Eq, PartialEq)]
+ struct DummyHigherOrderUDF;
+
+ impl HigherOrderUDF for DummyHigherOrderUDF {
+ fn name(&self) -> &str {
+ "dummy_higher_order_function"
+ }
+
+ fn signature(&self) -> &datafusion_expr::HigherOrderSignature {
+ unimplemented!()
+ }
+
+ fn lambda_parameters(
+ &self,
+ _step: usize,
+ _fields: &[ValueOrLambda<FieldRef, Option<FieldRef>>],
+ ) -> Result<LambdaParametersProgress> {
+ unimplemented!()
+ }
+
+ fn return_field_from_args(
+ &self,
+ _args: datafusion_expr::HigherOrderReturnFieldArgs,
+ ) -> Result<FieldRef> {
+ unimplemented!()
+ }
+
+ fn invoke_with_args(
+ &self,
+ _args: datafusion_expr::HigherOrderFunctionArgs,
+ ) -> Result<ColumnarValue> {
+ unimplemented!()
+ }
+ }
+
#[test]
fn expr_to_sql_ok() -> Result<()> {
let dummy_schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
@@ -1989,6 +2055,22 @@ mod tests {
.is_not_null(),
r#"dummy_udf(a, b) IS NOT NULL"#,
),
+ (
+ Expr::HigherOrderFunction(HigherOrderFunction::new(
+ Arc::new(DummyHigherOrderUDF),
+ vec![
+ col("a"),
+ lambda(
+ ["v"],
+ -Expr::LambdaVariable(LambdaVariable::new(
+ "v".to_string(),
+ Some(Arc::new(Field::new("", DataType::Null, true))),
+ )),
+ ),
+ ],
+ )),
+ r#"dummy_higher_order_function(a, (v) -> -v)"#,
+ ),
(
Expr::Like(Like {
negated: true,
diff --git a/datafusion/sql/tests/common/mod.rs b/datafusion/sql/tests/common/mod.rs
index b93f6135c1..71e864d2a7 100644
--- a/datafusion/sql/tests/common/mod.rs
+++ b/datafusion/sql/tests/common/mod.rs
@@ -27,7 +27,9 @@ use datafusion_common::datatype::DataTypeExt;
use datafusion_common::file_options::file_type::FileType;
use datafusion_common::{DFSchema, GetExt, Result, TableReference, plan_err};
use datafusion_expr::planner::{ExprPlanner, PlannerResult, TypePlanner};
-use datafusion_expr::{AggregateUDF, Expr, ScalarUDF, TableSource, WindowUDF};
+use datafusion_expr::{
+ AggregateUDF, Expr, HigherOrderUDF, ScalarUDF, TableSource, WindowUDF,
+};
use datafusion_functions_nested::expr_fn::make_array;
use datafusion_sql::planner::ContextProvider;
@@ -54,6 +56,7 @@ impl Display for MockCsvType {
#[derive(Default)]
pub(crate) struct MockSessionState {
scalar_functions: HashMap<String, Arc<ScalarUDF>>,
+ higher_order_functions: HashMap<String, Arc<dyn HigherOrderUDF>>,
aggregate_functions: HashMap<String, Arc<AggregateUDF>>,
expr_planners: Vec<Arc<dyn ExprPlanner>>,
type_planner: Option<Arc<dyn TypePlanner>>,
@@ -95,6 +98,17 @@ impl MockSessionState {
.insert(window_function.name().to_string(), window_function);
self
}
+
+ pub fn with_higher_order_function(
+ mut self,
+ higher_order_function: Arc<dyn HigherOrderUDF>,
+ ) -> Self {
+ self.higher_order_functions.insert(
+ higher_order_function.name().to_string(),
+ higher_order_function,
+ );
+ self
+ }
}
pub(crate) struct MockContextProvider {
@@ -277,6 +291,10 @@ impl ContextProvider for MockContextProvider {
self.state.scalar_functions.get(name).cloned()
}
+ fn get_higher_order_meta(&self, name: &str) -> Option<Arc<dyn HigherOrderUDF>> {
+ self.state.higher_order_functions.get(name).cloned()
+ }
+
fn get_aggregate_meta(&self, name: &str) -> Option<Arc<AggregateUDF>> {
self.state.aggregate_functions.get(name).cloned()
}
@@ -312,6 +330,10 @@ impl ContextProvider for MockContextProvider {
self.state.scalar_functions.keys().cloned().collect()
}
+ fn higher_order_function_names(&self) -> Vec<String> {
+ self.state.higher_order_functions.keys().cloned().collect()
+ }
+
fn udaf_names(&self) -> Vec<String> {
self.state.aggregate_functions.keys().cloned().collect()
}
diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs
index e0fc5422d6..724057318c 100644
--- a/datafusion/sql/tests/sql_integration.rs
+++ b/datafusion/sql/tests/sql_integration.rs
@@ -28,8 +28,13 @@ use arrow::datatypes::{TimeUnit::Nanosecond, *};
use common::MockContextProvider;
use datafusion_common::{DFSchema, DataFusionError, Result, assert_contains};
use datafusion_expr::{
- ColumnarValue, CreateIndex, DdlStatement, ScalarFunctionArgs, ScalarUDF,
- ScalarUDFImpl, Signature, Volatility, col, logical_plan::LogicalPlan,
+ ColumnarValue, CreateIndex, DdlStatement, Expr, HigherOrderFunctionArgs,
+ HigherOrderReturnFieldArgs, HigherOrderSignature, HigherOrderUDF,
+ LambdaParametersProgress, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature,
+ ValueOrLambda, Volatility, col,
+ expr::{HigherOrderFunction, LambdaVariable, ScalarFunction},
+ lambda,
+ logical_plan::LogicalPlan,
test::function_stub::sum_udaf,
};
use datafusion_functions::{string, unicode};
@@ -51,7 +56,9 @@ use datafusion_functions_nested::make_array::make_array_udf;
use datafusion_functions_window::{rank::rank_udwf, row_number::row_number_udwf};
use insta::{allow_duplicates, assert_snapshot};
use rstest::rstest;
-use sqlparser::dialect::{Dialect, GenericDialect, HiveDialect, MySqlDialect};
+use sqlparser::dialect::{
+ DatabricksDialect, Dialect, GenericDialect, HiveDialect, MySqlDialect,
+};
use sqlparser::parser::Parser;
mod cases;
@@ -3491,7 +3498,11 @@ fn logical_plan_with_options(sql: &str, options: ParserOptions) -> Result<Logica
}
fn logical_plan_with_dialect(sql: &str, dialect: &dyn Dialect) -> Result<LogicalPlan> {
- let state = MockSessionState::default().with_aggregate_function(sum_udaf());
+ let state = MockSessionState::default()
+ .with_aggregate_function(sum_udaf())
+ .with_higher_order_function(Arc::new(MockArrayReduce::new()))
+ .with_scalar_function(make_array_udf())
+ .with_expr_planner(Arc::new(CustomExprPlanner {})); // plan array literal
let context = MockContextProvider { state };
let planner = SqlToRel::new(&context);
let result = DFParser::parse_sql_with_dialect(sql, dialect);
@@ -5320,3 +5331,159 @@ Sort: customer.c_custkey ASC NULLS LAST
"#
);
}
+
+#[test]
+fn test_progressive_lambda_parameters() {
+ let sql = "select array_reduce([1.0, 2.0], 0, (acc, v) -> acc + v, v -> -v)";
+
+ let expr = logical_plan_with_dialect(sql, &DatabricksDialect)
+ .unwrap()
+ .head_output_expr()
+ .unwrap()
+ .unwrap();
+
+ // taking into account the user defined coercion that coerced the List(Float64) to List(Float32),
+ // test that merge accumulator parameter and finish parameter correctly got the type of the merge
+ // lambda output (Float32 instead of Float64), and not the initial value type (Int64)
+ assert_eq!(
+ expr,
+ Expr::HigherOrderFunction(HigherOrderFunction::new(
+ Arc::new(MockArrayReduce::new()),
+ vec![
+ Expr::ScalarFunction(ScalarFunction::new_udf(
+ make_array_udf(),
+ // note the array being reduced is List(Float64)
+ vec![
+ Expr::Literal(1.0f64.into(), None),
+ Expr::Literal(2.0f64.into(), None)
+ ]
+ )),
+ Expr::Literal(0i64.into(), None),
+ lambda(
+ ["acc", "v"],
+ // lambda vars are Float32
+ resolved_lambda_var("acc", DataType::Float32)
+ + resolved_lambda_var("v", DataType::Float32)
+ ),
+ lambda(["v"], -resolved_lambda_var("v", DataType::Float32)),
+ ]
+ ))
+ )
+}
+
+fn resolved_lambda_var(name: &str, dt: DataType) -> Expr {
+ Expr::LambdaVariable(LambdaVariable::new(
+ name.to_string(),
+ Some(Arc::new(Field::new(name.to_string(), dt, true))),
+ ))
+}
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub struct MockArrayReduce {
+ signature: HigherOrderSignature,
+}
+
+impl MockArrayReduce {
+ #[expect(clippy::new_without_default)]
+ pub fn new() -> Self {
+ Self {
+ signature: HigherOrderSignature::user_defined(Volatility::Immutable),
+ }
+ }
+}
+
+impl HigherOrderUDF for MockArrayReduce {
+ fn name(&self) -> &str {
+ "array_reduce"
+ }
+
+ fn aliases(&self) -> &[String] {
+ &[]
+ }
+
+ fn signature(&self) -> &HigherOrderSignature {
+ &self.signature
+ }
+
+ fn coerce_value_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
+ let [_list, initial] = arg_types else {
+ unreachable!()
+ };
+
+ Ok(vec![
+ DataType::new_list(DataType::Float32, true),
+ initial.clone(),
+ ])
+ }
+
+ fn lambda_parameters(
+ &self,
+ step: usize,
+ fields: &[ValueOrLambda<FieldRef, Option<FieldRef>>],
+ ) -> Result<LambdaParametersProgress> {
+ // optional finish not supported for simplicity
+ let [
+ ValueOrLambda::Value(list),
+ ValueOrLambda::Value(initial_value),
+ ValueOrLambda::Lambda(merge),
+ ValueOrLambda::Lambda(_finish),
+ ] = fields
+ else {
+ unreachable!()
+ };
+
+ let list_field = match list.data_type() {
+ DataType::List(field) => field,
+ _ => unreachable!(),
+ };
+
+ Ok(match (step, merge) {
+ (0, None) => {
+ // at the first step, we use the initial_value as merge accumulator,
+ // and return None for finish since we don't know the output of merge
+ LambdaParametersProgress::Partial(vec![
+ // merge
+ Some(vec![Arc::clone(initial_value), Arc::clone(list_field)]),
+ // finish
+ None,
+ ])
+ }
+ (1, Some(accumulator)) | (0, Some(accumulator)) => {
+ // now we can use the merge output as it's accumulator and
+ // as the finish parameter
+ LambdaParametersProgress::Complete(vec![
+ // merge
+ vec![Arc::clone(accumulator), Arc::clone(list_field)],
+ // finish
+ vec![Arc::clone(accumulator)],
+ ])
+ }
+ (1, None) => {
+ unreachable!()
+ }
+ _ => todo!(),
+ })
+ }
+
+ fn return_field_from_args(
+ &self,
+ args: HigherOrderReturnFieldArgs,
+ ) -> Result<FieldRef> {
+ // optional finish not supported for simplicity
+ let [
+ ValueOrLambda::Value(_list),
+ ValueOrLambda::Value(_initial_value),
+ ValueOrLambda::Lambda(_merge),
+ ValueOrLambda::Lambda(finish),
+ ] = args.arg_fields
+ else {
+ unreachable!()
+ };
+
+ Ok(Arc::clone(finish))
+ }
+
+ fn invoke_with_args(&self, _args: HigherOrderFunctionArgs) -> Result<ColumnarValue> {
+ unreachable!()
+ }
+}
diff --git a/datafusion/sqllogictest/test_files/higher_order.slt b/datafusion/sqllogictest/test_files/higher_order.slt
new file mode 100644
index 0000000000..86359194a5
--- /dev/null
+++ b/datafusion/sqllogictest/test_files/higher_order.slt
@@ -0,0 +1,286 @@
+# 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.
+
+#############
+## Higher Order Function Tests
+#############
+
+statement ok
+set datafusion.sql_parser.dialect = databricks;
+
+statement ok
+CREATE TABLE t (list array<int>, number int)
+AS VALUES
+([1, 50], 10),
+([4, 50], 40),
+([7, 50], 60);
+
+statement ok
+CREATE TABLE with_null_list (list array<int>)
+AS VALUES
+([1, 2]),
+(NULL);
+
+statement ok
+CREATE TABLE fully_null_list (list array<int>)
+AS VALUES
+(NULL),
+(NULL);
+
+query ?
+SELECT array_transform([1,2,3,4,5], v -> v*2);
+----
+[2, 4, 6, 8, 10]
+
+query ?
+SELECT array_transform([1,2,3,4,5], v -> repeat("a", v));
+----
+[a, aa, aaa, aaaa, aaaaa]
+
+query ?
+SELECT array_transform([1,2,3,4,5], v -> list_repeat("a", v));
+----
+[[a], [a, a], [a, a, a], [a, a, a, a], [a, a, a, a, a]]
+
+# return scalar
+query I?
+SELECT t.number, array_transform([1, 2], e1 -> 24) from t;
+----
+10 [24, 24]
+40 [24, 24]
+60 [24, 24]
+
+# shadows parent lambda variable
+query ?
+SELECT array_transform([[1, 2]], a -> array_transform(a, a -> a+1))
+----
+[[2, 3]]
+
+# multiple nesting
+query ?
+SELECT array_transform([[[1], [2], [3]]], a -> array_transform(a, b -> array_transform(b, c -> c*2)));
+----
+[[[2], [4], [6]]]
+
+# parameter shadows unqualified column
+query I?
+SELECT number, array_transform([1, 2], number -> number+1) from t;
+----
+10 [2, 3]
+40 [2, 3]
+60 [2, 3]
+
+# type coercion inside lambda body
+query ?
+SELECT array_transform([t.number], v -> v + 3.0) from t;
+----
+[13.0]
+[43.0]
+[63.0]
+
+query TT
+EXPLAIN SELECT array_transform([t.number], v -> v + 3.0) from t;
+----
+logical_plan
+01)Projection: array_transform(make_array(t.number), (v) -> CAST(v AS Float64) + Float64(3))
+02)--TableScan: t projection=[number]
+physical_plan
+01)ProjectionExec: expr=[array_transform(make_array(number@0), (v) -> CAST(v@0 AS Float64) + 3) as array_transform(make_array(t.number),(v) -> v + Float64(3))]
+02)--DataSourceExec: partitions=1, partition_sizes=[1]
+
+#cse should not eliminate subtrees containing lambdas
+query TT
+explain select array_transform([t.number], v -> 5), array_transform([t.number+1], v -> 5) from t;
+----
+logical_plan
+01)Projection: array_transform(make_array(t.number), (v) -> Int64(5)), array_transform(make_array(CAST(t.number AS Int64) + Int64(1)), (v) -> Int64(5))
+02)--TableScan: t projection=[number]
+physical_plan
+01)ProjectionExec: expr=[array_transform(make_array(number@0), (v) -> 5) as array_transform(make_array(t.number),(v) -> Int64(5)), array_transform(make_array(CAST(number@0 AS Int64) + 1), (v) -> 5) as array_transform(make_array(t.number + Int64(1)),(v) -> Int64(5))]
+02)--DataSourceExec: partitions=1, partition_sizes=[1]
+
+#cse should not eliminate subtrees containing lambda variables
+query TT
+explain select array_transform([t.number], v -> v*2), array_transform([t.number], v -> v*2-1) from t;
+----
+logical_plan
+01)Projection: array_transform(__common_expr_1 AS make_array(t.number), (v) -> CAST(v AS Int64) * Int64(2)), array_transform(__common_expr_1 AS make_array(t.number), (v) -> CAST(v AS Int64) * Int64(2) - Int64(1))
+02)--Projection: make_array(t.number) AS __common_expr_1
+03)----TableScan: t projection=[number]
+physical_plan
+01)ProjectionExec: expr=[array_transform(__common_expr_1@0, (v) -> CAST(v@0 AS Int64) * 2) as array_transform(make_array(t.number),(v) -> v * Int64(2)), array_transform(__common_expr_1@0, (v) -> CAST(v@0 AS Int64) * 2 - 1) as array_transform(make_array(t.number),(v) -> v * Int64(2) - Int64(1))]
+02)--ProjectionExec: expr=[make_array(number@0) as __common_expr_1]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
+
+
+# test that sql planner plans resolved lambda variables, as v[1] planning checks the datatype of lhs
+query ?
+SELECT array_transform([[10, 20]], v -> v[1]);
+----
+[10]
+
+
+# expr simplifier inside lambda body
+query TT
+EXPLAIN SELECT array_transform([t.number], v -> v = v) from t;
+----
+logical_plan
+01)Projection: array_transform(make_array(t.number), (v) -> v IS NOT NULL OR Boolean(NULL)) AS array_transform(make_array(t.number),(v) -> v = v)
+02)--TableScan: t projection=[number]
+physical_plan
+01)ProjectionExec: expr=[array_transform(make_array(number@0), (v) -> v@0 IS NOT NULL OR NULL) as array_transform(make_array(t.number),(v) -> v = v)]
+02)--DataSourceExec: partitions=1, partition_sizes=[1]
+
+
+# array_transform coercion rules
+query TT
+explain select array_transform(arrow_cast(t.list, 'ListView(Int32)'), a -> a+1) from t;
+----
+logical_plan
+01)Projection: array_transform(CAST(CAST(t.list AS ListView(Int32)) AS List(Int32)), (a) -> CAST(a AS Int64) + Int64(1)) AS array_transform(arrow_cast(t.list,Utf8("ListView(Int32)")),(a) -> a + Int64(1))
+02)--TableScan: t projection=[list]
+physical_plan
+01)ProjectionExec: expr=[array_transform(CAST(CAST(list@0 AS ListView(Int32)) AS List(Int32)), (a) -> CAST(a@0 AS Int64) + 1) as array_transform(arrow_cast(t.list,Utf8("ListView(Int32)")),(a) -> a + Int64(1))]
+02)--DataSourceExec: partitions=1, partition_sizes=[1]
+
+query ?
+select array_transform(arrow_cast(t.list, 'ListView(Int32)'), a -> a+1) from t;
+----
+[2, 51]
+[5, 51]
+[8, 51]
+
+# higher order function with inner case using lambda variables only
+query ?
+select array_transform([3, 5, 0], v -> case when v > 1 then 2 when v > 4 then 6 else 8 end);
+----
+[2, 2, 8]
+
+#case with inner higher order function
+query I??
+select
+ t.number,
+ t.list,
+ case
+ when t.number > 30 then array_transform(t.list, v -> v+1)
+ else array_transform(t.list, v -> v*2)
+ end
+from t
+order by t.number;
+----
+10 [1, 50] [2, 100]
+40 [4, 50] [5, 51]
+60 [7, 50] [8, 51]
+
+# array_transform coercion rules
+query TT
+explain select array_transform(arrow_cast(t.list, 'FixedSizeList(2, Int32)'), a -> a+1) from t;
+----
+logical_plan
+01)Projection: array_transform(CAST(CAST(t.list AS FixedSizeList(2 x Int32)) AS List(Int32)), (a) -> CAST(a AS Int64) + Int64(1)) AS array_transform(arrow_cast(t.list,Utf8("FixedSizeList(2, Int32)")),(a) -> a + Int64(1))
+02)--TableScan: t projection=[list]
+physical_plan
+01)ProjectionExec: expr=[array_transform(CAST(CAST(list@0 AS FixedSizeList(2 x Int32)) AS List(Int32)), (a) -> CAST(a@0 AS Int64) + 1) as array_transform(arrow_cast(t.list,Utf8("FixedSizeList(2, Int32)")),(a) -> a + Int64(1))]
+02)--DataSourceExec: partitions=1, partition_sizes=[1]
+
+query ?
+select array_transform(arrow_cast(t.list, 'FixedSizeList(2, Int32)'), a -> a+1) from t;
+----
+[2, 51]
+[5, 51]
+[8, 51]
+
+query ?
+select array_transform(arrow_cast(list, 'FixedSizeList(2, Int32)'), a -> a+1) from fully_null_list;
+----
+NULL
+NULL
+
+query ?
+select array_transform(arrow_cast(list, 'FixedSizeList(2, Int32)'), a -> a+1) from with_null_list;
+----
+[2, 3]
+NULL
+
+query ?
+select array_transform(list, a -> a+1) from with_null_list;
+----
+[2, 3]
+NULL
+
+# empty sublists fast path
+query ?
+select array_transform([], v -> v+1);
+----
+[]
+
+# all null fast path
+query ?
+select array_transform(list, v -> v+1) from fully_null_list;
+----
+NULL
+NULL
+
+# higher order function with inner case using lambda variables and captured column(capture not supported yet)
+query error DataFusion error: Error during planning: lambda doesn't support column capture
+select array_transform([3, 5, 9], v -> case when v > 1 then t.number when v > 4 then 6 else 8 end) from t;
+
+# higher order function with inner case using captured column only(capture not supported yet)
+query error DataFusion error: Error during planning: lambda doesn't support column capture
+select array_transform([3, 5, 9], v -> case when t.number > 1 then 2 when t.number > 4 then 6 else 8 end) from t;
+
+query error
+select array_transform();
+----
+DataFusion error: Error during planning: array_transform function requires 1 value arguments, got 0
+
+
+query error DataFusion error: Error during planning: array_transform expected a list as first argument, got Int64
+select array_transform(1, v -> v*2);
+
+query error DataFusion error: Error during planning: array_transform expects a value followed by a lambda, got Lambda\(None\) and Value\(Field \{ name: "make_array\(Int64\(1\),Int64\(2\)\)", data_type: List\(Field \{ data_type: Int64, nullable: true \}\), nullable: true \}\)
+select array_transform(v -> v*2, [1, 2]);
+
+query error DataFusion error: Error during planning: lambda defined 3 params but UDF support only 1
+SELECT array_transform([1, 2], (e, i, j) -> i);
+
+query error DataFusion error: Error during planning: lambda parameters names must be unique, got \(v, v\)
+SELECT array_transform([1], (v, v) -> v*2);
+
+query error DataFusion error: This feature is not implemented: Unsupported ast node in sqltorel: Lambda\(LambdaFunction \{ params: One\(Ident \{ value: "v", quote_style: None, span: Span\(Location\(1,12\)\.\.Location\(1,13\)\) \}\), body: Identifier\(Ident \{ value: "v", quote_style: None, span: Span\(Location\(1,17\)\.\.Location\(1,18\)\) \}\), syntax: Arrow \}\)
+SELECT abs(v -> v);
+
+query error DataFusion error: This feature is not implemented: Unsupported ast node in sqltorel: Lambda\(LambdaFunction \{ params: One\(Ident \{ value: "v", quote_style: None, span: Span\(Location\(1,8\)\.\.Location\(1,9\)\) \}\), body: Identifier\(Ident \{ value: "v", quote_style: None, span: Span\(Location\(1,13\)\.\.Location\(1,14\)\) \}\), syntax: Arrow \}\)
+SELECT v -> v;
+
+query error DataFusion error: This feature is not implemented: Unsupported ast node in sqltorel: Lambda\(LambdaFunction \{ params: One\(Ident \{ value: "v", quote_style: None, span: Span\(Location\(1,34\)\.\.Location\(1,35\)\) \}\), body: BinaryOp \{ left: Identifier\(Ident \{ value: "v", quote_style: None, span: Span\(Location\(1,39\)\.\.Location\(1,40\)\) \}\), op: Plus, right: Value\(ValueWithSpan \{ value: Number\("1", false\), span: Span\(Location\(1,41\)\.\.Location\(1,42\)\) \}\) [...]
+SELECT array_transform([1], v -> v -> v+1);
+
+query error DataFusion error: SQL error: ParserError\("Expected: an expression, found: \) at Line: 1, Column: 30"\)
+SELECT array_transform([1], () -> 1);
+
+statement ok
+drop table t;
+
+statement ok
+drop table with_null_list;
+
+statement ok
+drop table fully_null_list;
+
+statement ok
+set datafusion.sql_parser.dialect = generic;
diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs b/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs
index 10fe58862e..1a0fb3f55f 100644
--- a/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs
+++ b/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs
@@ -30,6 +30,7 @@ pub async fn from_scalar_function(
f: &ScalarFunction,
input_schema: &DFSchema,
) -> Result<Expr> {
+ //TODO: handle higher order functions, as they are also encoded as scalar functions
let Some(fn_signature) = consumer
.get_extensions()
.functions
diff --git a/datafusion/substrait/src/logical_plan/producer/expr/mod.rs b/datafusion/substrait/src/logical_plan/producer/expr/mod.rs
index b07a0206c5..6e053f0d90 100644
--- a/datafusion/substrait/src/logical_plan/producer/expr/mod.rs
+++ b/datafusion/substrait/src/logical_plan/producer/expr/mod.rs
@@ -151,6 +151,13 @@ pub fn to_substrait_rex(
not_impl_err!("Cannot convert {expr:?} to Substrait")
}
Expr::Unnest(expr) => not_impl_err!("Cannot convert {expr:?} to Substrait"),
+ Expr::HigherOrderFunction(expr) => {
+ producer.handle_higher_order_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")
+ }
}
}
diff --git a/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs b/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs
index 9f70e903a0..e36d5128cd 100644
--- a/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs
+++ b/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs
@@ -26,17 +26,34 @@ pub fn from_scalar_function(
producer: &mut impl SubstraitProducer,
fun: &expr::ScalarFunction,
schema: &DFSchemaRef,
+) -> datafusion::common::Result<Expression> {
+ from_function(producer, fun.name(), &fun.args, schema)
+}
+
+pub fn from_higher_order_function(
+ producer: &mut impl SubstraitProducer,
+ fun: &expr::HigherOrderFunction,
+ schema: &DFSchemaRef,
+) -> datafusion::common::Result<Expression> {
+ from_function(producer, fun.name(), &fun.args, schema)
+}
+
+fn from_function(
+ producer: &mut impl SubstraitProducer,
+ name: &str,
+ args: &[Expr],
+ schema: &DFSchemaRef,
) -> datafusion::common::Result<Expression> {
let mut arguments: Vec<FunctionArgument> = vec![];
- for arg in &fun.args {
+ for arg in args {
arguments.push(FunctionArgument {
arg_type: Some(ArgType::Value(producer.handle_expr(arg, schema)?)),
});
}
- let arguments = custom_argument_handler(fun.name(), arguments);
+ let arguments = custom_argument_handler(name, arguments);
- let function_anchor = producer.register_function(fun.name().to_string());
+ let function_anchor = producer.register_function(name.to_string());
#[expect(deprecated)]
Ok(Expression {
rex_type: Some(RexType::ScalarFunction(ScalarFunction {
diff --git a/datafusion/substrait/src/logical_plan/producer/substrait_producer.rs b/datafusion/substrait/src/logical_plan/producer/substrait_producer.rs
index 9a6c5317ec..4228c32435 100644
--- a/datafusion/substrait/src/logical_plan/producer/substrait_producer.rs
+++ b/datafusion/substrait/src/logical_plan/producer/substrait_producer.rs
@@ -19,11 +19,12 @@ use crate::extensions::Extensions;
use crate::logical_plan::producer::{
from_aggregate, from_aggregate_function, from_alias, from_between, from_binary_expr,
from_case, from_cast, from_column, from_distinct, from_empty_relation, from_exists,
- from_filter, from_in_list, from_in_subquery, from_join, from_like, from_limit,
- from_literal, from_placeholder, from_projection, from_repartition,
- from_scalar_function, from_scalar_subquery, from_set_comparison, from_sort,
- from_subquery_alias, from_table_scan, from_try_cast, from_unary_expr, from_union,
- from_values, from_window, from_window_function, to_substrait_rel, to_substrait_rex,
+ from_filter, from_higher_order_function, from_in_list, from_in_subquery, from_join,
+ from_like, from_limit, from_literal, from_placeholder, from_projection,
+ from_repartition, from_scalar_function, from_scalar_subquery, from_set_comparison,
+ from_sort, from_subquery_alias, from_table_scan, from_try_cast, from_unary_expr,
+ from_union, from_values, from_window, from_window_function, to_substrait_rel,
+ to_substrait_rex,
};
use datafusion::common::{Column, DFSchemaRef, ScalarValue, substrait_err};
use datafusion::execution::SessionState;
@@ -334,6 +335,14 @@ pub trait SubstraitProducer: Send + Sync + Sized {
from_scalar_function(self, scalar_fn, schema)
}
+ fn handle_higher_order_function(
+ &mut self,
+ scalar_fn: &expr::HigherOrderFunction,
+ schema: &DFSchemaRef,
+ ) -> datafusion::common::Result<Expression> {
+ from_higher_order_function(self, scalar_fn, schema)
+ }
+
fn handle_aggregate_function(
&mut self,
agg_fn: &expr::AggregateFunction,
diff --git a/docs/source/user-guide/sql/scalar_functions.md b/docs/source/user-guide/sql/scalar_functions.md
index ffa04c3013..1cb0ff8f4b 100644
--- a/docs/source/user-guide/sql/scalar_functions.md
+++ b/docs/source/user-guide/sql/scalar_functions.md
@@ -3276,6 +3276,7 @@ _Alias of [current_date](#current_date)._
- [array_slice](#array_slice)
- [array_sort](#array_sort)
- [array_to_string](#array_to_string)
+- [array_transform](#array_transform)
- [array_union](#array_union)
- [arrays_overlap](#arrays_overlap)
- [arrays_zip](#arrays_zip)
@@ -3325,6 +3326,7 @@ _Alias of [current_date](#current_date)._
- [list_slice](#list_slice)
- [list_sort](#list_sort)
- [list_to_string](#list_to_string)
+- [list_transform](#list_transform)
- [list_union](#list_union)
- [list_zip](#list_zip)
- [make_array](#make_array)
@@ -4376,6 +4378,34 @@ array_to_string(array, delimiter[, null_string])
- array_join
- list_join
+### `array_transform`
+
+transforms the values of an array
+
+```sql
+array_transform(array, x -> x*2)
+```
+
+#### Arguments
+
+- **array**: Array expression. Can be a constant, column, or function, and any combination of array operators.
+- **lambda**: Lambda
+
+#### Example
+
+```sql
+> select array_transform([1, 2, 3, 4, 5], x -> x*2);
++-------------------------------------------+
+| array_transform([1, 2, 3, 4, 5], x -> x*2) |
++-------------------------------------------+
+| [2, 4, 6, 8, 10] |
++-------------------------------------------+
+```
+
+#### Aliases
+
+- list_transform
+
### `array_union`
Returns an array of elements that are present in both arrays (all elements from both arrays) without duplicates.
@@ -4741,6 +4771,10 @@ _Alias of [array_sort](#array_sort)._
_Alias of [array_to_string](#array_to_string)._
+### `list_transform`
+
+_Alias of [array_transform](#array_transform)._
+
### `list_union`
_Alias of [array_union](#array_union)._
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org