diff --git a/datafusion/physical-expr/src/expressions/lambda.rs b/datafusion/physical-expr/src/expressions/lambda.rs index 95bb5db0b328e..2cb731811a8e2 100644 --- a/datafusion/physical-expr/src/expressions/lambda.rs +++ b/datafusion/physical-expr/src/expressions/lambda.rs @@ -31,7 +31,9 @@ use arrow::{ }; use datafusion_common::{ HashMap, plan_err, - tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeVisitor}, + tree_node::{ + Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter, TreeNodeVisitor, + }, }; use datafusion_common::{HashSet, Result, internal_err}; use datafusion_expr::ColumnarValue; @@ -91,6 +93,12 @@ impl LambdaExpr { .. } = visitor; + // Capture projection is *only* outer-batch columns (and captured + // outer lambda variables). This lambda's own parameters are appended + // after those captures by `LambdaArgument::new` and must not share + // the outer-batch index space — otherwise a later projection rewrite + // that changes the input schema width can make a stale param index + // look like a real input column (see #24372). let mut projection = used_indices.into_iter().collect::>(); projection.sort(); @@ -102,39 +110,33 @@ impl LambdaExpr { .map(|(new_idx, original)| (original, new_idx)) .collect::>(); - let projected_body = Arc::clone(&body) - .transform_down(|e| { - if let Some(column) = e.downcast_ref::() { - 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::() { - 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; - - let used_param_indices = params + let used_param_indices: Vec = params .iter() .enumerate() .filter(|(_, name)| used_param_names.contains(*name)) .map(|(i, _)| i) .collect(); + // Own params occupy slots after the captured columns: + // `captures ++ used_params`. Bind them by name so a stale + // planner-assigned combined-schema index cannot leak into the + // capture projection after a schema-boundary rewrite. + let n_captures = projection.len(); + let param_name_to_projected: HashMap = used_param_indices + .iter() + .enumerate() + .map(|(pos, param_i)| (params[*param_i].clone(), n_captures + pos)) + .collect(); + + let projected_body = Arc::clone(&body) + .rewrite(&mut ProjectLambdaBody { + column_index_map: &column_index_map, + param_name_to_projected: ¶m_name_to_projected, + shadow_stack: Vec::new(), + }) + .expect("rewriter is infallible") + .data; + Self { params, body, @@ -197,9 +199,12 @@ impl LambdaExpr { /// 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_indices` — every outer-batch `Column` index, plus indices of +/// captured *outer* `LambdaVariable`s (not this lambda's own params). +/// This drives the `projection` used to slice the outer batch. Own +/// parameters are deliberately excluded so their planner-assigned +/// combined-schema indexes cannot be mistaken for input columns after +/// a later projection rewrite changes the input schema width (#24372). /// * `used_param_names` — the subset of *this* lambda's `own_params` that /// the body actually references. /// @@ -230,12 +235,15 @@ impl TreeNodeVisitor<'_> for CollectUsedVisitor<'_> { 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 !shadowed { + // Captured outer lambda variable. Keep its index in the + // outer-batch projection. Nested lambdas' own params are + // shadowed here and remapped by the inner LambdaExpr. + self.used_indices.insert(var.index()); } } else if let Some(nested) = node.downcast_ref::() { self.shadow_stack @@ -253,6 +261,72 @@ impl TreeNodeVisitor<'_> for CollectUsedVisitor<'_> { } } +/// Rewrites a lambda body so captured columns occupy `0..n_captures` and +/// this lambda's own parameters occupy `n_captures..`. Nested lambdas are +/// tracked on `shadow_stack` so an inner parameter that reuses a name is +/// not remapped as if it were this lambda's param. +struct ProjectLambdaBody<'a> { + column_index_map: &'a HashMap, + param_name_to_projected: &'a HashMap, + shadow_stack: Vec>, +} + +impl TreeNodeRewriter for ProjectLambdaBody<'_> { + type Node = Arc; + + fn f_down(&mut self, node: Self::Node) -> Result> { + if let Some(nested) = node.downcast_ref::() { + self.shadow_stack + .push(nested.params.iter().cloned().collect()); + return Ok(Transformed::no(node)); + } + + if let Some(column) = node.downcast_ref::() { + let original = column.index(); + let projected = *self.column_index_map.get(&original).unwrap(); + if projected != original { + return Ok(Transformed::yes(Arc::new(Column::new( + column.name(), + projected, + )))); + } + return Ok(Transformed::no(node)); + } + + if let Some(lambda_variable) = node.downcast_ref::() { + let name = lambda_variable.name(); + let shadowed = self.shadow_stack.iter().any(|frame| frame.contains(name)); + if shadowed { + // Nested lambda's own param: the inner LambdaExpr remaps it. + return Ok(Transformed::no(node)); + } + let projected = + if let Some(&projected) = self.param_name_to_projected.get(name) { + projected + } else { + // Captured outer lambda variable: still lives in the same + // index space as the outer-batch columns. + *self.column_index_map.get(&lambda_variable.index()).unwrap() + }; + if projected != lambda_variable.index() { + return Ok(Transformed::yes(Arc::new(LambdaVariable::new( + projected, + Arc::clone(lambda_variable.field()), + )))); + } + } + + Ok(Transformed::no(node)) + } + + fn f_up(&mut self, node: Self::Node) -> Result> { + if node.downcast_ref::().is_some() { + self.shadow_stack.pop(); + } + Ok(Transformed::no(node)) + } +} + impl std::fmt::Display for LambdaExpr { fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { write!(f, "({}) -> {}", self.params.join(", "), self.body) @@ -388,7 +462,8 @@ mod tests { let lambda = LambdaExpr::try_new(vec!["k".to_string(), "v".to_string()], body).unwrap(); - assert_eq!(lambda.projection(), &[1]); + // Own params are not part of the outer-batch capture projection. + assert!(lambda.projection().is_empty()); assert_eq!(lambda.used_param_indices(), &[1]); } @@ -438,7 +513,7 @@ mod tests { let lambda = LambdaExpr::try_new(vec!["k".to_string(), "v".to_string()], body).unwrap(); - assert_eq!(lambda.projection(), &[0, 1]); + assert!(lambda.projection().is_empty()); assert_eq!(lambda.used_param_indices(), &[0, 1]); } @@ -495,4 +570,29 @@ mod tests { shadowed inside the nested lambda" ); } + + /// A planner-assigned combined-schema index that would collide with a + /// real input column after a later projection rewrite must not appear + /// in the capture projection. See #24372. + #[test] + fn test_own_param_index_not_captured_from_outer_batch() { + let x_field = Arc::new(Field::new("x", DataType::Int64, true)); + // Planned against a 1-column input, so param `x` sat at index 1. + // After collapse the input is wider, and index 1 is a real column. + let body = Arc::new(LambdaVariable::new(1, Arc::clone(&x_field))); + + let lambda = LambdaExpr::try_new(vec!["x".to_string()], body).unwrap(); + + assert!( + lambda.projection().is_empty(), + "own param must not be treated as an outer-batch capture, got {:?}", + lambda.projection() + ); + assert_eq!(lambda.used_param_indices(), &[0]); + assert_eq!( + lambda.projected_body().as_ref().to_string(), + "x@0", + "own param should be rebased to slot 0 of the captures++params batch" + ); + } } diff --git a/datafusion/sqllogictest/test_files/array/array_filter.slt b/datafusion/sqllogictest/test_files/array/array_filter.slt index b6d73fbe7d09d..b558ae4628cac 100644 --- a/datafusion/sqllogictest/test_files/array/array_filter.slt +++ b/datafusion/sqllogictest/test_files/array/array_filter.slt @@ -224,6 +224,30 @@ SELECT array_filter(NULL, x -> x > 2); ---- NULL +# Regression for https://github.com/apache/datafusion/issues/24372 +# NestedLoopJoin filter pushdown remaps the lambda across a schema +# boundary; the parameter must stay bound to the filter element. +statement ok +set datafusion.sql_parser.dialect = spark; + +query II +SELECT l.id, r.k +FROM ( + VALUES + (1, [1], 10, 20), + (2, arrow_cast([NULL], 'List(Int64)'), 30, 40) +) AS l(id, arr, padding_1, padding_2) +JOIN (VALUES (0), (1), (2)) AS r(k) +ON (cardinality(array_filter(l.arr, x -> x IS NOT NULL)) > 0) + = (l.id > r.k) +ORDER BY l.id, r.k; +---- +1 0 +2 2 + +statement ok +set datafusion.sql_parser.dialect = databricks; + statement ok drop table t; diff --git a/datafusion/sqllogictest/test_files/array/array_transform.slt b/datafusion/sqllogictest/test_files/array/array_transform.slt index 5439d7441155b..521d758488ce4 100644 --- a/datafusion/sqllogictest/test_files/array/array_transform.slt +++ b/datafusion/sqllogictest/test_files/array/array_transform.slt @@ -429,6 +429,27 @@ 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); +# Regression for https://github.com/apache/datafusion/issues/24372 +# Collapsing the inner projection must not treat the lambda parameter +# index as an outer-batch column after the input schema grows. +query ? +SELECT array_transform(arr, x -> x) +FROM ( + SELECT arr + FROM (VALUES ([1, 2], 7)) AS t(arr, padding) +) AS q; +---- +[1, 2] + +query ? +SELECT array_transform(arr, x -> x + pad) +FROM ( + SELECT arr, padding AS pad + FROM (VALUES ([1, 2], 7)) AS t(arr, padding) +) AS q; +---- +[8, 9] + statement ok drop table t;