perf: replace list_extract gather with take and zip

[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 can perform the gather in one operation while preserving the existing row-ordered validation and error behavior.

What changes are included in this PR?

• Build a UInt64Array of flat value indices and a separate out-of-bounds default mask.
• Gather elements with Arrow take, then substitute defaults with zip only for non-null-list, non-null-ordinal out-of-bounds rows.
• Strengthen the null-ordinal unit test with a non-null default value.

How are these changes tested?

• cargo fmt --all -- --check
• cargo test -p datafusion-comet-spark-expr --lib (566 tests passed)
• cargo clippy -p datafusion-comet-spark-expr --lib --tests -- -D warnings

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