fix(lambda): only push referenced params into the merged batch - #24162
fix(lambda): only push referenced params into the merged batch#24162LiaCastaneda wants to merge 3 commits into
Conversation
|
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 |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #24162 +/- ##
==========================================
+ Coverage 81.02% 81.04% +0.01%
==========================================
Files 1105 1106 +1
Lines 380669 381118 +449
Branches 380669 381118 +449
==========================================
+ Hits 308446 308883 +437
+ Misses 53994 53982 -12
- Partials 18229 18253 +24 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
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 params is 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_indicesfield,LambdaArgument::new, theused_paramsfield, andLambdaExpr::used_params. Keep one canonical version (theCollectUsedVisitordoc 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;Selfdoesn't resolve there. Use plain backticks.lambda.rs:357—use super::LambdaExpr;afteruse std::sync::Arc;, separated from the other imports. Fold it into the block above.- Tests reach for
crate::expressions::BinaryExpranddatafusion_expr::Operatorby full path inline; import them alongsideColumn/LambdaVariablefor consistency with the existing test style. - The
.copied()+(new_idx, original)rename incolumn_index_mapis a real readability improvement — the old(projected, original)binding onenumerate()was actively misleading. Good unrelated cleanup.
| /// Build a [`LambdaArgument`] for a lambda whose body references the | ||
| /// subset of `params` named in `used_params`. | ||
| /// | ||
| /// [`Self::evaluate`] only materialises the closures whose parameter name | ||
| /// appears in `used_params`, preserving the original declaration order of | ||
| /// `params`. Unused declared parameters therefore leave no slot in the | ||
| /// merged batch, so the body's compressed column indices line up directly | ||
| /// with the columns the evaluator built. | ||
| /// | ||
| /// Callers with a `LambdaExpr` in hand should pass `lambda.used_params()`; | ||
| /// that method already computes the exact set required here (with | ||
| /// nested-lambda shadow tracking). |
There was a problem hiding this comment.
Here's a problem I have with generated code documentation: LLMs of write this kind of documentation based on the diff of the code they're writing and not based on end user viewpoint. From a consumer of this function, the purpose of this is to create a new LambdaArgument. Instead if we read the docstring it's heavily focused on the purpose of used_params. Now this does have documentation where the prior new function didn't, so that's generally an improvement but the way this is written feels not obvious to me as a user of datafusion who isn't narrowly looking at the problem that this PR addresses.
I'm assuming this was generated documentation, so please correct me if I'm wrong. I've had this problem with my own PRs as well and it's one of the things I've had to add extra instructions in my agents' context just to avoid these kinds of documentation that are PR-facing rather than user-facing.
There was a problem hiding this comment.
yeah, I read through it and didn't find it irrelevant but its true it can be simplified since it only talks about the params param
There was a problem hiding this comment.
I don't even think we need doc for this function, the function itself is very self-explanatory (and also it will probably never be used outside datafusion)
| 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)), | ||
| )?); | ||
| } | ||
|
|
There was a problem hiding this comment.
Is this a separate issue that's caught and included here or was this introduced by the above changes?
| projected_body: Arc<dyn PhysicalExpr>, | ||
| projection: Vec<usize>, | ||
| /// Subset of `params` (by name) that the body actually references, | ||
| /// computed with nested-lambda shadow tracking. Empty when no parameter |
There was a problem hiding this comment.
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?
|
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. |
|
thanks for the review @timsaucer! I will take a look shortly.
Yep exactly, this came up while @Adam-Alani (the original author of the fix) was working on a PR to add |
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 thatLambdaExprmishandles lambdas with more than one parameter. The bug surfaced while working ontransform_values(#22689), which needs(k, v) -> exprtwo 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) -> vkalways goes into slot 0,valways into slot 1.LambdaExprseparately scans the body and renumbers whatever it finds referenced into a dense0..nrange, to avoid carrying around columns nothing uses (likevin 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) -> vvis 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, writesk's values into slot 0 andv's into slot 1. So the body ends up reading slot 0 expectingv— and getskinstead. So the results end up being incorrect.What changes are included in this PR?
LambdaExprnow computesused_params: which is the subset of its own declared parameters that are actually referenced in the body.LambdaArgument::newtakesused_paramsand 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::evaluateforwardslambda.used_params()toLambdaArgument::newAre 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::newwhich now requires a new argument:used_params: &HashSet<String>, however LambdaArgument::new is very unlikely to be called outside datafusion, see this comment