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.
• Retain the original single-pass MutableArrayData implementation whenever an explicit default is supplied, including an explicit null default.
• 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 (566 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.