fix(lambda): only push referenced params into the merged batch

`LambdaExpr` compresses the body's column-index projection by enumerating
every referenced `Column`/`LambdaVariable` index and packing them into a
dense range. That compression is correct for outer captures, but it
silently broke multi-parameter lambdas: a body like `(k, v) -> v` (where
`k` is unused) would have its `LambdaVariable("v")` re-projected from
index 1 to index 0 and then, at runtime, read the slot the higher-order
function had filled with `k`.

Per maintainer feedback on #22853, fix it without a breaking
change to `LambdaExpr::try_new` / `expressions::lambda(...)`:

* `LambdaExpr` now tracks `used_params: HashSet<String>` — the subset of
  its own declared parameters that the body actually references. The set
  is computed during a single walk of the body in `LambdaExpr::new`,
  with a shadow stack that ignores `LambdaVariable`s bound by nested
  lambdas. For
  `(k, v) -> func(col, (k, v2) -> k + v2 + v)` the inner `k` shadows the
  outer `k`, so only `v` flows up as used by the outer lambda.

* `LambdaArgument` gets an `Option<HashSet<String>>` for used parameter
  names plus a non-breaking `new_with_used_params(...)` constructor.
  The existing `new(...)` calls it with `None`, which preserves the old
  "push every declared parameter" behavior.

* `LambdaArgument::evaluate` (through `merge_captures_with_variables`)
  only evaluates and pushes the closures whose parameter name appears
  in `used_params`, preserving the original declaration order. Unused
  declared parameters therefore leave no slot in the merged batch, so
  the body's compressed indices line up directly with the columns the
  evaluator actually built.

* `HigherOrderFunctionExpr::evaluate` calls `new_with_used_params` and
  forwards `lambda.used_params().clone()`, so all in-tree higher-order
  UDFs benefit automatically without any callsite change.

No public API breakage: `LambdaExpr::try_new`, `expressions::lambda(...)`
and `LambdaArgument::new` keep their existing signatures. Two new
tests cover the unused-parameter case and the nested-lambda shadowing
case; existing tests in `physical-expr` and `expr` continue to pass.

Author: Adam-Alani

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,15 @@ pub struct LambdaArgument {
     /// For example, for `array_transform([2], v -> -v)`,
     /// this will be `vec![Field::new("v", DataType::Int32, true)]`
     params: Vec<FieldRef>,
+    /// Indices into `params` of the parameters that are actually referenced
+    /// by `body` (taking nested-lambda shadowing into account).
+    ///
+    /// `None` means "no information, assume every declared parameter is used"
+    /// — that is the backwards-compatible behavior of [`Self::new`]. When set,
+    /// [`Self::evaluate`] skips evaluating and pushing the closures for the
+    /// parameters not listed here, so unused declared parameters do not shift
+    /// the columns the body's compressed indices expect.
+    used_param_indices: Option<Vec<usize>>,
     /// The body of the lambda
     ///
     /// For example, for `array_transform([2], v -> -v)`,
@@ -257,26 +266,64 @@ pub struct LambdaArgument {
 }
 
 impl LambdaArgument {
+    /// Build a [`LambdaArgument`] that treats every declared parameter as
+    /// used. This is the backwards-compatible behavior. Prefer
+    /// [`Self::new_with_used_params`] when the caller knows which subset of
+    /// the lambda's parameters the body actually references — otherwise the
+    /// merged batch will still contain columns for unused parameters.
     pub fn new(
         params: Vec<FieldRef>,
         body: Arc<dyn PhysicalExpr>,
         captures: Option<RecordBatch>,
     ) -> Self {
+        Self::new_with_used_params(params, body, captures, None)
+    }
+
+    /// Build a [`LambdaArgument`] knowing which subset of `params` (by name)
+    /// the lambda body actually references.
+    ///
+    /// When `used_params` is `Some(set)`, [`Self::evaluate`] only evaluates
+    /// and pushes the closures whose corresponding parameter name is in
+    /// `set`, in the original declaration order of `params`. Unused declared
+    /// parameters leave no slot in the merged batch, so the body's compressed
+    /// column indices line up directly. When `used_params` is `None`,
+    /// behavior is identical to [`Self::new`].
+    pub fn new_with_used_params(
+        params: Vec<FieldRef>,
+        body: Arc<dyn PhysicalExpr>,
+        captures: Option<RecordBatch>,
+        used_params: Option<HashSet<String>>,
+    ) -> Self {
+        let used_param_indices = used_params.map(|set| {
+            params
+                .iter()
+                .enumerate()
+                .filter(|(_, f)| set.contains(f.name()))
+                .map(|(i, _)| i)
+                .collect::<Vec<_>>()
+        });
+
+        let effective_params: Vec<FieldRef> = match &used_param_indices {
+            Some(indices) => indices.iter().map(|i| Arc::clone(&params[*i])).collect(),
+            None => params.clone(),
+        };
+
         let fields = match &captures {
             Some(batch) => batch
                 .schema_ref()
                 .fields()
                 .iter()
                 .cloned()
-                .chain(params.clone())
+                .chain(effective_params.iter().cloned())
                 .collect(),
-            None => params.clone(),
+            None => effective_params,
         };
 
         let schema = Arc::new(Schema::new(fields));
 
         Self {
             params,
+            used_param_indices,
             body,
             schema,
             captures,
@@ -344,6 +391,7 @@ impl LambdaArgument {
             spread_captures.as_ref(),
             Arc::clone(&self.schema),
             &self.params,
+            self.used_param_indices.as_deref(),
             args,
         )?;
 
@@ -355,6 +403,7 @@ fn merge_captures_with_variables(
     captures: Option<&RecordBatch>,
     schema: SchemaRef,
     params: &[FieldRef],
+    used_param_indices: Option<&[usize]>,
     variables: &[&dyn Fn() -> Result<ArrayRef>],
 ) -> Result<RecordBatch> {
     if variables.len() < params.len() {
@@ -365,23 +414,56 @@ fn merge_captures_with_variables(
         );
     }
 
+    let push_param_arrays = |columns: &mut Vec<ArrayRef>| -> Result<()> {
+        match used_param_indices {
+            Some(indices) => {
+                for &i in indices {
+                    columns.push(variables[i]()?);
+                }
+            }
+            None => {
+                for arg in &variables[..params.len()] {
+                    columns.push(arg()?);
+                }
+            }
+        }
+        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
+                    .map(<[usize]>::len)
+                    .unwrap_or(params.len()),
+            );
+            push_param_arrays(&mut columns)?;
             columns
         }
-        None => variables
-            .iter()
-            .take(params.len())
-            .map(|arg| arg())
-            .collect::<Result<_>>()?,
     };
 
+    if columns.is_empty() {
+        // Constant lambda body with no captures and no used parameters. We
+        // still need a row count for the merged batch, so evaluate one
+        // variable just to derive it. This is essentially free in the common
+        // case (the variables already exist as closures over arrays the
+        // caller computed up front).
+        let row_count = match variables.first() {
+            Some(first) => first()?.len(),
+            None => 0,
+        };
+        return Ok(RecordBatch::try_new_with_options(
+            schema,
+            vec![],
+            &RecordBatchOptions::new().with_row_count(Some(row_count)),
+        )?);
+    }
+
     Ok(RecordBatch::try_new(schema, columns)?)
 }