Skip to content

Add ListTransform as a (higher order) scalar function - #9712

Draft
mhk197 wants to merge 1 commit into
mk/lambda-expressionsfrom
mk/scalar-list-transform
Draft

Add ListTransform as a (higher order) scalar function#9712
mhk197 wants to merge 1 commit into
mk/lambda-expressionsfrom
mk/scalar-list-transform

Conversation

@mhk197

@mhk197 mhk197 commented Sep 1, 2026

Copy link
Copy Markdown
Contributor

Summary

This PR adds array-level execution for list_transform. Since list_transform preserves input row cardinality, it can be represented as an ordinary ScalarFnArray:

ScalarFnArray {
    scalar_fn: ListTransform {
        options: LambdaBody,
    },
    children: [list, capture[0], capture[1], ...],
}

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

LambdaBody is the physical representation of the body of a lambda expression. It represents the body as a tree comprised of three types of nodes:

enum LambdaBody {
   Parameter(index: usize),
   Capture(index: usize),
   Scalar { scalar_fn: ScalarFnRef, children: Arc<[LambdaBody]> 
}
  • Parameter(i) is the ith lambda parameter.
  • Capture(i) refers to the ith captured array.
  • Scalar applies any registered scalar function to other lambda nodes.

For example:

x -> x + col("a")

is represented as:

LambdaBody::Scalar(AddRef, [LambdaBody::Parameter(0), LambdaBody::Capture(0)])

At evaluation time, Parameter and Capture nodes are replaced by arrays in the lambda invocation domain. Each Scalar node becomes a lazy ScalarFnArray and is passed through the existing array optimizer.

list_transform representation and type inference

ListTransform::try_new(list, lambda, captures) constructs a lazy ScalarFnArray whose first child is list and whose remaining children are captures in 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> becomes List<U, N>, where U is the lambda body dtype and N is 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:

  • List
  • ListView
  • FixedSizeList

The result is rebuilt using the same physical encoding family.

Preparing each physical encoding

For List:

  1. reset_offsets(false) slices away unreferenced leading/trailing elements and shifts offsets to
    start at zero.
  2. Per-row sizes are derived by subtracting adjacent offsets.
  3. The transformed elements are placed back under the reset offsets and original outer validity.

For FixedSizeList:

  1. Per-row sizes are represented by a constant array containing the fixed width.
  2. The transformed elements are rebuilt with the same width, outer length, and validity.

For ListView:

  1. The view is rebuilt with MakeZeroCopyToList before applying the lambda.
  2. This linearizes logical element occurrences, removes gaps and overlaps from the invocation
    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.
  3. The result remains a ListView, marked zero-copyable to List, with the rebuilt offsets, sizes,
    and validity.

Mapping outer rows to flattened invocations

list_transform has to bridge two domains:

outer row domain:          one value per list
flattened invocation domain: one value per logical list element

Given:

sizes = [2, 0, 3]

the implementation constructs two lazy PiecewiseSequenceArrays:

parent_indices = [0, 0, 2, 2, 2]
local_indices  = [0, 1, 0, 1, 2]

parent_indices maps every flattened invocation back to its containing outer row. It is encoded
as one constant run per list: starts are [0, 1, 2], run lengths are the list sizes, and every
multiplier is zero.

local_indices supplies optional lambda parameter 1. Each list contributes a sequence beginning
at 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_indices repeats it for every element of that row.

For example:

lists          = [[1, 2], [], [3, 4, 5]]
capture        = [10, 20, 30]
parent_indices = [0, 0, 2, 2, 2]
expanded       = capture.take(parent_indices)
               = [10, 10, 30, 30, 30]

The lambda x -> x + capture can then run as an ordinary array operation:

elements = [1, 2, 3, 4, 5]
capture  = [10, 10, 30, 30, 30]
result   = [11, 12, 33, 34, 35]

Captured children may themselves be lazy scalar-function arrays. The tests include capturing
list_length(input) without requiring it to be eagerly materialized before constructing
list_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:

logical lists  = [[1], null, [4]]
physical values = [1, 0, 4]
outer validity = [true, false, true]
lambda          = x -> 8 / x

The implementation expands outer validity with parent_indices:

invocation validity = [true, false, true]

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 the
hidden 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 the
body 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:

parameters[0] = valid flattened elements
parameters[1] = valid zero-based local indices, if referenced
captures[n]   = capture n expanded to valid flattened elements

The lambda recursively constructs and optimizes its scalar-function arrays. list_transform
checks 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_transform is itself just another scalar function, it can appear inside a
LambdaBody::Scalar node. Captures naturally spread again at each nested list boundary.

Example

input   = [[10, 20], [], [30, 40, 50]]
capture = [100, 200, 300]
lambda  = (x, i) -> x + capture + i

Execution prepares:

elements       = [10, 20, 30, 40, 50]
parent_indices = [0, 0, 2, 2, 2]
local_indices  = [0, 1, 0, 1, 2]
capture        = [100, 100, 300, 300, 300]

The flattened result is:

[110, 121, 330, 341, 352]

Reusing the original list structure produces:

[[110, 121], [], [330, 341, 352]]

@codspeed-hq

codspeed-hq Bot commented Sep 1, 2026

Copy link
Copy Markdown

Merging this PR will regress 2 benchmarks

⚠️ Unknown Walltime execution environment detected

Using the Walltime instrument on standard Hosted Runners will lead to inconsistent data.

For the most accurate results, we recommend using CodSpeed Macro Runners: bare-metal machines fine-tuned for performance measurement consistency.

⚡ 2 improved benchmarks
❌ 2 regressed benchmarks
✅ 2086 untouched benchmarks
⏩ 206 skipped benchmarks1
🗄️ 4 archived benchmarks run2

Warning

Please fix the performance issues or acknowledge them on CodSpeed.

Performance Changes

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)

Open in CodSpeed

Footnotes

  1. 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.

  2. 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.

@mhk197 mhk197 changed the title Add scalar list transform Add ListTransform as a (higher order) scalar function Sep 1, 2026
Signed-off-by: Matt Katz <mhkatz97@gmail.com>
@mhk197
mhk197 force-pushed the mk/scalar-list-transform branch from ba1bafe to c6f98c3 Compare September 1, 2026 05:08
@mhk197
mhk197 changed the base branch from develop to mk/lambda-expressions September 1, 2026 05:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant