Add ListTransform as a (higher order) scalar function - #9712
Conversation
Merging this PR will regress 2 benchmarks
|
| Mode | Benchmark | BASE |
HEAD |
Efficiency | |
|---|---|---|---|---|---|
| ❌ | WallTime | arrow_checked_add_u32_avx2[16384] |
17.7 µs | 21.3 µs | -17.01% |
| ❌ | WallTime | words_gather_scalar_avx2[65536] |
8.3 µs | 9.3 µs | -11.63% |
| ⚡ | WallTime | arrow_checked_add_u32_avx512[16384] |
21.3 µs | 17.7 µs | +20.58% |
| ⚡ | WallTime | mul_u32_nonnull_avx512 |
6.4 µs | 5.6 µs | +13.63% |
Tip
Investigate this regression by commenting @codspeedbot fix this regression on this PR, or directly use the CodSpeed MCP with your agent.
Comparing mk/scalar-list-transform (c6f98c3) with mk/lambda-expressions (d8e744f)
Footnotes
-
206 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. ↩
-
4 benchmarks were run, but are now archived. If they were deleted in another branch, consider rebasing to remove them from the report. Instead if they were added back, click here to restore them. ↩
ListTransform as a (higher order) scalar function
Signed-off-by: Matt Katz <mhkatz97@gmail.com>
ba1bafe to
c6f98c3
Compare
Summary
This PR adds array-level execution for
list_transform. Sincelist_transformpreserves input row cardinality, it can be represented as an ordinaryScalarFnArray:The lambda is structural metadata in the scalar-function options. The list and every captured array remain ordinary scalar-function children, so they participate in the existing lazy array execution model.
Lambda representation
LambdaBodyis the physical representation of the body of a lambda expression. It represents the body as a tree comprised of three types of nodes:Parameter(i)is theith lambda parameter.Capture(i)refers to theith captured array.Scalarapplies any registered scalar function to other lambda nodes.For example:
is represented as:
At evaluation time,
ParameterandCapturenodes are replaced by arrays in the lambda invocation domain. EachScalarnode becomes a lazyScalarFnArrayand is passed through the existing array optimizer.list_transformrepresentation and type inferenceListTransform::try_new(list, lambda, captures)constructs a lazyScalarFnArraywhose first child islistand whose remaining children arecapturesin the outer row domain. The scalar function arity is derived from the highest capture referenced by the lambda.The return dtype is derived from the lambda body while preserving the collection shape.
List<T, N>becomesList<U, N>, whereUis the lambda body dtype andNis the original outer nullability.Execution model
The lambda is evaluated once over arrays of flattened logical list elements. It is not called once
per scalar value.
Execution first resolves the input only as far as one of the supported physical list encodings:
ListListViewFixedSizeListThe result is rebuilt using the same physical encoding family.
Preparing each physical encoding
For
List:reset_offsets(false)slices away unreferenced leading/trailing elements and shifts offsets tostart at zero.
For
FixedSizeList:For
ListView:MakeZeroCopyToListbefore applying the lambda.domain, and gives null lists empty ranges. An element referenced by two views therefore appears
once for each logical occurrence, which is necessary because each occurrence can have a
different parent capture or local index.
ListView, marked zero-copyable toList, with the rebuilt offsets, sizes,and validity.
Mapping outer rows to flattened invocations
list_transformhas to bridge two domains:Given:
the implementation constructs two lazy
PiecewiseSequenceArrays:parent_indicesmaps every flattened invocation back to its containing outer row. It is encodedas one constant run per list: starts are
[0, 1, 2], run lengths are the list sizes, and everymultiplier is zero.
local_indicessupplies optional lambda parameter 1. Each list contributes a sequence beginningat zero: starts are all zero, run lengths are the list sizes, and every multiplier is one. This
array is only constructed when the lambda actually references parameter 1.
Expanding captures
Captures enter execution with one value per outer list row. Taking each capture by
parent_indicesrepeats it for every element of that row.For example:
The lambda
x -> x + capturecan then run as an ordinary array operation:Captured children may themselves be lazy scalar-function arrays. The tests include capturing
list_length(input)without requiring it to be eagerly materialized before constructinglist_transform.Outer nullability and hidden physical elements
A null outer list may still own physical element slots. Those values are not logically observable
and must not be passed to the lambda: doing so could produce an error or side effect from data that
does not exist logically.
For example:
The implementation expands outer validity with
parent_indices:It then filters the elements, parent indices, captures, and optional local indices before
evaluating the lambda. The lambda sees only
[1, 4], producing[8, 2]; it never evaluates thehidden zero.
For encodings whose physical shape still requires the filtered slot, the transformed values are
scattered back with an
InterleaveArray. Invalid positions receive a null placeholder when thebody dtype is nullable, or that dtype's zero value otherwise. The placeholder remains hidden by
the outer list validity. This preserves List offsets and FixedSizeList widths while maintaining
logical null semantics.
Evaluating and rebuilding
After domain preparation, the lambda receives:
The lambda recursively constructs and optimizes its scalar-function arrays.
list_transformchecks that the result has the inferred body dtype and exactly one value per invocation, restores
any filtered physical slots, and rebuilds the original list encoding around the transformed
elements.
Because nested
list_transformis itself just another scalar function, it can appear inside aLambdaBody::Scalarnode. Captures naturally spread again at each nested list boundary.Example
Execution prepares:
The flattened result is:
Reusing the original list structure produces: