Skip to content

[branch-54] Fix evaluation of lambda with multiple parameters when some are unused (backport of #22853, adapted) #23660

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
gstvg wants to merge 2 commits into
base: branch-54
Choose a base branch
from
+223 −107
Open

[branch-54] Fix evaluation of lambda with multiple parameters when some are unused (backport of #22853, adapted) #23660

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
63228c3
fix: lambda with multiple parameters, some unused
gstvg Jul 17, 2026
b1f83e8
clarify rationale for sqllogictest-only higher-order function
gstvg Jul 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
83 changes: 39 additions & 44 deletions datafusion/expr/src/higher_order_function.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ use crate::expr::{
use crate::type_coercion::functions::value_fields_with_higher_order_udf;
use crate::udf_eq::UdfEq;
use crate::{ColumnarValue, Documentation, Expr, ExprSchemable};
use arrow::array::{ArrayRef, RecordBatch};
use arrow::array::{ArrayRef, NullArray, RecordBatch};
use arrow::datatypes::{DataType, FieldRef, Schema};
use arrow_schema::SchemaRef;
use datafusion_common::config::ConfigOptions;
Expand All @@ -34,8 +34,8 @@ use datafusion_common::tree_node::{
Transformed, TreeNode, TreeNodeContainer, TreeNodeRecursion,
};
use datafusion_common::{
DFSchema, HashMap, HashSet, Result, ScalarValue, exec_err, internal_datafusion_err,
internal_err, not_impl_err, plan_datafusion_err, plan_err,
DFSchema, HashMap, HashSet, Result, ScalarValue, exec_datafusion_err, exec_err,
internal_datafusion_err, internal_err, not_impl_err, plan_datafusion_err, plan_err,
};
use datafusion_expr_common::dyn_eq::{DynEq, DynHash};
use datafusion_expr_common::signature::Volatility;
Expand Down Expand Up @@ -330,61 +330,56 @@ impl LambdaArgument {
args: &[&dyn Fn() -> Result<ArrayRef>],
spread_captures: impl FnOnce(&[ArrayRef]) -> Result<Vec<ArrayRef>>,
) -> Result<ColumnarValue> {
if args.len() < self.params.len() {
return exec_err!(
"expected at least {} lambda arguments to merge with captures, got {}",
self.params.len(),
args.len()
);
}

let args = args
.iter()
.take(self.params.len())
.map(|arg| arg())
.collect::<Result<Vec<_>>>()?;

let row_count = args.first().map(|arg| arg.len()).ok_or_else(|| {
exec_datafusion_err!("lambdas with zero parameters are not supported")
})?;

let spread_captures = self
.captures
.as_ref()
.filter(|b| !b.columns().is_empty())
.map(|captures| {
let spread_columns = spread_captures(captures.columns())?;
let spread_columns = if captures
.schema_ref()
.fields()
.iter()
.all(|f| f.data_type().is_null())
{
let null = Arc::new(NullArray::new(row_count)) as _;

vec![null; captures.num_columns()]
} else {
spread_captures(captures.columns())?
};

RecordBatch::try_new(captures.schema(), spread_columns)
})
.transpose()?;

let merged = merge_captures_with_variables(
spread_captures.as_ref(),
Arc::clone(&self.schema),
&self.params,
args,
)?;
let columns = match spread_captures {
Some(captures) => captures.columns().iter().cloned().chain(args).collect(),
None => args,
};
let merged = RecordBatch::try_new(Arc::clone(&self.schema), columns)?;

self.body.evaluate(&merged)
}
}

fn merge_captures_with_variables(
captures: Option<&RecordBatch>,
schema: SchemaRef,
params: &[FieldRef],
variables: &[&dyn Fn() -> Result<ArrayRef>],
) -> Result<RecordBatch> {
if variables.len() < params.len() {
return exec_err!(
"expected at least {} lambda arguments to merge with captures, got {}",
params.len(),
variables.len()
);
}

let columns = match captures {
Some(captures) => {
let mut columns = captures.columns().to_vec();

for arg in &variables[..params.len()] {
columns.push(arg()?);
}

columns
}
None => variables
.iter()
.take(params.len())
.map(|arg| arg())
.collect::<Result<_>>()?,
};

Ok(RecordBatch::try_new(schema, columns)?)
}

/// Information about arguments passed to the function
///
/// This structure contains metadata about how the function was called
Expand Down
56 changes: 7 additions & 49 deletions datafusion/physical-expr/src/expressions/lambda.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,20 +29,19 @@ use arrow::{
datatypes::{DataType, Schema},
record_batch::RecordBatch,
};
use datafusion_common::{HashSet, Result, internal_err};
use datafusion_common::{
HashMap, plan_err,
tree_node::{Transformed, TreeNode, TreeNodeRecursion},
plan_err,
tree_node::{TreeNode, TreeNodeRecursion},
};
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>,
projected_body: Arc<dyn PhysicalExpr>,
projection: Vec<usize>,
used_column_indices: HashSet<usize>,
}

// Manually derive PartialEq and Hash to work around https://github.com/rust-lang/rust/issues/78808 [https://github.com/apache/datafusion/issues/13196]
Expand Down Expand Up @@ -88,47 +87,10 @@ impl LambdaExpr {
})
.expect("closure should be infallible");

let mut projection = used_column_indices.into_iter().collect::<Vec<_>>();

projection.sort();

let column_index_map = projection
.iter()
.enumerate()
.map(|(projected, original)| (*original, projected))
.collect::<HashMap<_, _>>();

let projected_body = Arc::clone(&body)
.transform_down(|e| {
if let Some(column) = e.downcast_ref::<Column>() {
let original = column.index();
let projected = *column_index_map.get(&original).unwrap();
if projected != original {
return Ok(Transformed::yes(Arc::new(Column::new(
column.name(),
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()),
))));
}
}
Ok(Transformed::no(e))
})
.expect("closure should be infallible")
.data;

Self {
params,
body,
projected_body,
projection,
used_column_indices,
}
}

Expand All @@ -142,12 +104,8 @@ impl LambdaExpr {
&self.body
}

pub(crate) fn projection(&self) -> &[usize] {
&self.projection
}

pub(crate) fn projected_body(&self) -> &Arc<dyn PhysicalExpr> {
&self.projected_body
pub(crate) fn used_column_indices(&self) -> &HashSet<usize> {
&self.used_column_indices
}
}

Expand Down
31 changes: 19 additions & 12 deletions datafusion/physical-expr/src/higher_order_function.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,8 +36,8 @@ use std::sync::Arc;
use crate::PhysicalExpr;
use crate::expressions::{LambdaExpr, Literal};

use arrow::array::{Array, RecordBatch};
use arrow::datatypes::{DataType, FieldRef, Schema};
use arrow::array::{Array, NullArray, RecordBatch, RecordBatchOptions};
use arrow::datatypes::{DataType, Field, FieldRef, Schema};
use datafusion_common::config::{ConfigEntry, ConfigOptions};
use datafusion_common::datatype::FieldExt;
use datafusion_common::utils::remove_list_null_values;
Expand Down Expand Up @@ -337,21 +337,28 @@ impl PhysicalExpr for HigherOrderFunctionExpr {
.map(|(name, param)| param.renamed(name.as_str()))
.collect();

// lambda.projection may include indexes of nested lambda variables not present on this batch
let projection = lambda
.projection()
.iter()
.copied()
.filter(|i| *i < batch.num_columns())
.collect::<Vec<_>>();
let (fields, columns): (Vec<_>, _) = std::iter::zip(batch.schema_ref().fields(), batch.columns())
.enumerate()
.map(|(i, (field, array))| if lambda.used_column_indices().contains(&i) {
(Arc::clone(field), Arc::clone(array))
} else {
(Arc::new(Field::new(field.name(), DataType::Null, true)), Arc::new(NullArray::new(batch.num_rows())) as _)
})
.collect();

let batch= RecordBatch::try_new_with_options(
Arc::new(Schema::new(fields)),
columns,
&RecordBatchOptions::new().with_row_count(Some(batch.num_rows())),
)?;

Ok(ValueOrLambda::Lambda(LambdaArgument::new(
params,
Arc::clone(lambda.projected_body()),
if projection.is_empty() {
Arc::clone(lambda.body()),
if batch.num_columns() == 0 {
None
} else {
Some(batch.project(&projection)?)
Some(batch)
},
)))
}
Expand Down
98 changes: 96 additions & 2 deletions datafusion/sqllogictest/src/test_context.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,8 +41,10 @@ use datafusion::functions::math::abs;
use datafusion::logical_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl};
use datafusion::logical_expr::planner::TypePlanner;
use datafusion::logical_expr::{
ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature,
Volatility, create_udf,
ColumnarValue, Expr, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs,
HigherOrderSignature, HigherOrderUDF, HigherOrderUDFImpl, LambdaParametersProgress,
ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, ValueOrLambda, Volatility,
create_udf,
};
use datafusion::physical_plan::ExecutionPlan;
use datafusion::prelude::*;
Expand Down Expand Up @@ -183,6 +185,10 @@ impl TestContext {
info!("Registering dummy async udf");
register_async_abs_udf(test_ctx.session_ctx())
}
"multi_arg_lambda.slt" => {
info!("Registering dummy multi arg higher-order udf");
register_multi_arg_lambda_udf(test_ctx.session_ctx())
}
_ => {
info!("Using default SessionContext");
}
Expand Down Expand Up @@ -712,3 +718,91 @@ fn register_async_abs_udf(ctx: &SessionContext) {
let udf = AsyncScalarUDF::new(Arc::new(async_abs));
ctx.register_udf(udf.into_scalar_udf());
}

fn register_multi_arg_lambda_udf(ctx: &SessionContext) {
// In 54 release, all built-in higher-order functions are single-parameter, but the public higher-order trait supports
// user defined functions with multi-parameter lambdas. To test the lambda handling of multiple parameters, we create
// a dummy udf that provides every value argument it receives as a parameter to it's lambda, the last argument:
// multi_arg_higher_order_udf([..value_args], ([..value_args_as_params]) -> ...)
#[derive(Debug, PartialEq, Eq, Hash)]

@alamb alamb Jul 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we please add comments about the rationale for this function? It isn't at all clear to me

@gstvg gstvg Jul 18, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

b1f83e8

struct MultiArgHigherOrderUdf {
signature: HigherOrderSignature,
}

impl HigherOrderUDFImpl for MultiArgHigherOrderUdf {
fn name(&self) -> &str {
"multi_arg_higher_order_udf"
}

fn signature(&self) -> &HigherOrderSignature {
&self.signature
}

fn lambda_parameters(
&self,
_step: usize,
fields: &[ValueOrLambda<FieldRef, Option<FieldRef>>],
) -> Result<LambdaParametersProgress> {
let fields = fields
.iter()
.filter_map(|f| match f {
ValueOrLambda::Value(v) => Some(v),
ValueOrLambda::Lambda(_) => None,
})
.cloned()
.collect();

Ok(LambdaParametersProgress::Complete(vec![fields]))
}

fn return_field_from_args(
&self,
args: HigherOrderReturnFieldArgs,
) -> Result<FieldRef> {
Ok(match args.arg_fields.last().unwrap() {
ValueOrLambda::Value(_) => unreachable!(),
ValueOrLambda::Lambda(l) => Arc::clone(l),
})
}

fn invoke_with_args(
&self,
args: HigherOrderFunctionArgs,
) -> Result<ColumnarValue> {
let lambda = match args.args.last().unwrap() {
ValueOrLambda::Value(_) => unreachable!(),
ValueOrLambda::Lambda(l) => l,
};

let num_rows = args
.args
.iter()
.filter_map(|a| match a {
ValueOrLambda::Value(ColumnarValue::Array(a)) => Some(a.len()),
_ => None,
})
.max()
.unwrap_or(1);

let args = args
.args
.iter()
.filter_map(|arg| match arg {
ValueOrLambda::Value(v) => Some(Box::new(|| v.to_array(num_rows))),
ValueOrLambda::Lambda(_) => None,
})
.collect::<Vec<_>>();

let args: Vec<&dyn Fn() -> Result<ArrayRef>> =
args.iter().map(|v| v as _).collect();

lambda.evaluate(&args, |arrays| Ok(arrays.to_vec()))
}
}

ctx.register_higher_order_function(Arc::new(HigherOrderUDF::new_from_impl(
MultiArgHigherOrderUdf {
signature: HigherOrderSignature::variadic_any(Volatility::Immutable),
},
)));
}
Loading
Loading