perf: optimize list_extract without defaults using Arrow take

[peterxcli]

Which issue does this PR close?

Closes #5100.

Rationale for this change

list_extract gathered one element per row with repeated MutableArrayData::extend calls. Arrow's take kernel is substantially faster when no out-of-bounds default expression is supplied, but combining take with zip regressed workloads with explicit defaults, especially at high out-of-bounds rates.

What changes are included in this PR?

• Use one Arrow take for the common no-default path; null indices directly produce null output, so no zip pass is needed.
• Route explicit null defaults through the same take fast path (a null default is indistinguishable from no default); retain the original single-pass MutableArrayData implementation only for non-null defaults, where it benchmarks better at high out-of-bounds rates.
• Preserve null-list, null-ordinal, one-based/zero-based, and ANSI error behavior.
• Add a Criterion microbenchmark covering Int32 and UTF-8 values, null and non-null defaults, and 0% and 50% out-of-bounds ordinals.

How are these changes tested?

• cargo fmt --all -- --check
• cargo test -p datafusion-comet-spark-expr --lib (628 passed)
• cargo clippy -p datafusion-comet-spark-expr --all-targets -- -D warnings
• cargo bench -p datafusion-comet-spark-expr --bench list_extract --no-run

Criterion medians for 8,192 five-element lists, comparing base dba2ce49e with the patched implementation (milliseconds; lower is better):

Type	Input	Default	Base	Patched	Change
Int32	0% OOB	none	0.049228 ms	0.025821 ms	-47.5% (1.9x)
Int32	0% OOB	non-null	0.048056 ms	0.047547 ms	-1.1% (1.0x)
UTF-8	0% OOB	none	0.070956 ms	0.047984 ms	-32.4% (1.5x)
UTF-8	0% OOB	non-null	0.070416 ms	0.068071 ms	-3.3% (1.0x)
Int32	50% OOB	none	0.056537 ms	0.025699 ms	-54.5% (2.2x)
Int32	50% OOB	non-null	0.051957 ms	0.047144 ms	-9.3% (1.1x)
UTF-8	50% OOB	none	0.075447 ms	0.040718 ms	-46.0% (1.9x)
UTF-8	50% OOB	non-null	0.074368 ms	0.069192 ms	-7.0% (1.1x)
Int32	25% null lists + 25% null ordinals	none	0.051064 ms	0.025365 ms	-50.3% (2.0x)
UTF-8	25% null lists + 25% null ordinals	none	0.055546 ms	0.039790 ms	-28.4% (1.4x)

[andygrove]

Thanks for taking this on. The structure is exactly what #5100 asked for, and I like that the row-by-row index pass preserves the error-before-result ordering for fail_on_error. Null propagation looks unchanged to me as well, and strengthening test_list_extract_null_index with a non-null default is a nice touch, since it proves the default mask does not leak into null-list and null-ordinal rows.

Could you add a microbenchmark and post before and after numbers? There is no list_extract bench in native/spark-expr/benches/ yet, and array_size.rs is a close template. It would be good to include a case with a high proportion of out-of-bounds rows rather than only all-in-bounds, and a variable-width element type such as utf8 alongside int.

The reason I ask is that arrow's zip is itself built on MutableArrayData. When truthy is a scalar it calls mutable.extend(0, 0, 1) once per set bit in the mask (see arrow-select/src/zip.rs). So a batch with many out-of-bounds rows ends up paying the old per-row cost and a full take on top of it. A batch with no out-of-bounds rows at all still pays one extra whole-array allocation and copy that the old single-pass version did not.

I measured this locally. My machine had other work running, so rather than compare separate runs I put the old implementation, this PR, and a patched version all in a single bench binary. That keeps the ratios comparable under load. I also added assertions that all three variants produce identical output, which they do. Numbers below are 8192 rows at 5 elements per list, reproduced across three runs.

With a null default, which covers GetArrayItem, element_at, try_element_at, and the map_extract wrapper in planner.rs:

case	main	this PR	PR + guard below
int, 0% out of bounds	53.8 µs	30.9 µs (−43%)	27.1 µs (−50%)
int, 50% out of bounds	67.0 µs	88.6 µs (+32%)	26.7 µs (−60%)
utf8, 0% out of bounds	77.4 µs	68.2 µs (−12%)	57.3 µs (−26%)
utf8, 50% out of bounds	84.6 µs	134.6 µs (+59%)	46.9 µs (−45%)

For a null default the zip cannot change anything, because a null index already gathers as null through take. Would you consider skipping it in that case?

let taken = take(values.as_ref(), &UInt64Array::from(indices), None)?;
// A null index already gathers as null, so the zip is only needed when some row
// actually has to be replaced by a non-null default.
if default_value.is_null() || !use_default.iter().any(|b| *b) {
    return Ok(ColumnarValue::Array(taken));
}

That turns both regressions above into solid wins and improves the in-bounds cases further. All 566 datafusion-comet-spark-expr lib tests pass with it, and cargo clippy -D warnings is clean.

That still leaves the case of a non-null default, where the zip is genuinely needed and the guard does not help. As far as I can tell the only Spark path that sets defaultValueOutOfBound is split_part (stringExpressions.scala, ElementAt(StringSplitSQL(...), partNum, Some(Literal.create("", ...)), ...)). There the crossover sits somewhere between 10% and 25% out-of-bounds rows:

case	main	this PR	PR + guard
int, 1% out of bounds	53.4 µs	30.6 µs (−43%)	29.1 µs (−45%)
int, 10% out of bounds	53.1 µs	29.9 µs (−44%)	28.6 µs (−46%)
int, 25% out of bounds	52.3 µs	62.7 µs (+20%)	62.7 µs (+20%)
int, 50% out of bounds	54.1 µs	83.2 µs (+54%)	84.3 µs (+56%)
utf8, 25% out of bounds	74.0 µs	107.4 µs (+45%)	107.6 µs (+45%)
utf8, 50% out of bounds	77.4 µs	125.4 µs (+62%)	124.0 µs (+60%)

split_part(str, delim, n) with a constant n over rows that have fewer than n parts can land well past that crossover, so this looks worth handling. Would you be up for either keeping the MutableArrayData path when the default is non-null, or finding something cheaper than zip for it? I tried appending the default to the values array so a single take does everything, and it looked better at high out-of-bounds rates and worse at low ones, but my numbers for that variant were not self-consistent, so I would not treat it as a validated option without more careful measurement.

One thing I checked that is not a problem, just noting it so nobody else has to dig. zip rejects a default whose data type differs from the values type, and the equals_datatype guard in evaluate is looser than that, since it ignores nested field names. That looked like it might be a new failure mode, but MutableArrayData::new asserts on the same mismatch, so main panics where this PR returns an error. This PR is the better behavior. I also confirmed nested list and struct element types both work correctly on the new code.

[peterxcli]

@andygrove thanks for the review! review change is pushed. please take another look, thanks!

Could you add a microbenchmark and post before and after numbers? Include high out-of-bounds rates and utf8 alongside int.

Added an 8-case Criterion benchmark: Int32/UTF-8 × null/non-null default × 0%/50% OOB. Posted before/after numbers in the PR description.

For a null default the zip cannot change anything. Would you consider skipping it?

When no default expression is supplied, list_extract now uses a single Arrow take; the mask and zip are removed.

Would you be up for either keeping the MutableArrayData path when the default is non-null, or finding something cheaper than zip?

All explicit defaults, including split_part’s non-null default and explicit null defaults—retain the original single-pass MutableArrayData path.

[andygrove]

Thanks for the revision. You took both asks: the zip is gone entirely for the no-default case, and the non-null-default regression is avoided by simply not touching that path. The benchmark is a good addition and the numbers look right to me.

I checked that the fast path actually covers the callers that matter. CometGetArrayItem never sets default_value, CometElementAt only sets it when expr.defaultValueOutOfBound is present, and the map_extract wrapper in planner.rs:629 passes None. So arr[i], element_at, try_element_at, and map extraction all land on take.

I could not find a behavior change. The row loop in list_extract_without_default is a faithful copy of the original. Null list wins over null ordinal, fail_on_error still raises before any result is produced, and the error variant selection is now shared through out_of_bounds_error. Dropping .unwrap_or(element_type.try_into())? is a real improvement too, since the no-default path no longer has to materialize a ScalarValue from the element DataType just to throw it away.

Also worth saying for anyone reading the table later: the non-null-default deltas of -1.1% to -9.3% are noise, since that code path is unchanged. That is fine and expected.

A few things I would like to see addressed.

1. An explicit null default could take the fast path too (list_extract.rs:145)

The dispatch is on whether a default expression exists, not on whether it is null. A null default is semantically identical to no default. Out of bounds yields null either way, and take already gathers null for a null index. Right now Some(ScalarValue::Int32(None)) pays the MutableArrayData cost for nothing, which by your own table is roughly 2x on Int32.

Would you consider filtering it out after evaluation?

// A null default is indistinguishable from no default: an out-of-bounds index
// yields null either way, and `take` already gathers null for a null index.
let default_value = default_value.filter(|d| !d.is_null());

No Spark path sets a null defaultValueOutOfBound today, so this is not a live performance bug. It matters because the next expression wired through ListExtract with a null default literal would quietly land on the slow path with nothing to warn the author.

2. Build the indices with UInt64Builder rather than Vec<Option<u64>> (list_extract.rs:343)

Vec<Option<u64>> is 16 bytes per row, and UInt64Array::from(indices) then walks the whole thing again to split it into a values buffer and a null buffer. That is a temporary allocation twice the size of the final values buffer, plus an extra pass, in the hot loop of a PR whose point is to remove passes.

UInt64Builder::with_capacity(index_array.len()) writes straight into the final buffers, and the loop body becomes append_value and append_null instead of push. Could you try it and see what it does to the Int32 0% out-of-bounds case? That case is down to about 26 microseconds for 8192 rows now, so a 128 KB temporary should be measurable.

3. The two row loops have to stay in sync (list_extract.rs:309 and list_extract.rs:345)

The index resolution logic exists twice now, and the copies have to agree on four subtle things: null list taking precedence over null ordinal, the one-based versus zero-based adjustment, raising the error before producing any output, and which SparkError variant is used. out_of_bounds_error factors out the last one. The precedence and the error ordering are still duplicated by hand.

Would you be up for pulling the per-row decision into a small helper returning something like Gather(usize), Null, or Default, so each loop is just a match on the result? The two callers then differ only in how they consume it.

Related to that, the fail_on_error branch inside list_extract_without_default has no direct unit test. The ANSI behavior is covered end to end by element_at_ansi.sql and get_array_item_ansi.sql, so this is not a correctness hole today. It is a hole in the fast feedback loop for exactly the branch that is now duplicated. A short test asserting InvalidElementAtIndex for the one-based case and InvalidArrayIndex for the zero-based case on the no-default path would cover it.

4. SQL test coverage for complex element types (get_array_item.sql)

Every no-default extraction now gathers through take instead of MutableArrayData, including nested lists, structs, and maps. I checked nested list and struct elements by hand on the earlier revision and they were fine, and I have no reason to think they are not fine here. I would rather have that written down as a test than resting on a manual check. Could you add a couple of queries over array<array<int>> and array<struct<a:int,b:string>>, including a null list row and a null ordinal row, so the null-gather behavior is pinned for those types as well?

5. The benchmark lists have no nulls (benches/list_extract.rs:36)

Every list row is valid and every ordinal is non-null, which is the friendliest possible input for take. Null handling is the thing this path changes most, so a case with a meaningful fraction of null list rows and null ordinals would make the numbers more representative. It should still win, and it is useful to know by how much.

[peterxcli]

@andygrove thanks for another round of review, addressed all of your review. please take another look. TIA!

An explicit null default could take the fast path too

Explicit-null defaults now use the Arrow take fast path. Added OOB/null regression coverage and explicit-null benchmarks.

Build the indices with UInt64Builder rather than Vec<Option<u64>>

Replaced the temporary vector and conversion with direct builder writes. Isolated performance varied from +14.9% to −17.0%, so there is no stable builder-only claim.

The two row loops have to stay in sync

Added shared Gather/Null/Default row resolution. Added direct no-default ANSI tests for InvalidElementAtIndex and InvalidArrayIndex.

SQL test coverage for complex element types

Added nested-array and struct-array cases covering valid extraction, null lists, and null ordinals.

The benchmark lists have no nulls

Added disjoint 25% null-list and 25% null-ordinal workloads. result is updated in PR description.

[andygrove]

Note on this review: this was generated by an LLM (Claude Code) at my request while I worked through a review backlog. I have not verified the individual findings myself. Please treat everything below as suggestions to evaluate rather than as authoritative review feedback, and push back on anything that is wrong or already handled.

Nice result, and I appreciate that the benchmark table shows the non-null-default path too rather than only the case that improves. Factoring the per-row decision into resolve_row so both implementations share the index arithmetic and the ANSI error handling is a good way to keep them from drifting.

Three things.

The description contradicts the code on explicit null defaults

The description says the original MutableArrayData path is retained "whenever an explicit default is supplied, including an explicit null default". The code does the opposite:

let Some(default_value) = default_value.filter(|default| !default.is_null()) else {
    return list_extract_without_default(...)
};

An explicit null default falls through to the take path. That is the right behavior, since a null default and no default produce the same output, and it means the fast path applies more often. But the description should say so, because as written a reader would go looking for the branch that does the opposite.

Is check_bounds worth disabling?

take(values.as_ref(), &indices, None) uses default TakeOptions. Every index handed to take has already been validated by adjust_index against the row's list length, so the bounds check is redundant work proportional to the batch size. If Arrow's default is bounds-checking, passing TakeOptions { check_bounds: false } should be free speed. If the default is already unchecked, a one-line comment saying the indices are pre-validated would be worth having anyway.

Coverage of the new split

get_array_item.sql gains 15 lines. Since the change introduces two implementations that must agree, it would be good to know the SQL fixture exercises both. Specifically: a query with an explicit non-null default and out-of-bounds ordinals, and the same query without a default, against the same data. If those two are not both there, a future change to resolve_row could break one path silently.

Related, does anything cover element_at (one-based) versus arr[i] (zero-based) through the new take path, in ANSI mode, where the two produce different error types? out_of_bounds_error is only reachable from resolve_row now, so one test per branch would pin it.

[peterxcli]

@andygrove thanks for the pass — evaluated all three.

1. Description vs. code on null defaults — you're right, the description was stale. Explicit null defaults were moved onto the take fast path in the second revision (per the earlier review's point 1), and the description still described the first revision. I've updated it to match the code.

2. check_bounds — arrow's TakeOptions derives Default, so check_bounds is already false when passing None (verified in arrow-select 58.4.0, take.rs); there's no bounds-checking work to remove. I added a comment at the take call noting that every gathered index is pre-validated by resolve_row and that take masks null index slots (take_native returns T::default() for an out-of-bounds placeholder under a null slot rather than panicking), so the unchecked default is safe.

3. Coverage — the checks you describe exist, mostly from the previous round:

• test_list_extract_default_value and test_list_extract_null_index run the same data through all three configurations (no default, explicit null default, non-null default) and assert the outputs, so the two implementations are pinned against each other at the unit level.
• test_list_extract_without_default_ansi_errors asserts InvalidElementAtIndex (one-based) and InvalidArrayIndex (zero-based) specifically on the no-default/take path; element_at_ansi.sql and get_array_item_ansi.sql cover both end to end.
• At the SQL level a non-null default can't be expressed through arr[i] or element_at — the only Spark path that sets defaultValueOutOfBound is split_part, and split_part.sql already exercises out-of-bounds part numbers (e.g. split_part('a.b', '.', 4)).

[sunchao]

Summary

• Prior state and problem: list_extract gathered elements through per-row MutableArrayData calls, adding overhead to array indexing.
• Design approach: Use Arrow take for absent or null defaults. Retain MutableArrayData for non-null defaults, addressing the performance regression discussed in earlier reviews.
• Correctness / compatibility analysis: No introduced P1/P2 issues found within this review. Checked Spark implementations and relevant tests across supported versions 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0, plus the audit skill’s earlier reference versions. Index adjustment, null propagation, ANSI error selection and split_part defaults remain consistent with the prior implementation.
• Key design decisions: Shared resolve_row logic keeps both execution paths aligned. UInt64Builder avoids an intermediate Vec<Option<u64>>. The added abstraction is small and directly supports the two gathering strategies.
• Implementation sketch: Evaluate and validate the default, resolve each row, then gather through either take or the existing mutable-array path. The diff also extends benchmarks and nested-array/struct SQL coverage.
• Behavioral changes worth calling out: Explicit null defaults now use the fast path. No result or error changes were observed in validation. Earlier substantiated review concerns are addressed, with no unresolved P1/P2 blockers identified.
• Suggested improvements: None at the requested P1/P2 threshold.

Reviewed the entire three-file diff from 3810936b40af3562003120742f7d258b234e4dc5 to 78d65bd6a625b7616201daebc2f4bb5d0099264b. The PR is not a draft. Read existing reviews, issue comments and the snapshot’s empty inline-comment/thread collections.

Routed skills: review-comet-pr, audit-comet-expression, and optimize-comet-expression.

Exact-head CI: 56 successful checks and 10 skipped checks, with no failures. Required Checks, Rust tests and Linux expression suites for Spark 3.4–4.2 passed. Skipped checks include macOS, benchmark CI and upstream Spark SQL suites for 3.4–4.0.

Local validation: All four focused list_extract unit tests passed. A disposable comparison of the exact base and head implementations passed 11,520 cases covering List/LargeList, sliced and empty arrays, primitive and nested values, defaults, nulls and ANSI errors. The checkout remains clean.

Validation limits: Did not independently rerun release benchmarks or JVM/Spark suites locally. Performance assessment uses the implementation, benchmark coverage and existing measurements. JVM validation relies on exact-head CI.

[sunchao]

Merging, thanks @peterxcli !

[sunchao]

Filed #6225 to track the nested-array retained-buffer allocation increase as a separate follow-up. It includes the measured base/head comparison and suggested regression coverage, with the goal of preserving the primitive/string extraction gains.

[sunchao]

Summary

• Prior state and problem: list_extract gathered elements through per-row MutableArrayData calls, adding overhead to array indexing.
• Design approach: Use Arrow take for absent or null defaults and retain the original gathering strategy for non-null defaults.
• Correctness / compatibility analysis: No additional introduced P1/P2 issues found within this review. Index adjustment, null propagation, defaults and ANSI error selection remain consistent with the base implementation and relevant Spark semantics.
• Key design decisions: Shared resolve_row logic keeps both paths aligned. UInt64Builder avoids an intermediate allocation. These abstractions are small and directly support the two gathering strategies.
• Implementation sketch: Evaluate the default, resolve each row, then gather through take or MutableArrayData. Benchmarks cover defaults, out-of-bounds indices and null inputs. SQL tests add nested-array and struct extraction.
• Behavioral changes worth calling out: The existing P2 allocation concern in #6225 remains unresolved at native/spark-expr/src/array_funcs/list_extract.rs:392. For 8,192 rows containing an empty inner array followed by 128 integers, selecting the empty array retains 2,129,924 output-buffer bytes versus 65,600 at base, approximately 32.5× more. Both outputs contain the same empty arrays. This reproduces with absent and explicit null defaults. Native shuffle charges retained buffer capacity against its reservation.
• Suggested improvements: Address the tracked nested-array allocation regression through tighter capacity sizing or a targeted gathering path. No duplicate inline finding is included.

Reviewed the entire three-file diff from 3810936b40af3562003120742f7d258b234e4dc5 to 78d65bd6a625b7616201daebc2f4bb5d0099264b. The PR remains open and non-draft. Read existing reviews, issue comments, inline-comment/thread collections and linked issue #6225.

Routed skills: review-comet-pr, audit-comet-expression, and optimize-comet-expression. Compared relevant Spark sources and tests across supported versions 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0, plus the audit skill’s earlier reference versions.

Exact-head CI: 56 successful checks and 10 skipped, with no failures. Required Checks, Rust tests and Linux expression suites for Spark 3.4–4.2 passed. Benchmark CI, macOS and upstream Spark SQL suites for 3.4–4.0 were skipped.

Local validation: All four focused unit tests passed. A freshly compiled disposable comparison of the exact base/head implementations passed 11,520 cases covering values, nulls, defaults, nested types, slices and errors. Independently reproduced #6225 through ListExtract::evaluate. The checkout remains clean.

Validation limits: Did not rerun JVM suites or release benchmarks locally. Several existing ANSI SQL error cases are ignored, so those fixtures are not passing error-path evidence. Memory measurements are retained buffer capacity, not RSS. No end-to-end spill or runtime increase was measured.