apache / apache/datafusion-comet
list_extract: replace per-row MutableArrayData gather with take and zip
- Dominant language
- Scala
- Stars
- 1.3k
- Forks
- 373
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 198
Description
### What is the problem the feature request solves?
`list_extract` in `native/spark-expr/src/array_funcs/list_extract.rs` gathers one element per row from the flat values buffer with a per-row `MutableArrayData` loop (extend one element, or extend from the one-row default array, or extend nulls). This is a hand-rolled `take`, and per-row `extend` calls are slower than a single gather.
### Describe the potential solution
Two-pass structure:
1. Compute an index array per row: `Some(start + adjusted_index)`, `None` for null list, null ordinal, or out-of-bounds, erroring early for `fail_on_error` exactly as today (the row-by-row index computation preserves error-before-result ordering).
2. One `arrow::compute::take(values, &indices, None)` (take handles null indices as null output and works on any element type, including nested).
3. Default substitution with `arrow::compute::kernels::zip::zip(mask, default_scalar, taken)`, where the mask fires only on non-null-list, non-null-index out-of-bounds rows.
### Additional context
Null propagation is identical; the mask construction for default substitution is the only subtle part. Likely a performance win (one gather versus per-row extends).
Found during an audit of native code that replicates existing arrow-rs kernels.
Contributor guide
Assessment
This issue has not been assessed yet.