fix: rebase lambda params after projection schema changes

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::<Vec<_>>();
 
         projection.sort();
@@ -102,39 +110,33 @@ impl LambdaExpr {
             .map(|(new_idx, original)| (original, new_idx))
             .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;
-
-        let used_param_indices = params
+        let used_param_indices: Vec<usize> = 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<String, usize> = 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: &param_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::<Column>() {
             self.used_indices.insert(col.index());
         } else if let Some(var) = node.downcast_ref::<LambdaVariable>() {
-            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::<LambdaExpr>() {
             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<usize, usize>,
+    param_name_to_projected: &'a HashMap<String, usize>,
+    shadow_stack: Vec<HashSet<String>>,
+}
+
+impl TreeNodeRewriter for ProjectLambdaBody<'_> {
+    type Node = Arc<dyn PhysicalExpr>;
+
+    fn f_down(&mut self, node: Self::Node) -> Result<Transformed<Self::Node>> {
+        if let Some(nested) = node.downcast_ref::<LambdaExpr>() {
+            self.shadow_stack
+                .push(nested.params.iter().cloned().collect());
+            return Ok(Transformed::no(node));
+        }
+
+        if let Some(column) = node.downcast_ref::<Column>() {
+            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::<LambdaVariable>() {
+            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<Transformed<Self::Node>> {
+        if node.downcast_ref::<LambdaExpr>().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"
+        );
+    }
 }

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;
 

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;