Skip to content

Add lambdas for higher-order scalar functions - #10502

Open
mhk197 wants to merge 1 commit into
developfrom
mk/lambdas-v3
Open

mhk197 wants to merge 1 commit into
developfrom
mk/lambdas-v3

Conversation

@mhk197

@mhk197 mhk197 commented Oct 14, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Adds lambdas, the function arguments of higher-order scalar functions such as list_transform (#10503). A lambda is a list of parameter names and a body that reads them through variables. A function that takes a lambda keeps it in its options, binds it once when its node is bound, and runs the bound lambda over values it makes from its inputs.

Variables

  • var("x") builds Expression::Variable("x"): a leaf, like $, that reads the parameter x of the enclosing lambda.
  • It displays as @x, so @x (a parameter) and $.x (a field of the input) never look alike.
  • A variable exists only in a written body. Binding the lambda turns @x into $.x, a field of the struct of the lambda's inputs (see below), so BoundExpression is unchanged.
  • Anywhere else it fails to bind, type, apply and serialize with "variable @x is not a parameter of an enclosing lambda".

Lambdas

let lambda = Lambda::new(["x"], checked_add(var("x"), lit(1_i32)))?;
// (x) -> (@x + 1i32)
let bound = lambda.bind(&[DType::Primitive(PType::I32, NonNullable)])?;
// (x) -> ($.x + 1i32): an ordinary bound expression over the struct {x: i32}
  • Lambda is the lambda as written: params() and body(), an Expression.
  • Lambda::new rejects a repeated parameter, and a body that reads from outside the lambda: $ or a column, another lambda's parameter, or the scan's row_idx() and row_count(). Constants are fine, such as the literal set in x -> list_contains([2, 3], x).
  • bind(param_dtypes) turns the lambda into an ordinary expression over a struct of its inputs, one field per parameter: each @x becomes $.x, and the body is bound to that struct. It returns a BoundLambda.
  • return_dtype(param_dtypes) gives the body's dtype over the same struct, without binding it, the way Expression::return_dtype does for a whole expression.
  • BoundLambda is the lambda bound: params() with their dtypes, body(), an ordinary BoundExpression over the struct, dtype(), the body's, and is_bound_to(param_dtypes).
  • BoundLambda::evaluate(args, len) puts one array per parameter in that struct and applies the bound body to it. The arrays must have the dtypes the lambda is bound to. It binds nothing, and like any applied expression, the result is computed later.
  • is_infallible(), on both, says whether nothing in the body can fail, so a function knows whether the body may run on values whose results are never read, such as dictionary values no row uses.

Using a lambda in a function

A function's options are an enum of the lambda as written and as bound. Its children are its ordinary inputs.

enum MyFnOptions {
    Lambda(Lambda),
    Bound(BoundLambda),
}

impl ScalarFnVTable for MyFn {
    type Options = MyFnOptions;

    fn bind_options(&self, options: &MyFnOptions, args: &[DType]) -> VortexResult<Option<MyFnOptions>> {
        match options {
            MyFnOptions::Lambda(lambda) => Ok(Some(MyFnOptions::Bound(lambda.bind(/* param dtypes, from args */)?))),
            MyFnOptions::Bound(lambda) => {
                vortex_ensure!(lambda.is_bound_to(/* param dtypes, from args */), "bound to other dtypes");
                Ok(None)
            }
        }
    }

    fn execute(&self, options: &MyFnOptions, args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
        let MyFnOptions::Bound(lambda) = options else { vortex_bail!("not bound") };
        lambda.evaluate(vec![/* the parameters' arrays, made from args */], len)
    }
}
  • ScalarFnVTable::bind_options(options, arg_dtypes) is new, and does nothing by default. It runs whenever a bound node is built, and whenever a ScalarFnArray is built from its children, before return_dtype. That's where a function binds its lambda, once, so execute finds it bound.
  • A rebuilt node calls bind_options again, possibly with other dtypes, so a function checks that options already bound still fit, and errors if they don't. No rewrite today changes the dtype a lambda is bound to.
  • An Expression holds the written lambda, and a BoundExpression or an array holds the bound one.
  • The body reads only its parameters, so the function's children are everything it reads. Binding, traversal, replace, field pruning and pushdown see an ordinary node and never look inside the lambda.
  • Lambdas nest: a call that takes a lambda is an ordinary node inside another lambda's body. Its children are part of the outer body, so they can read the outer parameters, and its own lambda reads only its own. Binding the outer lambda binds the inner one.
(xs) -> vortex.list.transform(@xs, (y) -> (@y + 1i32))

Not supported yet

  • Reading anything outside the lambda, such as x -> x + $.a.
  • Serializing a lambda: a function that takes one is not serializable yet.

Testing

  • cargo +nightly-2026-09-10 fmt -p vortex-array -p vortex-layout
  • cargo clippy -p vortex-array -p vortex-layout --all-targets --all-features -- -D warnings
  • cargo nextest run -p vortex-array -p vortex-layout (the full suites), including a vortex-layout test that a lambda can't read row_idx()
  • cargo test --doc -p vortex-array -- lambda exprs::var

Follow-up

  • Reading values from outside a lambda, passed to the function as extra inputs and to the lambda as extra parameters.
  • Serialization: the written lambda in the function's options, with a wire id for variables inside bodies. The bound lambda is rebuilt when the expression is bound.
  • A lambda as a node of the written expression tree, folded into the function's options when it is bound. Binding and execution would not change.
  • Rebinding a bound lambda when a rewrite changes its parameters' dtypes, which captures will need: a captured value's nullability follows the root's.
  • Letting passes look inside lambda bodies, for example to read only the element fields a body uses.

@mhk197 mhk197 added the changelog/feature A new feature label Oct 14, 2026
@mhk197 mhk197 mentioned this pull request Oct 14, 2026
@codspeed

codspeed Bot commented Oct 14, 2026 •

Copy link
Copy Markdown

Merging this PR will degrade performance by 10.74%

3 warnings, 1 note

⚠️ Walltime results may be unreliable

Some benchmarks used the Walltime instrument outside CodSpeed Macro Runners, which can make the results inconsistent.

Learn more about Macro Runners

See affected benchmarks

⚠️ 16 benchmarks are dominated by system calls

System calls cannot be simulated consistently and are left out of the measure. The reported value understates the real cost. Switch to the Walltime instrument to measure them.

Learn more about system calls

See affected benchmarks

⚠️ Base and head ran on different runner types

Some benchmarks with significant performance changes were compared across different runner types, so the changes might not come from the code. Run base and head on the same runner type, or use CodSpeed Macro Runners for more stable results.

Learn more about Macro Runners

Open the report in CodSpeed to investigate

ℹ️ 40 benchmarks have no flamegraph

The flamegraph wasn't generated: the symbols of the profiling trace could not be resolved, because no information about the language frames was collected.

See affected benchmarks

⚡ 3 improved benchmarks
❌ 8 regressed benchmarks
✅ 2397 untouched benchmarks
⏩ 12 skipped benchmarks1

Warning

Please fix the performance issues or acknowledge them on CodSpeed.

Performance Changes

Mode Benchmark BASE HEAD Efficiency
❌ Walltime mul_u64_nonnull_neon 21 µs 39.3 µs -46.71%
❌ CPU Simulation test_scan_filter[int] 2.8 ms 3.4 ms -19.49%
❌ CPU Simulation test_scan_filter_to_arrow_array[int_float] 3.5 ms 4.1 ms -13.67%
❌ Walltime bitpack_blocked_compress_avx2 6.8 µs 7.6 µs -11.29%
❌ CPU Simulation test_repeated_scan_filter[int_float] 3.4 ms 3.9 ms -10.93%
❌ CPU Simulation random_i8[0.01] 104 µs 116.7 µs -10.91%
❌ CPU Simulation test_repeated_scan_filter[int] 2.4 ms 2.7 ms -10.46%
❌ CPU Simulation test_scan_filter_to_arrow_array[int] 2.5 ms 2.8 ms -10.21%
⚡ Walltime add_constant_shapes_neon[(32768, PerRowConstant)] 16.1 µs 14.4 µs +11.45%
⚡ Walltime add_constant_shapes_neon[(32768, PerRowNullableConstant)] 16.6 µs 15 µs +10.85%
⚡ Walltime add_constant_shapes_neon[(32768, ConstantPerRow)] 16 µs 14.4 µs +10.6%

Tip

Investigate this regression by commenting @codspeedbot fix this regression on this PR, or directly use the CodSpeed MCP with your agent.


Comparing mk/lambdas-v3 (aafd2d3) with develop (2e41681)

Open in CodSpeed

Footnotes

  1. 12 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩

@mhk197 mhk197 mentioned this pull request Oct 19, 2026
@mhk197
mhk197 added this pull request to stack #10614 October 19, 2026 14:32
@mhk197
mhk197 force-pushed the mk/lambdas-v3 branch 5 times, most recently from 4b7dd4a to 700462a Compare October 20, 2026 19:55
@mhk197
mhk197 marked this pull request as ready for review October 21, 2026 15:29
@mhk197
mhk197 requested a review from joseph-isaacs October 21, 2026 15:29
Signed-off-by: Matt Katz <mhkatz97@gmail.com>

This branch has not been deployed

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

Labels

changelog/feature A new feature

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant