apache / apache/datafusion-comet

list_extract: replace per-row MutableArrayData gather with take and zip

Open
#5,100 1 comment 0 reactions 1 assignee Claimed by @peterxcli View on GitHub
area:expressions enhancement
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.