From 0552b4ff4fdf0a31e15d612437e68ba46b77a659 Mon Sep 17 00:00:00 2001 From: Jayant Shrivastava Date: Thu, 13 Aug 2026 21:36:32 +0000 Subject: [PATCH 1/3] Revert "fix(lambda): only push referenced params into the merged batch (#24162) (#166)" This reverts commit 79de5e90f9c454225d8b1648ef6f743a1d32e824. --- datafusion/expr/src/higher_order_function.rs | 198 +------------- .../physical-expr/src/expressions/lambda.rs | 243 ++---------------- .../src/higher_order_function.rs | 75 +----- 3 files changed, 37 insertions(+), 479 deletions(-) diff --git a/datafusion/expr/src/higher_order_function.rs b/datafusion/expr/src/higher_order_function.rs index c300be8f6fcfe..413714f498164 100644 --- a/datafusion/expr/src/higher_order_function.rs +++ b/datafusion/expr/src/higher_order_function.rs @@ -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, RecordBatchOptions}; +use arrow::array::{ArrayRef, RecordBatch}; use arrow::datatypes::{DataType, FieldRef, Schema}; use arrow_schema::SchemaRef; use datafusion_common::config::ConfigOptions; @@ -239,26 +239,6 @@ pub struct LambdaArgument { /// For example, for `array_transform([2], v -> -v)`, /// this will be `vec![Field::new("v", DataType::Int32, true)]` params: Vec, - /// Indices into [`Self::params`] of the parameters that are actually - /// referenced by [`Self::body`] (taking nested-lambda shadowing into - /// account), in the original declaration order of `params`. - /// - /// [`Self::evaluate`] only evaluates and pushes the closures whose - /// corresponding parameter index appears here, so unused declared - /// parameters leave no slot in the merged batch and the body's compressed - /// column indices line up directly with what the evaluator built. - /// - /// Callers who already have a `LambdaExpr` should pass - /// `LambdaExpr::used_param_indices()` directly to [`Self::new`] — both - /// are indices into the same positionally-aligned `params` list. - /// - /// Every index here must be `< params.len()`; see the precondition on - /// [`Self::new`]. - /// - /// Relies on captures sorting before this lambda's own params in the - /// planner's (un-projected) index space, which is what makes - /// `captures ++ used_params` below line up with the projected body. - used_param_indices: Vec, /// The body of the lambda /// /// For example, for `array_transform([2], v -> -v)`, @@ -277,45 +257,26 @@ pub struct LambdaArgument { } impl LambdaArgument { - /// # Preconditions - /// - /// Every index in `used_param_indices` must be `< params.len()`; - /// violating this panics on out-of-bounds indexing below. Callers should - /// pass `LambdaExpr::used_param_indices()`, which always indexes into the - /// same `params` list, rather than constructing indices by hand. pub fn new( params: Vec, body: Arc, captures: Option, - used_param_indices: &[usize], ) -> Self { - debug_assert!( - used_param_indices.iter().all(|i| *i < params.len()), - "used_param_indices contains an index out of bounds for params \ - (len {}): {:?}", - params.len(), - used_param_indices - ); - - let used_param_indices = used_param_indices.to_vec(); - let effective_params = used_param_indices.iter().map(|i| Arc::clone(¶ms[*i])); - - let fields: Vec = match &captures { + let fields = match &captures { Some(batch) => batch .schema_ref() .fields() .iter() .cloned() - .chain(effective_params) + .chain(params.clone()) .collect(), - None => effective_params.collect(), + None => params.clone(), }; let schema = Arc::new(Schema::new(fields)); Self { params, - used_param_indices, body, schema, captures, @@ -326,11 +287,6 @@ impl LambdaArgument { /// `args` should evaluate to the value of each parameter /// of the correspondent lambda returned in [HigherOrderUDFImpl::lambda_parameters]. /// - /// Only the closures in `args` for parameters the lambda body actually - /// references are called; closures for declared-but-unused parameters - /// are skipped entirely. Callers should not rely on every closure in - /// `args` being invoked. - /// /// `spread_captures` is responsible for transforming the captured column arrays /// so they align with the evaluation batch. Captures are snapshotted from the /// outer batch at construction time, giving one value per outer row, but the @@ -388,7 +344,6 @@ impl LambdaArgument { spread_captures.as_ref(), Arc::clone(&self.schema), &self.params, - &self.used_param_indices, args, )?; @@ -400,7 +355,6 @@ fn merge_captures_with_variables( captures: Option<&RecordBatch>, schema: SchemaRef, params: &[FieldRef], - used_param_indices: &[usize], variables: &[&dyn Fn() -> Result], ) -> Result { if variables.len() < params.len() { @@ -411,42 +365,23 @@ fn merge_captures_with_variables( ); } - let push_param_arrays = |columns: &mut Vec| -> Result<()> { - for &i in used_param_indices { - columns.push(variables[i]()?); - } - Ok(()) - }; - let columns = match captures { Some(captures) => { let mut columns = captures.columns().to_vec(); - push_param_arrays(&mut columns)?; - columns - } - None => { - let mut columns = Vec::with_capacity(used_param_indices.len()); - push_param_arrays(&mut columns)?; + + for arg in &variables[..params.len()] { + columns.push(arg()?); + } + columns } + None => variables + .iter() + .take(params.len()) + .map(|arg| arg()) + .collect::>()?, }; - if columns.is_empty() { - // No columns to derive a row count from, so borrow one variable's - // array length instead (all variables have the same length). - let row_count = variables.first().ok_or_else(|| { - internal_datafusion_err!( - "merge_captures_with_variables: no variables to derive a row count from" - ) - })?()? - .len(); - return Ok(RecordBatch::try_new_with_options( - schema, - vec![], - &RecordBatchOptions::new().with_row_count(Some(row_count)), - )?); - } - Ok(RecordBatch::try_new(schema, columns)?) } @@ -1746,109 +1681,4 @@ mod tests { Some(Arc::new(Field::new(name, dt, nullable))), )) } - - /// A physical expression that reads the column at a fixed index of the - /// batch it is evaluated against, for exercising [`LambdaArgument`] - /// directly without depending on `datafusion-physical-expr`. - #[derive(Debug, Eq, PartialEq, Hash)] - struct ColumnAt(usize); - - impl std::fmt::Display for ColumnAt { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "column_at({})", self.0) - } - } - - impl PhysicalExpr for ColumnAt { - fn evaluate(&self, batch: &RecordBatch) -> Result { - Ok(ColumnarValue::Array(Arc::clone(batch.column(self.0)))) - } - - fn children(&self) -> Vec<&Arc> { - vec![] - } - - fn with_new_children( - self: Arc, - _children: Vec>, - ) -> Result> { - Ok(self) - } - - fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{self}") - } - } - - /// `(k, v) -> v` with only `v` used must push `v`'s array, not `k`'s. - #[test] - fn test_lambda_argument_evaluate_pushes_only_used_param() { - use arrow::array::Int32Array; - - let k_field = Arc::new(Field::new("k", DataType::Int32, true)); - let v_field = Arc::new(Field::new("v", DataType::Int32, true)); - - let body = Arc::new(ColumnAt(0)) as Arc; - let lambda_arg = LambdaArgument::new(vec![k_field, v_field], body, None, &[1]); - - let k_values: ArrayRef = Arc::new(Int32Array::from(vec![100, 200, 300])); - let v_values: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); - let k_closure = || -> Result { Ok(Arc::clone(&k_values)) }; - let v_closure = || -> Result { Ok(Arc::clone(&v_values)) }; - let args: Vec<&dyn Fn() -> Result> = vec![&k_closure, &v_closure]; - - let result = lambda_arg - .evaluate(&args, |arrays| Ok(arrays.to_vec())) - .unwrap(); - let ColumnarValue::Array(result) = result else { - unreachable!() - }; - - assert_eq!( - result.as_any().downcast_ref::().unwrap(), - &Int32Array::from(vec![1, 2, 3]), - "body should read v's values, not k's" - ); - } - - /// Same as above, but with a capture occupying the leading slot. - #[test] - fn test_lambda_argument_evaluate_pushes_only_used_param_with_captures() { - use arrow::array::Int32Array; - - let cap_field = Arc::new(Field::new("cap", DataType::Int32, true)); - let k_field = Arc::new(Field::new("k", DataType::Int32, true)); - let v_field = Arc::new(Field::new("v", DataType::Int32, true)); - - let body = Arc::new(ColumnAt(1)) as Arc; - - let cap_values: ArrayRef = Arc::new(Int32Array::from(vec![9, 9, 9])); - let captures = RecordBatch::try_new( - Arc::new(Schema::new(vec![cap_field])), - vec![cap_values], - ) - .unwrap(); - - let lambda_arg = - LambdaArgument::new(vec![k_field, v_field], body, Some(captures), &[1]); - - let k_values: ArrayRef = Arc::new(Int32Array::from(vec![100, 200, 300])); - let v_values: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); - let k_closure = || -> Result { Ok(Arc::clone(&k_values)) }; - let v_closure = || -> Result { Ok(Arc::clone(&v_values)) }; - let args: Vec<&dyn Fn() -> Result> = vec![&k_closure, &v_closure]; - - let result = lambda_arg - .evaluate(&args, |arrays| Ok(arrays.to_vec())) - .unwrap(); - let ColumnarValue::Array(result) = result else { - unreachable!() - }; - - assert_eq!( - result.as_any().downcast_ref::().unwrap(), - &Int32Array::from(vec![1, 2, 3]), - "body should read v's values, not k's or the capture's" - ); - } } diff --git a/datafusion/physical-expr/src/expressions/lambda.rs b/datafusion/physical-expr/src/expressions/lambda.rs index 02b6c7363adb6..9275821ae9150 100644 --- a/datafusion/physical-expr/src/expressions/lambda.rs +++ b/datafusion/physical-expr/src/expressions/lambda.rs @@ -31,7 +31,7 @@ use arrow::{ }; use datafusion_common::{ HashMap, plan_err, - tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeVisitor}, + tree_node::{Transformed, TreeNode, TreeNodeRecursion}, }; use datafusion_common::{HashSet, Result, internal_err}; use datafusion_expr::ColumnarValue; @@ -43,7 +43,6 @@ pub struct LambdaExpr { body: Arc, projected_body: Arc, projection: Vec, - used_param_indices: Vec, } // Manually derive PartialEq and Hash to work around https://github.com/rust-lang/rust/issues/78808 [https://github.com/apache/datafusion/issues/13196] @@ -61,7 +60,7 @@ impl Hash for LambdaExpr { } impl LambdaExpr { - /// Create a new lambda expression with the given parameters and body. + /// Create a new lambda expression with the given parameters and body pub fn try_new(params: Vec, body: Arc) -> Result { if !all_unique(¶ms) { return plan_err!( @@ -76,30 +75,27 @@ impl LambdaExpr { } fn new(params: Vec, body: Arc) -> Self { - let own_params: HashSet = params.iter().cloned().collect(); + let mut used_column_indices = HashSet::new(); - let mut visitor = CollectUsedVisitor { - own_params: &own_params, - used_indices: HashSet::new(), - used_param_names: HashSet::new(), - shadow_stack: Vec::new(), - }; - body.visit(&mut visitor).expect("visitor is infallible"); - let CollectUsedVisitor { - used_indices, - used_param_names, - .. - } = visitor; + body.apply(|node| { + if let Some(col) = node.downcast_ref::() { + used_column_indices.insert(col.index()); + } else if let Some(var) = node.downcast_ref::() { + used_column_indices.insert(var.index()); + } + + Ok(TreeNodeRecursion::Continue) + }) + .expect("closure should be infallible"); - let mut projection = used_indices.into_iter().collect::>(); + let mut projection = used_column_indices.into_iter().collect::>(); projection.sort(); let column_index_map = projection .iter() - .copied() .enumerate() - .map(|(new_idx, original)| (original, new_idx)) + .map(|(projected, original)| (*original, projected)) .collect::>(); let projected_body = Arc::clone(&body) @@ -128,19 +124,11 @@ impl LambdaExpr { .expect("closure should be infallible") .data; - let used_param_indices = params - .iter() - .enumerate() - .filter(|(_, name)| used_param_names.contains(*name)) - .map(|(i, _)| i) - .collect(); - Self { params, body, projected_body, projection, - used_param_indices, } } @@ -161,75 +149,6 @@ impl LambdaExpr { pub(crate) fn projected_body(&self) -> &Arc { &self.projected_body } - - /// Indices into [`params`](Self::params) of the parameters the body - /// actually references, in declaration order. See `CollectUsedVisitor` - /// in this module. - /// - /// Relies on the planner appending each lambda's own params after - /// captures, matching the `captures ++ used_params` layout - /// `LambdaArgument::new` builds. - pub fn used_param_indices(&self) -> &[usize] { - &self.used_param_indices - } -} - -/// Walks the body of a [`LambdaExpr`] and collects, on a single pass: -/// -/// * `used_indices` — every `Column` / `LambdaVariable` index referenced -/// anywhere in the tree (including inside nested lambdas). This drives -/// the `projection` used to slice the outer batch. -/// * `used_param_names` — the subset of *this* lambda's `own_params` that -/// the body actually references. -/// -/// A nested lambda can declare its own parameter with the same name as -/// one of `own_params` — a distinct variable that happens to reuse the -/// name (variable shadowing). E.g. in -/// `(k, v) -> func(col, (k, v2) -> k + v2 + v)`, the inner `k` is not -/// `own_params`' `k`; only `v` should flow up as used, not `k`. -/// -/// `shadow_stack` holds one frame per nested `LambdaExpr` currently being -/// visited, each frame being that lambda's own parameter names. A -/// `LambdaVariable` only counts toward `used_param_names` if its name -/// isn't in any active frame (i.e. not shadowed). -/// -/// The stack is maintained via `TreeNodeVisitor`'s `f_down` / `f_up`: -/// push a frame when entering a nested [`LambdaExpr`], pop it when leaving. -struct CollectUsedVisitor<'a> { - own_params: &'a HashSet, - used_indices: HashSet, - used_param_names: HashSet, - shadow_stack: Vec>, -} - -impl TreeNodeVisitor<'_> for CollectUsedVisitor<'_> { - type Node = Arc; - - fn f_down(&mut self, node: &Self::Node) -> Result { - if let Some(col) = node.downcast_ref::() { - self.used_indices.insert(col.index()); - } else if let Some(var) = node.downcast_ref::() { - self.used_indices.insert(var.index()); - - let name = var.name(); - let shadowed = self.shadow_stack.iter().any(|frame| frame.contains(name)); - if !shadowed && self.own_params.contains(name) { - self.used_param_names.insert(name.to_string()); - } - } else if let Some(nested) = node.downcast_ref::() { - self.shadow_stack - .push(nested.params.iter().cloned().collect()); - } - - Ok(TreeNodeRecursion::Continue) - } - - fn f_up(&mut self, node: &Self::Node) -> Result { - if node.downcast_ref::().is_some() { - self.shadow_stack.pop(); - } - Ok(TreeNodeRecursion::Continue) - } } impl std::fmt::Display for LambdaExpr { @@ -276,7 +195,7 @@ impl PhysicalExpr for LambdaExpr { } } -/// Create a lambda expression. +/// Create a lambda expression pub fn lambda( params: impl IntoIterator>, body: Arc, @@ -315,15 +234,10 @@ fn check_async_udf(body: &Arc) -> Result<()> { #[cfg(test)] mod tests { - use crate::expressions::{Column, LambdaVariable, NoOp, lambda::lambda}; - use arrow::{ - array::RecordBatch, - datatypes::{DataType, Field, Schema}, - }; + use crate::expressions::{NoOp, lambda::lambda}; + use arrow::{array::RecordBatch, datatypes::Schema}; use std::sync::Arc; - use super::LambdaExpr; - #[test] fn test_lambda_evaluate() { let lambda = lambda(["a"], Arc::new(NoOp::new())).unwrap(); @@ -335,125 +249,4 @@ mod tests { fn test_lambda_duplicate_name() { assert!(lambda(["a", "a"], Arc::new(NoOp::new())).is_err()); } - - /// A two-parameter lambda whose body only references the second - /// parameter (`v`) must report only `v` as used. The higher-order - /// function uses this set to push only `v` into the merged batch, so - /// the body's compressed `LambdaVariable` index for `v` lines up with - /// the batch layout. - #[test] - fn test_used_params_collects_only_referenced_param() { - let v_field = Arc::new(Field::new("v", DataType::Int32, true)); - let body = Arc::new(LambdaVariable::new(1, Arc::clone(&v_field))); - - let lambda = - LambdaExpr::try_new(vec!["k".to_string(), "v".to_string()], body).unwrap(); - - assert_eq!(lambda.projection(), &[1]); - assert_eq!(lambda.used_param_indices(), &[1]); - } - - /// A body that references neither declared parameter reports no used params. - #[test] - fn test_used_params_all_unused() { - let body = Arc::new(NoOp::new()); - - let lambda = - LambdaExpr::try_new(vec!["k".to_string(), "v".to_string()], body).unwrap(); - - assert!(lambda.projection().is_empty()); - assert!(lambda.used_param_indices().is_empty()); - } - - /// A three-parameter lambda that skips the middle parameter reports only the ends as used. - #[test] - fn test_used_params_three_params_middle_unused() { - let a_field = Arc::new(Field::new("a", DataType::Int32, true)); - let c_field = Arc::new(Field::new("c", DataType::Int32, true)); - let body = Arc::new(crate::expressions::BinaryExpr::new( - Arc::new(LambdaVariable::new(0, Arc::clone(&a_field))), - datafusion_expr::Operator::Plus, - Arc::new(LambdaVariable::new(2, Arc::clone(&c_field))), - )); - - let lambda = LambdaExpr::try_new( - vec!["a".to_string(), "b".to_string(), "c".to_string()], - body, - ) - .unwrap(); - - assert_eq!(lambda.used_param_indices(), &[0, 2]); - } - - /// Referencing params out of declaration order still reports both as used. - #[test] - fn test_used_params_both_used_in_reverse_reference_order() { - let k_field = Arc::new(Field::new("k", DataType::Int32, true)); - let v_field = Arc::new(Field::new("v", DataType::Int32, true)); - let body = Arc::new(crate::expressions::BinaryExpr::new( - Arc::new(LambdaVariable::new(1, Arc::clone(&v_field))), - datafusion_expr::Operator::Plus, - Arc::new(LambdaVariable::new(0, Arc::clone(&k_field))), - )); - - let lambda = - LambdaExpr::try_new(vec!["k".to_string(), "v".to_string()], body).unwrap(); - - assert_eq!(lambda.projection(), &[0, 1]); - assert_eq!(lambda.used_param_indices(), &[0, 1]); - } - - /// Inside a nested lambda that re-declares one of the outer parameter - /// names, only the non-shadowed outer references should be reported as - /// used by the outer lambda. In - /// `(k, v) -> func(col, (k, v2) -> k + v2 + v)` the inner `k` shadows - /// the outer `k`, so the outer lambda must only see `v` as used. - #[test] - fn test_used_params_handles_shadowing_inside_nested_lambda() { - let outer_k_field = Arc::new(Field::new("k", DataType::Int32, true)); - let outer_v_field = Arc::new(Field::new("v", DataType::Int32, true)); - let inner_v2_field = Arc::new(Field::new("v2", DataType::Int32, true)); - - // Inner lambda body references "k" (inner's), "v2" (inner's), and - // "v" (outer's). Build it directly with the dense compressed - // indices the inner LambdaExpr::new would produce: sorted referenced - // indices, so the names alone matter here — what matters for - // shadow tracking is the names, not the indices. - let inner_body: Arc = - Arc::new(crate::expressions::BinaryExpr::new( - Arc::new(crate::expressions::BinaryExpr::new( - Arc::new(LambdaVariable::new(1, Arc::clone(&outer_k_field))), - datafusion_expr::Operator::Plus, - Arc::new(LambdaVariable::new(2, Arc::clone(&inner_v2_field))), - )), - datafusion_expr::Operator::Plus, - Arc::new(LambdaVariable::new(0, Arc::clone(&outer_v_field))), - )); - let inner_lambda = Arc::new( - LambdaExpr::try_new(vec!["k".to_string(), "v2".to_string()], inner_body) - .unwrap(), - ); - - // Outer body wraps the inner lambda in a binary op next to a - // regular column reference so the walk has something non-trivial - // to descend through. The outer body references the inner lambda - // via `inner_lambda`. - let outer_body: Arc = - Arc::new(crate::expressions::BinaryExpr::new( - Arc::new(Column::new("col", 0)), - datafusion_expr::Operator::Plus, - inner_lambda, - )); - - let outer_lambda = - LambdaExpr::try_new(vec!["k".to_string(), "v".to_string()], outer_body) - .unwrap(); - - assert_eq!( - outer_lambda.used_param_indices(), - &[1], - "only outer's `v` (index 1) should be reported as used; `k` (index 0) is \ - shadowed inside the nested lambda" - ); - } } diff --git a/datafusion/physical-expr/src/higher_order_function.rs b/datafusion/physical-expr/src/higher_order_function.rs index 53f99a8895349..7390eb33a0922 100644 --- a/datafusion/physical-expr/src/higher_order_function.rs +++ b/datafusion/physical-expr/src/higher_order_function.rs @@ -353,7 +353,6 @@ impl PhysicalExpr for HigherOrderFunctionExpr { } else { Some(batch.project(&projection)?) }, - lambda.used_param_indices(), ))) } ArgSlot::Value => { @@ -510,17 +509,15 @@ mod tests { use super::*; use crate::HigherOrderFunctionExpr; - use crate::create_physical_expr; 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::array::{ArrayRef, Int32Array}; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::Result; use datafusion_common::assert_contains; - use datafusion_expr::execution_props::ExecutionProps; use datafusion_expr::{ HigherOrderFunctionArgs, HigherOrderSignature, HigherOrderUDF, HigherOrderUDFImpl, }; @@ -548,11 +545,9 @@ mod tests { _step: usize, _fields: &[ValueOrLambda>], ) -> Result { - // Offer two params; single-param lambdas just ignore the second. - Ok(LambdaParametersProgress::Complete(vec![vec![ - Arc::new(Field::new("", DataType::Int32, true)), - Arc::new(Field::new("", DataType::Int32, true)), - ]])) + Ok(LambdaParametersProgress::Complete(vec![vec![Arc::new( + Field::new("", DataType::Null, true), + )]])) } fn return_field_from_args( @@ -572,18 +567,7 @@ mod tests { ) -> Result { match &args.args[0] { ValueOrLambda::Lambda(lambda) => lambda.evaluate( - &[ - // Sentinel for the first param, distinct from the second's value. - &|| { - Ok(Arc::new(Int32Array::from(vec![-1000; args.number_rows])) - as ArrayRef) - }, - &|| { - Ok(Arc::new(Int32Array::from_iter_values( - (0..args.number_rows as i32).map(|i| 10 * (i + 1)), - )) as ArrayRef) - }, - ], + &[&|| Ok(Arc::new(NullArray::new(args.number_rows)))], |arrays| Ok(arrays.to_vec()), ), ValueOrLambda::Value(value) => Ok(value.clone()), @@ -731,53 +715,4 @@ mod tests { "mock_function received a lambda via with_new_children at position 0 that wasn't a lambda before" ); } - - /// Exercises the real planner end to end (not hand-picked indices) to - /// check the "captures before own-params" layout invariant. - #[test] - fn test_higher_order_function_two_lambda_params_capture_and_unused_param() { - use datafusion_common::DFSchema; - use datafusion_expr::expr::{HigherOrderFunction, LambdaVariable}; - use datafusion_expr::{Expr, col, lambda as logical_lambda}; - - let fun = Arc::new(HigherOrderUDF::new_from_impl(MockHigherOrderUDF { - signature: HigherOrderSignature::variadic_any(Volatility::Stable), - })); - - // Body uses capture "a" and param "v"; param "k" is left unused. - let v = Expr::LambdaVariable(LambdaVariable::new( - "v".to_string(), - Some(Arc::new(Field::new("v", DataType::Int32, true))), - )); - let body = col("a") + v; - let lambda_expr = logical_lambda(["k", "v"], body); - - let schema = DFSchema::from_unqualified_fields( - vec![Field::new("a", DataType::Int32, false)].into(), - std::collections::HashMap::new(), - ) - .unwrap(); - - let physical_expr = create_physical_expr( - &Expr::HigherOrderFunction(HigherOrderFunction::new(fun, vec![lambda_expr])), - &schema, - &ExecutionProps::new(), - ) - .unwrap(); - - let batch = RecordBatch::try_new( - Arc::clone(schema.inner()), - vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef], - ) - .unwrap(); - - let result = physical_expr.evaluate(&batch).unwrap(); - let ColumnarValue::Array(result) = result else { - unreachable!() - }; - - // a + v; k's sentinel (-1000) must not leak into the result. - let expected = Int32Array::from(vec![11, 22, 33]); - assert_eq!(result.as_ref(), &expected); - } } From f1c9994deb3a206d78965f98008e239ec670a809 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Fri, 31 Jul 2026 01:02:34 -0500 Subject: [PATCH 2/3] fix: keep a CoalescePartitionsExec required by a SinglePartition child (#23948) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - None filed; happy to open one if preferred. A valid query can be planned into a physical plan that `SanityCheckPlan` then rejects: ``` SanityCheckPlan caused by Error during planning: Plan: ["HashJoinExec: mode=CollectLeft, join_type=Left, on=[(id@0, id@0)], projection=[id@0]", " DataSourceExec: file_groups={4 groups: [...]}, projection=[id], file_type=parquet", " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", " CoalescePartitionsExec", " ProjectionExec: expr=[first_value(t.id) ORDER BY [...]@1 as id]", " AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[first_value(t.id) ORDER BY [...]]", " RepartitionExec: partitioning=Hash([id@0], 8), input_partitions=4", " AggregateExec: mode=Partial, gby=[id@1 as id], aggr=[first_value(t.id) ORDER BY [...]]", " DataSourceExec: file_groups={4 groups: [...]}, projection=[ts, id], file_type=parquet"] does not satisfy distribution requirements: SinglePartition. Child-0 output partitioning: UnknownPartitioning(4) ``` The `HashJoinExec` is in `CollectLeft` mode, which requires `Distribution::SinglePartition` on its build (left) child, but child 0 is a bare 4-partition `DataSourceExec` with no `CoalescePartitionsExec` above it. Self-contained reproducer with `datafusion-cli` (the four `COPY` statements are what make the scan multi-partition): ```sql set datafusion.execution.target_partitions = 8; set datafusion.optimizer.repartition_file_scans = false; create table src (id int, ts int) as values (1, 10), (2, 20), (3, 30); copy (select * from src) to 'data/0.parquet' stored as parquet; copy (select * from src) to 'data/1.parquet' stored as parquet; copy (select * from src) to 'data/2.parquet' stored as parquet; copy (select * from src) to 'data/3.parquet' stored as parquet; create external table t stored as parquet location 'data/'; select a.id from t a left join (select distinct on (id) id, ts from t order by id, ts) f on a.id = f.id order by a.id; ``` Setting `datafusion.optimizer.repartition_sorts = false` makes it plan fine, which points at the sort-parallelization phase. `EnsureRequirements` does insert the coalesce for the `SinglePartition` requirement (`enforce_distribution.rs`, `Distribution::SinglePartition => add_merge_on_top(...)`). Its own phase 3a (`parallelize_sorts`) then takes it back out: `remove_bottleneck_in_subplan` removes a `CoalescePartitionsExec` found at `children[0]` positionally, without consulting the parent's distribution requirement for that child. That parent is reached because `update_coalesce_ctx_children` marks a node as connected when *any* child qualifies. It correctly excludes a `SinglePartition`-requiring child from *setting* the flag, but the join's other child (`UnspecifiedDistribution`, connected to a coalesce below) sets it, so the traversal descends into the join and rewrites child 0 anyway. Nothing re-enforces distribution afterwards, so `SanityCheckPlan` is the first thing to notice. Note the surviving `CoalescePartitionsExec` on the probe side in the plan above: it is what propagated the flag, and it is untouched because the `if` returns without recursing into child 1. The sibling helper on the phase 2b path already does consult the requirement (`update_child_to_remove_unnecessary_sort` / `remove_corresponding_sort_from_sub_plan` re-add a merge using the per-child `child_distribution(child_idx)`); only this path is missing it. The same failure shows up with a build child that is already hash-partitioned on the join key (`Child-0 output partitioning: Hash([k@0], 8)`), which is what a `JoinSelection` input swap leaves behind — a `CollectLeft` join reported as `join_type=Right` with an embedded projection. `remove_bottleneck_in_subplan` now checks the parent's per-child distribution requirement before removing a coalesce, both for `children[0]` and when recursing into the other children. The node `parallelize_sorts` is itself rewriting (the root of the call) is exempt, since the caller drops that node and rebuilds the sort cascade around the result — that is the rule's intended transformation, and gating it too would disable sort parallelization below a global sort. This is threaded through as an `is_root` flag on a private `_impl` function; the public entry point keeps its signature. Yes, at two levels: - An end-to-end sqllogictest in `datafusion/sqllogictest/test_files/joins.slt` reproducing it from SQL (the reproducer above, with the data written by `COPY` inside the test). On `main` it fails with exactly the distribution error above. - Two tests in `datafusion/core/tests/physical_optimizer/ensure_requirements.rs` covering both shapes of the build child (`UnknownPartitioning(n)` and `Hash([k], n)`), running the full `EnsureRequirements` rule and then `SanityCheckPlan` via the existing `optimize_and_sanity_check` helper, plus the idempotency check. `cargo test -p datafusion-physical-optimizer`, `cargo test -p datafusion --test core_integration -- physical_optimizer` (530 tests) and the full `sqllogictest` suite (498 files) pass. No API changes. Plans that were previously rejected by `SanityCheckPlan` now plan and execute; a coalesce that is genuinely required is retained where it was previously (incorrectly) removed. --------- Co-authored-by: Claude Opus 5 --- .../src/enforce_sorting/mod.rs | 47 +++++++++++- datafusion/sqllogictest/test_files/joins.slt | 74 +++++++++++++++++++ 2 files changed, 117 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-optimizer/src/enforce_sorting/mod.rs b/datafusion/physical-optimizer/src/enforce_sorting/mod.rs index 241c843556e1c..2adc702e6b1c1 100644 --- a/datafusion/physical-optimizer/src/enforce_sorting/mod.rs +++ b/datafusion/physical-optimizer/src/enforce_sorting/mod.rs @@ -676,11 +676,45 @@ fn adjust_window_sort_removal( /// the plan, some of the remaining `RepartitionExec`s might become unnecessary. /// Removes such `RepartitionExec`s from the plan as well. fn remove_bottleneck_in_subplan( + requirements: PlanWithCorrespondingCoalescePartitions, +) -> Result { + // The root is the node `parallelize_sorts` is rewriting (a `SortExec`, + // `SortPreservingMergeExec` or `CoalescePartitionsExec`). Its own distribution + // requirement does not constrain the removal, because the caller drops the node and + // rebuilds the cascade around the result. + remove_bottleneck_in_subplan_impl(requirements, true) +} + +fn remove_bottleneck_in_subplan_impl( mut requirements: PlanWithCorrespondingCoalescePartitions, + is_root: bool, ) -> Result { let plan = &requirements.plan; + // Below the root, a `CoalescePartitionsExec` feeding a child that requires + // `Distribution::SinglePartition` is not an avoidable bottleneck: it is what satisfies + // that requirement. Removing it leaves the parent with a multi-partition input it cannot + // accept, and nothing re-runs distribution enforcement afterwards, so the plan reaches + // `SanityCheckPlan` invalid. The traversal reaches such a node because + // `update_coalesce_ctx_children` marks a node as connected when *any* child qualifies: + // a `CollectLeft` `HashJoinExec` whose probe side is connected is descended into even + // though its build side must stay single-partition. + // + // Only `SinglePartition` is protected. A `HashPartitioned` child is in principle in the + // same position — a single-partition input trivially satisfies a hash requirement, so a + // coalesce below one is also load-bearing — but nothing puts a coalesce there: + // `ensure_distribution` satisfies a hash requirement with a `RepartitionExec`, never a + // `CoalescePartitionsExec`. Widening the check would be dead code today. + let dist_reqs = plan.required_input_distribution(); + let removable = |idx: usize| { + is_root || !matches!(dist_reqs.get(idx), Some(Distribution::SinglePartition)) + }; + let remove_from_first_child = requirements + .children + .first() + .is_some_and(|child| is_coalesce_partitions(&child.plan)) + && removable(0); let children = &mut requirements.children; - if is_coalesce_partitions(&children[0].plan) { + if remove_from_first_child { // We can safely use the 0th index since we have a `CoalescePartitionsExec`. let mut new_child_node = children[0].children.swap_remove(0); while new_child_node.plan.output_partitioning() == plan.output_partitioning() @@ -694,9 +728,14 @@ fn remove_bottleneck_in_subplan( requirements.children = requirements .children .into_iter() - .map(|node| { - if node.data { - remove_bottleneck_in_subplan(node) + .enumerate() + .map(|(idx, node)| { + // Deliberately conservative: not descending at all also skips legitimate + // cleanups *below* a protected child (a redundant second coalesce under the + // load-bearing one, say). This could later be narrowed to "descend, but + // protect only the topmost coalesce" if that turns out to matter. + if node.data && removable(idx) { + remove_bottleneck_in_subplan_impl(node, false) } else { Ok(node) } diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt index e0be63fe71525..3495158fc3c65 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/datafusion/sqllogictest/test_files/joins.slt @@ -5527,3 +5527,77 @@ DROP TABLE t1; statement ok DROP TABLE t2; + +# Regression test: a `CollectLeft` `HashJoinExec` requires `SinglePartition` on its build +# (left) child, and the `CoalescePartitionsExec` that satisfies it must survive the +# sort-parallelization phase of `EnsureRequirements`. It used to be removed positionally +# (the traversal descends into the join because the *probe* side is linked to a coalesce), +# leaving a multi-partition build side that `SanityCheckPlan` rejects with +# "does not satisfy distribution requirements: SinglePartition". + +statement ok +set datafusion.execution.target_partitions = 8; + +# Keep the scan multi-partition as written, i.e. one partition per file. +statement ok +set datafusion.optimizer.repartition_file_scans = false; + +statement ok +CREATE TABLE collect_left_src (id INT, ts INT) AS VALUES (1, 10), (2, 20), (3, 30); + +query I +COPY (SELECT * FROM collect_left_src) TO 'test_files/scratch/joins/collect_left/0.parquet' STORED AS PARQUET; +---- +3 + +query I +COPY (SELECT * FROM collect_left_src) TO 'test_files/scratch/joins/collect_left/1.parquet' STORED AS PARQUET; +---- +3 + +query I +COPY (SELECT * FROM collect_left_src) TO 'test_files/scratch/joins/collect_left/2.parquet' STORED AS PARQUET; +---- +3 + +query I +COPY (SELECT * FROM collect_left_src) TO 'test_files/scratch/joins/collect_left/3.parquet' STORED AS PARQUET; +---- +3 + +statement ok +CREATE EXTERNAL TABLE collect_left STORED AS PARQUET LOCATION 'test_files/scratch/joins/collect_left/'; + +# The build side is the 4-partition scan; the probe side is the `DISTINCT ON` aggregate, +# whose `CoalescePartitionsExec` is what makes the traversal reach the join. +query I +SELECT a.id +FROM collect_left a +LEFT JOIN (SELECT DISTINCT ON (id) id, ts FROM collect_left ORDER BY id, ts) f + ON a.id = f.id +ORDER BY a.id; +---- +1 +1 +1 +1 +2 +2 +2 +2 +3 +3 +3 +3 + +statement ok +DROP TABLE collect_left; + +statement ok +DROP TABLE collect_left_src; + +statement ok +reset datafusion.optimizer.repartition_file_scans; + +statement ok +set datafusion.execution.target_partitions = 4; From 0d6aff7bd96b1381b37d9a67b31d392ed22936ea Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=ADa=20Adriana?= Date: Wed, 12 Aug 2026 10:54:38 +0200 Subject: [PATCH 3/3] fix(lambda): only push referenced params into the merged batch (#24162) (#166) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(lambda): only push referenced params into the merged batch (#24162) ## Which issue does this PR close? basically this PR https://github.com/apache/datafusion/pull/22853 + a few more tests ## Rationale for this change The current lambdas in DF only take a single parameter `(v -> ...)`, so nobody had noticed that `LambdaExpr` mishandles lambdas with more than one parameter. The bug surfaced while working on `transform_values` (#22689), which needs `(k, v) -> expr ` two parameters, one of which is very often unused (e.g. `(k, v) -> v * 2`, k never referenced). The bug is that when a higher order function with more than 1 param evaluates a lambda, it fills each parameter into a slot based on its declared position — for example for `(k, v) -> v` `k` always goes into slot 0, `v` always into slot 1. `LambdaExpr` separately scans the body and renumbers whatever it finds referenced into a dense `0..n` range, to avoid carrying around columns nothing uses (like `v` in this case). That renumbering is fine for outer captures, but applying it to the lambda's own parameters is wrong, because it changes where the body looks for a value without changing where the evaluator put it. ### Example: in `(k, v) -> v` `v` is declared second (slot 1), but since it's the only parameter the body references, the renumbering logic reassigns it to slot 0. The evaluator, unaware of this, writes `k`'s values into slot 0 and `v`'s into slot 1. So the body ends up reading slot 0 expecting `v` — and gets `k` instead. So the results end up being incorrect. ## What changes are included in this PR? - `LambdaExpr` now computes `used_params`: which is the subset of its own declared parameters that are actually referenced in the body. - `LambdaArgument::new` takes `used_params` and only pushes the referenced parameters in the body into the merged batch, in original declaration order — so the body's indices always line up with what's actually built. - `HigherOrderFunctionExpr::evaluate` forwards `lambda.used_params()` to `LambdaArgument::new` ## Are these changes tested? yes, added two new tests one for the unused-parameter case and nested-lambda for the shadowing case. ## Are there any user-facing changes? The only public api change is on `LambdaArgument::new ` which now requires a new argument: `used_params: &HashSet`, however LambdaArgument::new is very unlikely to be called outside datafusion, see [this](https://github.com/apache/datafusion/pull/22853#discussion_r3527236493) comment (cherry picked from commit 4e6acfe8ee7e5da38f1f1d55427989b86f70a775) * Adjust to API change --- datafusion/expr/src/higher_order_function.rs | 198 +++++++++++++- .../physical-expr/src/expressions/lambda.rs | 243 ++++++++++++++++-- .../src/higher_order_function.rs | 75 +++++- 3 files changed, 479 insertions(+), 37 deletions(-) diff --git a/datafusion/expr/src/higher_order_function.rs b/datafusion/expr/src/higher_order_function.rs index 413714f498164..c300be8f6fcfe 100644 --- a/datafusion/expr/src/higher_order_function.rs +++ b/datafusion/expr/src/higher_order_function.rs @@ -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, RecordBatch, RecordBatchOptions}; use arrow::datatypes::{DataType, FieldRef, Schema}; use arrow_schema::SchemaRef; use datafusion_common::config::ConfigOptions; @@ -239,6 +239,26 @@ pub struct LambdaArgument { /// For example, for `array_transform([2], v -> -v)`, /// this will be `vec![Field::new("v", DataType::Int32, true)]` params: Vec, + /// Indices into [`Self::params`] of the parameters that are actually + /// referenced by [`Self::body`] (taking nested-lambda shadowing into + /// account), in the original declaration order of `params`. + /// + /// [`Self::evaluate`] only evaluates and pushes the closures whose + /// corresponding parameter index appears here, so unused declared + /// parameters leave no slot in the merged batch and the body's compressed + /// column indices line up directly with what the evaluator built. + /// + /// Callers who already have a `LambdaExpr` should pass + /// `LambdaExpr::used_param_indices()` directly to [`Self::new`] — both + /// are indices into the same positionally-aligned `params` list. + /// + /// Every index here must be `< params.len()`; see the precondition on + /// [`Self::new`]. + /// + /// Relies on captures sorting before this lambda's own params in the + /// planner's (un-projected) index space, which is what makes + /// `captures ++ used_params` below line up with the projected body. + used_param_indices: Vec, /// The body of the lambda /// /// For example, for `array_transform([2], v -> -v)`, @@ -257,26 +277,45 @@ pub struct LambdaArgument { } impl LambdaArgument { + /// # Preconditions + /// + /// Every index in `used_param_indices` must be `< params.len()`; + /// violating this panics on out-of-bounds indexing below. Callers should + /// pass `LambdaExpr::used_param_indices()`, which always indexes into the + /// same `params` list, rather than constructing indices by hand. pub fn new( params: Vec, body: Arc, captures: Option, + used_param_indices: &[usize], ) -> Self { - let fields = match &captures { + debug_assert!( + used_param_indices.iter().all(|i| *i < params.len()), + "used_param_indices contains an index out of bounds for params \ + (len {}): {:?}", + params.len(), + used_param_indices + ); + + let used_param_indices = used_param_indices.to_vec(); + let effective_params = used_param_indices.iter().map(|i| Arc::clone(¶ms[*i])); + + let fields: Vec = match &captures { Some(batch) => batch .schema_ref() .fields() .iter() .cloned() - .chain(params.clone()) + .chain(effective_params) .collect(), - None => params.clone(), + None => effective_params.collect(), }; let schema = Arc::new(Schema::new(fields)); Self { params, + used_param_indices, body, schema, captures, @@ -287,6 +326,11 @@ impl LambdaArgument { /// `args` should evaluate to the value of each parameter /// of the correspondent lambda returned in [HigherOrderUDFImpl::lambda_parameters]. /// + /// Only the closures in `args` for parameters the lambda body actually + /// references are called; closures for declared-but-unused parameters + /// are skipped entirely. Callers should not rely on every closure in + /// `args` being invoked. + /// /// `spread_captures` is responsible for transforming the captured column arrays /// so they align with the evaluation batch. Captures are snapshotted from the /// outer batch at construction time, giving one value per outer row, but the @@ -344,6 +388,7 @@ impl LambdaArgument { spread_captures.as_ref(), Arc::clone(&self.schema), &self.params, + &self.used_param_indices, args, )?; @@ -355,6 +400,7 @@ fn merge_captures_with_variables( captures: Option<&RecordBatch>, schema: SchemaRef, params: &[FieldRef], + used_param_indices: &[usize], variables: &[&dyn Fn() -> Result], ) -> Result { if variables.len() < params.len() { @@ -365,23 +411,42 @@ fn merge_captures_with_variables( ); } + let push_param_arrays = |columns: &mut Vec| -> Result<()> { + for &i in used_param_indices { + columns.push(variables[i]()?); + } + Ok(()) + }; + let columns = match captures { Some(captures) => { let mut columns = captures.columns().to_vec(); - - for arg in &variables[..params.len()] { - columns.push(arg()?); - } - + push_param_arrays(&mut columns)?; + columns + } + None => { + let mut columns = Vec::with_capacity(used_param_indices.len()); + push_param_arrays(&mut columns)?; columns } - None => variables - .iter() - .take(params.len()) - .map(|arg| arg()) - .collect::>()?, }; + if columns.is_empty() { + // No columns to derive a row count from, so borrow one variable's + // array length instead (all variables have the same length). + let row_count = variables.first().ok_or_else(|| { + internal_datafusion_err!( + "merge_captures_with_variables: no variables to derive a row count from" + ) + })?()? + .len(); + return Ok(RecordBatch::try_new_with_options( + schema, + vec![], + &RecordBatchOptions::new().with_row_count(Some(row_count)), + )?); + } + Ok(RecordBatch::try_new(schema, columns)?) } @@ -1681,4 +1746,109 @@ mod tests { Some(Arc::new(Field::new(name, dt, nullable))), )) } + + /// A physical expression that reads the column at a fixed index of the + /// batch it is evaluated against, for exercising [`LambdaArgument`] + /// directly without depending on `datafusion-physical-expr`. + #[derive(Debug, Eq, PartialEq, Hash)] + struct ColumnAt(usize); + + impl std::fmt::Display for ColumnAt { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "column_at({})", self.0) + } + } + + impl PhysicalExpr for ColumnAt { + fn evaluate(&self, batch: &RecordBatch) -> Result { + Ok(ColumnarValue::Array(Arc::clone(batch.column(self.0)))) + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Ok(self) + } + + fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{self}") + } + } + + /// `(k, v) -> v` with only `v` used must push `v`'s array, not `k`'s. + #[test] + fn test_lambda_argument_evaluate_pushes_only_used_param() { + use arrow::array::Int32Array; + + let k_field = Arc::new(Field::new("k", DataType::Int32, true)); + let v_field = Arc::new(Field::new("v", DataType::Int32, true)); + + let body = Arc::new(ColumnAt(0)) as Arc; + let lambda_arg = LambdaArgument::new(vec![k_field, v_field], body, None, &[1]); + + let k_values: ArrayRef = Arc::new(Int32Array::from(vec![100, 200, 300])); + let v_values: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let k_closure = || -> Result { Ok(Arc::clone(&k_values)) }; + let v_closure = || -> Result { Ok(Arc::clone(&v_values)) }; + let args: Vec<&dyn Fn() -> Result> = vec![&k_closure, &v_closure]; + + let result = lambda_arg + .evaluate(&args, |arrays| Ok(arrays.to_vec())) + .unwrap(); + let ColumnarValue::Array(result) = result else { + unreachable!() + }; + + assert_eq!( + result.as_any().downcast_ref::().unwrap(), + &Int32Array::from(vec![1, 2, 3]), + "body should read v's values, not k's" + ); + } + + /// Same as above, but with a capture occupying the leading slot. + #[test] + fn test_lambda_argument_evaluate_pushes_only_used_param_with_captures() { + use arrow::array::Int32Array; + + let cap_field = Arc::new(Field::new("cap", DataType::Int32, true)); + let k_field = Arc::new(Field::new("k", DataType::Int32, true)); + let v_field = Arc::new(Field::new("v", DataType::Int32, true)); + + let body = Arc::new(ColumnAt(1)) as Arc; + + let cap_values: ArrayRef = Arc::new(Int32Array::from(vec![9, 9, 9])); + let captures = RecordBatch::try_new( + Arc::new(Schema::new(vec![cap_field])), + vec![cap_values], + ) + .unwrap(); + + let lambda_arg = + LambdaArgument::new(vec![k_field, v_field], body, Some(captures), &[1]); + + let k_values: ArrayRef = Arc::new(Int32Array::from(vec![100, 200, 300])); + let v_values: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let k_closure = || -> Result { Ok(Arc::clone(&k_values)) }; + let v_closure = || -> Result { Ok(Arc::clone(&v_values)) }; + let args: Vec<&dyn Fn() -> Result> = vec![&k_closure, &v_closure]; + + let result = lambda_arg + .evaluate(&args, |arrays| Ok(arrays.to_vec())) + .unwrap(); + let ColumnarValue::Array(result) = result else { + unreachable!() + }; + + assert_eq!( + result.as_any().downcast_ref::().unwrap(), + &Int32Array::from(vec![1, 2, 3]), + "body should read v's values, not k's or the capture's" + ); + } } diff --git a/datafusion/physical-expr/src/expressions/lambda.rs b/datafusion/physical-expr/src/expressions/lambda.rs index 9275821ae9150..02b6c7363adb6 100644 --- a/datafusion/physical-expr/src/expressions/lambda.rs +++ b/datafusion/physical-expr/src/expressions/lambda.rs @@ -31,7 +31,7 @@ use arrow::{ }; use datafusion_common::{ HashMap, plan_err, - tree_node::{Transformed, TreeNode, TreeNodeRecursion}, + tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeVisitor}, }; use datafusion_common::{HashSet, Result, internal_err}; use datafusion_expr::ColumnarValue; @@ -43,6 +43,7 @@ pub struct LambdaExpr { body: Arc, projected_body: Arc, projection: Vec, + used_param_indices: Vec, } // Manually derive PartialEq and Hash to work around https://github.com/rust-lang/rust/issues/78808 [https://github.com/apache/datafusion/issues/13196] @@ -60,7 +61,7 @@ impl Hash for LambdaExpr { } impl LambdaExpr { - /// Create a new lambda expression with the given parameters and body + /// Create a new lambda expression with the given parameters and body. pub fn try_new(params: Vec, body: Arc) -> Result { if !all_unique(¶ms) { return plan_err!( @@ -75,27 +76,30 @@ impl LambdaExpr { } fn new(params: Vec, body: Arc) -> Self { - let mut used_column_indices = HashSet::new(); + let own_params: HashSet = params.iter().cloned().collect(); - body.apply(|node| { - if let Some(col) = node.downcast_ref::() { - used_column_indices.insert(col.index()); - } else if let Some(var) = node.downcast_ref::() { - used_column_indices.insert(var.index()); - } - - Ok(TreeNodeRecursion::Continue) - }) - .expect("closure should be infallible"); + let mut visitor = CollectUsedVisitor { + own_params: &own_params, + used_indices: HashSet::new(), + used_param_names: HashSet::new(), + shadow_stack: Vec::new(), + }; + body.visit(&mut visitor).expect("visitor is infallible"); + let CollectUsedVisitor { + used_indices, + used_param_names, + .. + } = visitor; - let mut projection = used_column_indices.into_iter().collect::>(); + let mut projection = used_indices.into_iter().collect::>(); projection.sort(); let column_index_map = projection .iter() + .copied() .enumerate() - .map(|(projected, original)| (*original, projected)) + .map(|(new_idx, original)| (original, new_idx)) .collect::>(); let projected_body = Arc::clone(&body) @@ -124,11 +128,19 @@ impl LambdaExpr { .expect("closure should be infallible") .data; + let used_param_indices = params + .iter() + .enumerate() + .filter(|(_, name)| used_param_names.contains(*name)) + .map(|(i, _)| i) + .collect(); + Self { params, body, projected_body, projection, + used_param_indices, } } @@ -149,6 +161,75 @@ impl LambdaExpr { pub(crate) fn projected_body(&self) -> &Arc { &self.projected_body } + + /// Indices into [`params`](Self::params) of the parameters the body + /// actually references, in declaration order. See `CollectUsedVisitor` + /// in this module. + /// + /// Relies on the planner appending each lambda's own params after + /// captures, matching the `captures ++ used_params` layout + /// `LambdaArgument::new` builds. + pub fn used_param_indices(&self) -> &[usize] { + &self.used_param_indices + } +} + +/// Walks the body of a [`LambdaExpr`] and collects, on a single pass: +/// +/// * `used_indices` — every `Column` / `LambdaVariable` index referenced +/// anywhere in the tree (including inside nested lambdas). This drives +/// the `projection` used to slice the outer batch. +/// * `used_param_names` — the subset of *this* lambda's `own_params` that +/// the body actually references. +/// +/// A nested lambda can declare its own parameter with the same name as +/// one of `own_params` — a distinct variable that happens to reuse the +/// name (variable shadowing). E.g. in +/// `(k, v) -> func(col, (k, v2) -> k + v2 + v)`, the inner `k` is not +/// `own_params`' `k`; only `v` should flow up as used, not `k`. +/// +/// `shadow_stack` holds one frame per nested `LambdaExpr` currently being +/// visited, each frame being that lambda's own parameter names. A +/// `LambdaVariable` only counts toward `used_param_names` if its name +/// isn't in any active frame (i.e. not shadowed). +/// +/// The stack is maintained via `TreeNodeVisitor`'s `f_down` / `f_up`: +/// push a frame when entering a nested [`LambdaExpr`], pop it when leaving. +struct CollectUsedVisitor<'a> { + own_params: &'a HashSet, + used_indices: HashSet, + used_param_names: HashSet, + shadow_stack: Vec>, +} + +impl TreeNodeVisitor<'_> for CollectUsedVisitor<'_> { + type Node = Arc; + + fn f_down(&mut self, node: &Self::Node) -> Result { + if let Some(col) = node.downcast_ref::() { + self.used_indices.insert(col.index()); + } else if let Some(var) = node.downcast_ref::() { + self.used_indices.insert(var.index()); + + let name = var.name(); + let shadowed = self.shadow_stack.iter().any(|frame| frame.contains(name)); + if !shadowed && self.own_params.contains(name) { + self.used_param_names.insert(name.to_string()); + } + } else if let Some(nested) = node.downcast_ref::() { + self.shadow_stack + .push(nested.params.iter().cloned().collect()); + } + + Ok(TreeNodeRecursion::Continue) + } + + fn f_up(&mut self, node: &Self::Node) -> Result { + if node.downcast_ref::().is_some() { + self.shadow_stack.pop(); + } + Ok(TreeNodeRecursion::Continue) + } } impl std::fmt::Display for LambdaExpr { @@ -195,7 +276,7 @@ impl PhysicalExpr for LambdaExpr { } } -/// Create a lambda expression +/// Create a lambda expression. pub fn lambda( params: impl IntoIterator>, body: Arc, @@ -234,10 +315,15 @@ fn check_async_udf(body: &Arc) -> Result<()> { #[cfg(test)] mod tests { - use crate::expressions::{NoOp, lambda::lambda}; - use arrow::{array::RecordBatch, datatypes::Schema}; + use crate::expressions::{Column, LambdaVariable, NoOp, lambda::lambda}; + use arrow::{ + array::RecordBatch, + datatypes::{DataType, Field, Schema}, + }; use std::sync::Arc; + use super::LambdaExpr; + #[test] fn test_lambda_evaluate() { let lambda = lambda(["a"], Arc::new(NoOp::new())).unwrap(); @@ -249,4 +335,125 @@ mod tests { fn test_lambda_duplicate_name() { assert!(lambda(["a", "a"], Arc::new(NoOp::new())).is_err()); } + + /// A two-parameter lambda whose body only references the second + /// parameter (`v`) must report only `v` as used. The higher-order + /// function uses this set to push only `v` into the merged batch, so + /// the body's compressed `LambdaVariable` index for `v` lines up with + /// the batch layout. + #[test] + fn test_used_params_collects_only_referenced_param() { + let v_field = Arc::new(Field::new("v", DataType::Int32, true)); + let body = Arc::new(LambdaVariable::new(1, Arc::clone(&v_field))); + + let lambda = + LambdaExpr::try_new(vec!["k".to_string(), "v".to_string()], body).unwrap(); + + assert_eq!(lambda.projection(), &[1]); + assert_eq!(lambda.used_param_indices(), &[1]); + } + + /// A body that references neither declared parameter reports no used params. + #[test] + fn test_used_params_all_unused() { + let body = Arc::new(NoOp::new()); + + let lambda = + LambdaExpr::try_new(vec!["k".to_string(), "v".to_string()], body).unwrap(); + + assert!(lambda.projection().is_empty()); + assert!(lambda.used_param_indices().is_empty()); + } + + /// A three-parameter lambda that skips the middle parameter reports only the ends as used. + #[test] + fn test_used_params_three_params_middle_unused() { + let a_field = Arc::new(Field::new("a", DataType::Int32, true)); + let c_field = Arc::new(Field::new("c", DataType::Int32, true)); + let body = Arc::new(crate::expressions::BinaryExpr::new( + Arc::new(LambdaVariable::new(0, Arc::clone(&a_field))), + datafusion_expr::Operator::Plus, + Arc::new(LambdaVariable::new(2, Arc::clone(&c_field))), + )); + + let lambda = LambdaExpr::try_new( + vec!["a".to_string(), "b".to_string(), "c".to_string()], + body, + ) + .unwrap(); + + assert_eq!(lambda.used_param_indices(), &[0, 2]); + } + + /// Referencing params out of declaration order still reports both as used. + #[test] + fn test_used_params_both_used_in_reverse_reference_order() { + let k_field = Arc::new(Field::new("k", DataType::Int32, true)); + let v_field = Arc::new(Field::new("v", DataType::Int32, true)); + let body = Arc::new(crate::expressions::BinaryExpr::new( + Arc::new(LambdaVariable::new(1, Arc::clone(&v_field))), + datafusion_expr::Operator::Plus, + Arc::new(LambdaVariable::new(0, Arc::clone(&k_field))), + )); + + let lambda = + LambdaExpr::try_new(vec!["k".to_string(), "v".to_string()], body).unwrap(); + + assert_eq!(lambda.projection(), &[0, 1]); + assert_eq!(lambda.used_param_indices(), &[0, 1]); + } + + /// Inside a nested lambda that re-declares one of the outer parameter + /// names, only the non-shadowed outer references should be reported as + /// used by the outer lambda. In + /// `(k, v) -> func(col, (k, v2) -> k + v2 + v)` the inner `k` shadows + /// the outer `k`, so the outer lambda must only see `v` as used. + #[test] + fn test_used_params_handles_shadowing_inside_nested_lambda() { + let outer_k_field = Arc::new(Field::new("k", DataType::Int32, true)); + let outer_v_field = Arc::new(Field::new("v", DataType::Int32, true)); + let inner_v2_field = Arc::new(Field::new("v2", DataType::Int32, true)); + + // Inner lambda body references "k" (inner's), "v2" (inner's), and + // "v" (outer's). Build it directly with the dense compressed + // indices the inner LambdaExpr::new would produce: sorted referenced + // indices, so the names alone matter here — what matters for + // shadow tracking is the names, not the indices. + let inner_body: Arc = + Arc::new(crate::expressions::BinaryExpr::new( + Arc::new(crate::expressions::BinaryExpr::new( + Arc::new(LambdaVariable::new(1, Arc::clone(&outer_k_field))), + datafusion_expr::Operator::Plus, + Arc::new(LambdaVariable::new(2, Arc::clone(&inner_v2_field))), + )), + datafusion_expr::Operator::Plus, + Arc::new(LambdaVariable::new(0, Arc::clone(&outer_v_field))), + )); + let inner_lambda = Arc::new( + LambdaExpr::try_new(vec!["k".to_string(), "v2".to_string()], inner_body) + .unwrap(), + ); + + // Outer body wraps the inner lambda in a binary op next to a + // regular column reference so the walk has something non-trivial + // to descend through. The outer body references the inner lambda + // via `inner_lambda`. + let outer_body: Arc = + Arc::new(crate::expressions::BinaryExpr::new( + Arc::new(Column::new("col", 0)), + datafusion_expr::Operator::Plus, + inner_lambda, + )); + + let outer_lambda = + LambdaExpr::try_new(vec!["k".to_string(), "v".to_string()], outer_body) + .unwrap(); + + assert_eq!( + outer_lambda.used_param_indices(), + &[1], + "only outer's `v` (index 1) should be reported as used; `k` (index 0) is \ + shadowed inside the nested lambda" + ); + } } diff --git a/datafusion/physical-expr/src/higher_order_function.rs b/datafusion/physical-expr/src/higher_order_function.rs index 7390eb33a0922..53f99a8895349 100644 --- a/datafusion/physical-expr/src/higher_order_function.rs +++ b/datafusion/physical-expr/src/higher_order_function.rs @@ -353,6 +353,7 @@ impl PhysicalExpr for HigherOrderFunctionExpr { } else { Some(batch.project(&projection)?) }, + lambda.used_param_indices(), ))) } ArgSlot::Value => { @@ -509,15 +510,17 @@ mod tests { use super::*; use crate::HigherOrderFunctionExpr; + use crate::create_physical_expr; 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::array::{ArrayRef, Int32Array}; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::Result; use datafusion_common::assert_contains; + use datafusion_expr::execution_props::ExecutionProps; use datafusion_expr::{ HigherOrderFunctionArgs, HigherOrderSignature, HigherOrderUDF, HigherOrderUDFImpl, }; @@ -545,9 +548,11 @@ mod tests { _step: usize, _fields: &[ValueOrLambda>], ) -> Result { - Ok(LambdaParametersProgress::Complete(vec![vec![Arc::new( - Field::new("", DataType::Null, true), - )]])) + // Offer two params; single-param lambdas just ignore the second. + Ok(LambdaParametersProgress::Complete(vec![vec![ + Arc::new(Field::new("", DataType::Int32, true)), + Arc::new(Field::new("", DataType::Int32, true)), + ]])) } fn return_field_from_args( @@ -567,7 +572,18 @@ mod tests { ) -> Result { match &args.args[0] { ValueOrLambda::Lambda(lambda) => lambda.evaluate( - &[&|| Ok(Arc::new(NullArray::new(args.number_rows)))], + &[ + // Sentinel for the first param, distinct from the second's value. + &|| { + Ok(Arc::new(Int32Array::from(vec![-1000; args.number_rows])) + as ArrayRef) + }, + &|| { + Ok(Arc::new(Int32Array::from_iter_values( + (0..args.number_rows as i32).map(|i| 10 * (i + 1)), + )) as ArrayRef) + }, + ], |arrays| Ok(arrays.to_vec()), ), ValueOrLambda::Value(value) => Ok(value.clone()), @@ -715,4 +731,53 @@ mod tests { "mock_function received a lambda via with_new_children at position 0 that wasn't a lambda before" ); } + + /// Exercises the real planner end to end (not hand-picked indices) to + /// check the "captures before own-params" layout invariant. + #[test] + fn test_higher_order_function_two_lambda_params_capture_and_unused_param() { + use datafusion_common::DFSchema; + use datafusion_expr::expr::{HigherOrderFunction, LambdaVariable}; + use datafusion_expr::{Expr, col, lambda as logical_lambda}; + + let fun = Arc::new(HigherOrderUDF::new_from_impl(MockHigherOrderUDF { + signature: HigherOrderSignature::variadic_any(Volatility::Stable), + })); + + // Body uses capture "a" and param "v"; param "k" is left unused. + let v = Expr::LambdaVariable(LambdaVariable::new( + "v".to_string(), + Some(Arc::new(Field::new("v", DataType::Int32, true))), + )); + let body = col("a") + v; + let lambda_expr = logical_lambda(["k", "v"], body); + + let schema = DFSchema::from_unqualified_fields( + vec![Field::new("a", DataType::Int32, false)].into(), + std::collections::HashMap::new(), + ) + .unwrap(); + + let physical_expr = create_physical_expr( + &Expr::HigherOrderFunction(HigherOrderFunction::new(fun, vec![lambda_expr])), + &schema, + &ExecutionProps::new(), + ) + .unwrap(); + + let batch = RecordBatch::try_new( + Arc::clone(schema.inner()), + vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef], + ) + .unwrap(); + + let result = physical_expr.evaluate(&batch).unwrap(); + let ColumnarValue::Array(result) = result else { + unreachable!() + }; + + // a + v; k's sentinel (-1000) must not leak into the result. + let expected = Int32Array::from(vec![11, 22, 33]); + assert_eq!(result.as_ref(), &expected); + } }