Skip to content

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

Open
LiaCastaneda wants to merge 3 commits into
apache:mainfrom
LiaCastaneda:lambda-multi-param-fix
Open

fix(lambda): only push referenced params into the merged batch#24162
LiaCastaneda wants to merge 3 commits into
apache:mainfrom
LiaCastaneda:lambda-multi-param-fix

Conversation

@LiaCastaneda

@LiaCastaneda LiaCastaneda commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

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) -> v k 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) -> v v 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 github-actions Bot added logical-expr Logical plan and expressions physical-expr Changes to the physical-expr crates labels Aug 7, 2026
@github-actions

github-actions Bot commented Aug 7, 2026

Copy link
Copy Markdown

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 [  30.059s] (current)
     Parsing datafusion-expr v54.1.0 (current)
      Parsed [   0.079s] (current)
    Building datafusion-expr v54.1.0 (baseline)
       Built [  31.153s] (baseline)
     Parsing datafusion-expr v54.1.0 (baseline)
      Parsed [   0.078s] (baseline)
    Checking datafusion-expr v54.1.0 -> v54.1.0 (no change; assume patch)
     Checked [   1.431s] 223 checks: 222 pass, 1 fail, 0 warn, 30 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.49.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/1251ae8c900020c872bd8dcd238407e32cdfd763/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:285

     Summary semver requires new major version: 1 major and 0 minor checks failed
    Finished [  64.022s] datafusion-expr
    Building datafusion-physical-expr v54.1.0 (current)
       Built [  32.113s] (current)
     Parsing datafusion-physical-expr v54.1.0 (current)
      Parsed [   0.050s] (current)
    Building datafusion-physical-expr v54.1.0 (baseline)
       Built [  31.956s] (baseline)
     Parsing datafusion-physical-expr v54.1.0 (baseline)
      Parsed [   0.050s] (baseline)
    Checking datafusion-physical-expr v54.1.0 -> v54.1.0 (no change; assume patch)
     Checked [   0.337s] 223 checks: 223 pass, 30 skip
     Summary no semver update required
    Finished [  65.379s] datafusion-physical-expr

@codecov-commenter

codecov-commenter commented Aug 7, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 96.15385% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.04%. Comparing base (0646a31) to head (f230015).
⚠️ Report is 3 commits behind head on main.

Files with missing lines Patch % Lines
datafusion/expr/src/higher_order_function.rs 81.81% 2 Missing and 4 partials ⚠️
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.
📢 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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

auto detected api change Auto detected API change logical-expr Logical plan and expressions physical-expr Changes to the physical-expr crates

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants