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 2a59e0899 (microseconds; lower is better):

Type	OOB rows	Default	Base	Patched	Change
Int32	0%	null	49.228	25.821	-47.5%
Int32	0%	non-null	48.056	47.547	-1.1%
UTF-8	0%	null	70.956	47.984	-32.4%
UTF-8	0%	non-null	70.416	68.071	-3.3%
Int32	50%	null	56.537	25.699	-54.5%
Int32	50%	non-null	51.957	47.144	-9.3%
UTF-8	50%	null	75.447	40.718	-46.0%
UTF-8	50%	non-null	74.368	69.192	-7.0%

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