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

[Adam-Alani]

Which issue does this PR close?

• Closes #.

Rationale for this change

Extracted from #22689 per maintainer request: the original UDF and the underlying lambda fix deserve separate review threads. This PR carries only the lambda fix.

LambdaExpr previously compressed the column-index projection by enumerating every referenced Column/LambdaVariable index and packing them into a dense range. That collapse is correct for outer captures (and is a no-op for single-parameter lambdas, which is why array_transform was never affected), but it also moves lambda parameters around. A two-parameter lambda like (k, v) -> v (with k unused) would have its LambdaVariable for v re-projected from index 1 to index 0 — so at runtime the body reads the slot the higher-order function had filled with k and silently returns the wrong column.

This is a latent bug today — no in-tree higher-order function exercises it — but it blocks #22689 (transform_values, which uses (k, v) -> body lambdas) and any future HOF that takes more than one parameter.

What changes are included in this PR?

Per @gstvg's suggested non-breaking approach in the discussion thread:

• LambdaExpr now tracks used_params: HashSet<String> — the subset of its own declared parameters that the body actually references. It is computed during a single walk of the body in LambdaExpr::new, with a shadow stack that ignores LambdaVariables 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 the used parameter names plus a non-breaking LambdaArgument::new_with_used_params(...) constructor. The existing LambdaArgument::new(...) calls it with None, which preserves the old "push every declared parameter" behavior — so external callers that build LambdaArgument directly keep working unchanged.
• 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 LambdaArgument::new_with_used_params(...) and forwards lambda.used_params().clone(), so all in-tree higher-order UDFs pick up the fix automatically with no callsite change.

Compared to the previous revision of this PR (which added an outer_columns_count: usize parameter to LambdaExpr::try_new and expressions::lambda(...)), this revision:

• Has no breaking API change. LambdaExpr::try_new, expressions::lambda(...) and LambdaArgument::new keep their existing signatures. cargo-semver-checks should be clean now.
• Should be straightforwardly backportable to the 54 release branch.
• Also has the nice side effect of skipping evaluation of declared-but-unused parameters entirely (the higher-order function never invokes their closures), which avoids materializing arrays the lambda body never reads.

Are these changes tested?

Yes, two new unit tests in datafusion/physical-expr/src/expressions/lambda.rs:

• test_used_params_collects_only_referenced_param — a (k, v) -> v lambda reports only {\"v\"} as used.
• test_used_params_handles_shadowing_inside_nested_lambda — for (k, v) -> col + (k, v2) -> k + v2 + v, the outer lambda's used_params is {\"v\"} only; the inner k does not flag the outer k as used.

The existing test_lambda_evaluate, test_lambda_duplicate_name, and test_higher_order_function_* tests continue to pass. cargo test -p datafusion-expr higher_order (11 tests) and cargo test -p datafusion-physical-expr lambda (7 tests) both pass.

Are there any user-facing changes?

No breaking changes. LambdaArgument gains a new non-breaking constructor and a new optional field; the rest is purely internal correctness.

[rluvaton]

@gstvg can you please take a look?

[gstvg]

Thanks @Adam-Alani for unconvering and fixing this. This looks great to me, but I wonder if we want to fix this without breaking changes so that it can be trivially backported to the 54 branch? If you and @rluvaton agree on this, is possible to:

• Remove the body projection in LambdaExpr::new, keeping only the used columns indices collection
• In HigherOrderFunctionExpr::evaluate, instead of projecting the batch, derive a new batch with uncaptured columns swapped with a cheap NullArray
• In LambdaArgument::evaluate, inline merge_captures_with_variables, evaluate the parameters before spreading captures, get the len of the first evaluated parameter (0-param lambdas should return an error), then, if all columns in captures are DataType::Null, directly create a new batch of null arrays with the len of the first evaluated parameter, instead of calling spread_captures

That's how it was implemented in the first version of the lambda PR

[Adam-Alani]

Thanks @Adam-Alani for unconvering and fixing this. This looks great to me, but I wonder if we want to fix this without breaking changes so that it can be trivially backported to the 54 branch? If you and @rluvaton agree on this, is possible to:

• Remove the body projection in LambdaExpr::new, keeping only the used columns indices collection
• In HigherOrderFunctionExpr::evaluate, instead of projecting the batch, derive a new batch with uncaptured columns swapped with a cheap NullArray
• In LambdaArgument::evaluate, inline merge_captures_with_variables, evaluate the parameters before spreading captures, get the len of the first evaluated parameter (0-param lambdas should return an error), then, if all columns in captures are DataType::Null, directly create a new batch of null arrays with the len of the first evaluated parameter, instead of calling spread_captures

That's how it was implemented in the first version of the lambda PR

@rluvaton let me know if this plan sounds good to you, happy to do that.

[gstvg]

I imagined another approach:

• Add used_params as HashSet<String> or Option<HashSet<String>> to LambdaArgument
• Add LambdaArgument::new_with_used_params, the existing new invokes this new method with None or the hash set of all params names
• While traversing the lambda body to collect used indices, also collect used lambda variables names of the given lambda (ignoring vars of others lambda), and provide them to LambdaArgument::new_with_used_params. It must take into account shadowing: (k, v) -> func(column, (k, v2) -> k + v2 + v), the innermost lambda param k should not flag the outermost one as used. Only the outermost v is used
• In LambdaArgument::evaluate, just evaluates and push to the batch the parameters that are actually used, so unused ones don't shift the others.

Because this skips the evaluation of declared but not used parameters, I personally prefer this one, WDYT?

[Adam-Alani]

@gstvg applied the changes.

[gstvg]

Thanks @Adam-Alani, LGTM
cc @rluvaton

[Adam-Alani]

@rluvaton small bump on this, thanks!

[LiaCastaneda]

I don't want to block the fix, but wondering if there is a non user facing way to fix the issue @Adam-Alani @gstvg lmk what you think and if it makes sense!

[LiaCastaneda]

If we pass None in new() and assume all parameters are used, would it be better to compute the indices directly from params: Vec<FieldRef>? That way this field doesn't need to be an Option it always holds the used parameter indices, which happen to be all of them when called via new()

[LiaCastaneda]

Should we just require used_params to be set in this function? Otherwise the behavior is the same as new, so I'm not sure why we'd need both for that case

[LiaCastaneda]

Providing an api where we have to specify the indices of the parameters that will be used in the body feels a bit unergonomic. Have you considered somehow extracting this from the body and handling this internally so it's invisible to the caller?

[LiaCastaneda]

Also, the body is dynamic, which technically means for higher order functions indices positions can differ for each query no? for example if you have a Higher Order function with parameters (x,y,x) for a given query you can use x,y or y or all of them. This essentially means external callers of LambdaArgument::new_with_used_params would have to walk the body themselves to figure out which params are referenced.

edit: is it actually possible to build a LambdaArgument like as an api? they are technically build in DF before calling the invoke_with_args api 🤔 so I guess new_with_used_params would only be called in higher_order_function.rs in evaluate inside DataFusion. If that's the case, I would consider resolving this inside LambdaExpr instead, keeping the fix self-contained there

[gstvg]

Yes, is dynamic and require walking the body, but at least for this PR, it's folded with the walk that collect used indices.

is it actually possible to build a LambdaArgument like as an api?

Technically yes but not because I imagined a usecase for it outside datafusion, but only to be able to use it in evaluate which is within another crate

Some other ideas besides #22853 (comment):

1. Providing the LambdaExpr itself (instead of it's body) as the Arc<dyn PhyiscalExpr> parameter of LambdaArgument::new, so evaluate can get the used params by downcasting the expr to LambdaExpr (or returning an error if it fails)
2. computing the lambda used params in LambdaArgument::evaluate, at the small cost of a tree traversal per evaluation instead of only during planning

[LiaCastaneda]

Since used_param_names is already computed here during construction, could the fix live entirely in projected_body instead? Rather than threading used_params through to LambdaArgument, then LambdaArgument stays unchanged and external callers don't need to think about used params at all
Left a rough idea here LiaCastaneda@993ac37

[gstvg]

The initial version of this PR b9f816f did that, but introduced a breaking change by adding a new argument outer_columns_count in LambdaExpr::try_new. I then proposed non-breaking alternatives #22853 (comment) (which is self-contained but uses null arrays to avoid uncaptured columns copies) and #22853 (comment) (the current version, which has the small advantage of skipping the evaluation of declared but unused parameters, for example (x, y) -> x+1). Another self-contained fix that stills uses projection, and maybe also skip unused params would be great. Using the first suggestion is fine by me as well. WDYT?

[LiaCastaneda]

ok thanks for explaining, so the correctness fix can be self contained in datafusion-physical-expr with no new public API, but the optimization of skipping unused param evaluation requires crossing into datafusion-expr, which forces new_with_used_params to be public (even if it's only meant to be used internally).

My approach is non-breaking but does not skip unused param evaluation. I'm not sure this is achievable while keeping the fix self contained in datafusion-physical-expr, since the evaluate lives in datafusion-expr.

I think the optimization (therefore this approach) is worth keeping, the cost of evaluating the unused param will grow with the number of map entries, so for large inputs skipping an unused param closure might be meaningful