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

[LiaCastaneda]

Which issue does this PR close?

basically this PR #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) -> vk 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) -> vv 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<String>, however LambdaArgument::new is very unlikely to be called outside datafusion, see this comment

[github-actions]

Thank you for opening this pull request!

Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch).

Details

 Cloning apache/main
Building datafusion-expr v54.1.0 (current)
Built [ 28.299s] (current)
Parsing datafusion-expr v54.1.0 (current)
Parsed [ 0.074s] (current)
Building datafusion-expr v54.1.0 (baseline)
Built [ 27.008s] (baseline)
Parsing datafusion-expr v54.1.0 (baseline)
Parsed [ 0.073s] (baseline)
Checking datafusion-expr v54.1.0 -> v54.1.0 (no change; assume patch)
Checked [ 1.432s] 223 checks: 222 pass, 1 fail, 0 warn, 31 skip
--- failure method_parameter_count_changed: pub method parameter count changed ---
Description:
A publicly-visible method now takes a different number of parameters, not counting the receiver (self) parameter.
ref: https://doc.rust-lang.org/cargo/reference/semver.html#fn-change-arity
impl: https://github.com/obi1kenobi/cargo-semver-checks/tree/v0.50.0/src/lints/method_parameter_count_changed.ron
Failed in:
datafusion_expr::LambdaArgument::new takes 3 parameters in /home/runner/work/datafusion/datafusion/target/semver-checks/git-apache_main/f88709796761ac0d8ef24e2c482c7d5a1420f47a/datafusion/expr/src/higher_order_function.rs:260, but now takes 4 parameters in /home/runner/work/datafusion/datafusion/datafusion/expr/src/higher_order_function.rs:286
Summary semver requires new major version: 1 major and 0 minor checks failed
Finished [ 58.241s] datafusion-expr
Building datafusion-physical-expr v54.1.0 (current)
Built [ 28.744s] (current)
Parsing datafusion-physical-expr v54.1.0 (current)
Parsed [ 0.047s] (current)
Building datafusion-physical-expr v54.1.0 (baseline)
Built [ 28.758s] (baseline)
Parsing datafusion-physical-expr v54.1.0 (baseline)
Parsed [ 0.048s] (baseline)
Checking datafusion-physical-expr v54.1.0 -> v54.1.0 (no change; assume patch)
Checked [ 0.332s] 223 checks: 223 pass, 31 skip
Summary no semver update required
Finished [ 59.445s] datafusion-physical-expr

[codecov-commenter]

Codecov Report

❌ Patch coverage is 88.58268% with 29 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.03%. Comparing base (0646a31) to head (67b0300).
⚠️ Report is 61 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/expr/src/higher_order_function.rs	70.52%	21 Missing and 7 partials ⚠️
...afusion/physical-expr/src/higher_order_function.rs	97.72%	1 Missing ⚠️

Additional details and impacted files

@@ Coverage Diff @@## main #24162 +/- ##
==========================================
+ Coverage 81.02% 81.03% +0.01% 
==========================================
Files 1105 1107 +2 Lines 380669 385292 +4623 Branches 380669 385292 +4623 ==========================================
+ Hits 308446 312233 +3787 - Misses 53994 54641 +647 - Partials 18229 18418 +189 

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[timsaucer]

This is a nice PR, and thank you for fixing the bug!

In addition to the things I noted, I also directed an AI review and it came up with below. #1, #2, and #5 look important to me, but I haven't dug deeper yet.

Issues

1. No test covers the actual fix (blocking)

All five new tests assert LambdaExpr::used_params() — the metadata. Nothing exercises LambdaArgument::new → merge_captures_with_variables, which is where the bug lived and where the layout contract is enforced. The PR body says "added two new tests… for the unused-parameter case"; those tests would still pass if merge_captures_with_variables ignored used_param_indices entirely.

This is testable without a real multi-param HOF: in datafusion/expr/src/higher_order_function.rs build a LambdaArgument with params = [k, v], used_params = {"v"}, body = LambdaVariable::new(0, v_field), and call evaluate with two distinguishable closures. Assert you get v's values, not k's. Add a captures variant too, since the captures ++ params offset is the other half of the contract.

2. Name-based matching is redundant and now a silent-corruption hazard

LambdaArgument::new takes &HashSet<String> and re-derives positions by name — but the caller already has positional alignment: params is built at physical-expr/higher_order_function.rs:336 by zip(lambda.params(), lambda_params), so index i of paramsis index i of lambda.params(). The round-trip name → index buys nothing.

The cost is a new failure mode. LambdaVariable can be built directly (proto decode, lambda_variable.rs:189, tests) with a field name that doesn't match the declared param. Before this PR a mismatch was harmless; now it silently drops the param from used_params, the column never gets pushed, and the body reads the wrong slot — no error, wrong answers.

Suggest: cache used_param_indices: Vec<usize> on LambdaExpr, expose used_param_indices() -> &[usize], and have LambdaArgument::new take &[usize]. That removes the name coupling, removes per-batch string hashing (LambdaArgument::new runs per evaluate, i.e. per batch), makes LambdaExpr::clone cheaper, and avoids putting datafusion_common::HashSet (a hashbrown re-export) in the public API.

If you keep names, at least assert in LambdaExpr::try_new that every name in used_params is in params — currently an unmatched name is a silent no-op.

3. Missing test for the f_up pop

shadow_stack.pop() is never verified. Every shadowing test has a single nested lambda, so a missing pop would pass all five. Add: outer (k, v), body = f(g(arr, (k) -> k), k) — the second k sits after the nested lambda at the same level, so a leaked frame would wrongly mark outer k unused.

4. variables.first() row-count derivation

let row_count = match variables.first(){Some(first) => first()?.len(),None => 0,};

Two things. The comment says evaluating a variable is "essentially free," but that's an assumption about HOF impls, not a guarantee — a closure that builds an index/range array is not free, and here it's built purely to be discarded. Passing the row count down from LambdaArgument::evaluate would be exact.

The None => 0 arm is unreachable in practice and untested; consider internal_err! instead of silently producing a 0-row batch, which would be very hard to debug if it ever fires.

5. Undocumented behavior change for HOF authors

Unused params' closures are no longer invoked at all. That's a genuine perf win (skips materializing e.g. an unused index array per batch), but it's an observable contract change for HigherOrderUDFImpl implementors who assumed every closure in args gets called. Worth a line in the evaluate doc and in the PR description.

6. Missing upgrade-guide entry (project convention)

The PR carries the auto detected api change label for the LambdaArgument::new signature break. docs/source/library-user-guide/upgrading/55.0.0.md is the active page and gets entries for exactly this. Add one, even if short — the "unlikely to be called externally" argument justifies making the break, not skipping the note. (Adopting #2 changes the new arg's type, so write the entry after settling that.)

---

Smaller things

• Doc duplication. The same three-paragraph explanation appears four times: the used_param_indices field, LambdaArgument::new, the used_params field, and LambdaExpr::used_params. Keep one canonical version (the CollectUsedVisitor doc is the best-written) and make the rest one-liners pointing at it.
• lambda.rs:408 — [Self::used_params] inside a #[cfg(test)] doc comment; Self doesn't resolve there. Use plain backticks.
• lambda.rs:357 — use super::LambdaExpr; after use std::sync::Arc;, separated from the other imports. Fold it into the block above.
• Tests reach for crate::expressions::BinaryExpr and datafusion_expr::Operator by full path inline; import them alongside Column/LambdaVariable for consistency with the existing test style.
• The .copied() + (new_idx, original) rename in column_index_map is a real readability improvement — the old (projected, original) binding on enumerate() was actively misleading. Good unrelated cleanup.

[timsaucer]

Is this a separate issue that's caught and included here or was this introduced by the above changes?

[LiaCastaneda]

sort of introduced by our changes. Before, we pushed all params (used or not) into columns so the variable was never empty. Now we just push the used params only, so for lambdas that use 0 params (like (k,v) -> constant) columns is empty, so we need to derive the row count manually (otherwise RecordBatch::try_new(schema, columns) wont do it, like before)

[timsaucer]

TBH the term "nested-lambda shadow tracking" doesn't have obvious meaning to me. Is there an easy way to make the meaning more clear, or somewhere else in the code I should have looked to understand what it means?

[LiaCastaneda]

when a lambda has nested lambdas like (k, v) -> func(col, (k, v2) -> k + v2 + v), the innermost k "shadows"/"overrides" the outermost k bc even though they share a name, the innermost k is its own separate parameter, not a reference to the outer one. So when collecting which parameters a lambda's body references n a nested lambda, we need to take this possible name conflict into account, otherwise we'd wrongly count the outer k as used just because the name appears in the body of the nested lambda.

Its true its a bit unclear by just reading it in the param description, I moved the explanation to CollectUsedVisitor that actually handles the logic for "shadowed" params

[timsaucer]

This is a case where the current lambda functions are all working because they're single variable, right? I am wondering if this is necessary to get into #22393 or if it's okay going in the next release.

[LiaCastaneda]

thanks for the review @timsaucer! I will take a look shortly.

This is a case where the current lambda functions are all working because they're single variable, right?

Yep exactly, this came up while @Adam-Alani (the original author of the fix) was working on a PR to add transform_values

[LiaCastaneda]

I am wondering if this is necessary to get into #22393

it would be nice if its on the new release

[LiaCastaneda]

The cost is a new failure mode. LambdaVariable can be built directly (proto decode, lambda_variable.rs:189, tests) with a field name that doesn't match the declared param. Before this PR a mismatch was harmless; now it silently drops the param from used_params, the column never gets pushed, and the body reads the wrong slot — no error, wrong answers.

this is unlikely but not impossible, if it does differ it would be a coding bug, however its true its probably cleanest to just handle the indices instead of names

[LiaCastaneda]

@timsaucer I addressed most of the review points :)

[timsaucer]

Looks great, thanks again! One minor comment that can be fixed with documentation. Also my agent suggests one additional test:

LLM generated evaluation

The layout contract is that in the un-projected body's index space, every capture index sorts before every own-param index, which sorts before every nested-param index. That's what makes column_index_map's dense rank for a used param equal n_captures + rank_among_used_params, which is what makes captures ++ used_params in LambdaArgument::new line up. It holds because the planner appends each lambda's params to the enclosing scope schema — but nothing states it, nothing checks it, and all seven tests hand-write indices chosen to be consistent with it. If the planner's numbering ever changed, every test here would still pass and results would be silently wrong.

At minimum, a comment on used_param_indices stating the invariant. Better: an end-to-end test, which is cheaper than it looks — MockHigherOrderUDF in datafusion/physical-expr/src/higher_order_function.rs already calls lambda.evaluate(...) from invoke_with_args. Widening it (or adding a two-param sibling) to return two lambda_parameters and two closures gives you planner → LambdaExpr → used_param_indices → LambdaArgument → merge_captures_with_variables in one test, with the indices produced by the real planner instead of by hand. That's the test that would actually have caught the original bug.

[timsaucer]

There is a potential for a panic here if used_param_indices ever contains an index that is larger than the length of params. In practice, not a problem but it looks like this is a pub function so we at least need to document the invariant. Maybe even a debug_assert.

[LiaCastaneda]

make sense, done.