0lai0 opened a new pull request, #5300:
URL: https://github.com/apache/datafusion-comet/pull/5300
## Which issue does this PR close?
Closes #5272.
## Rationale for this change
Follow-up to #5233. That PR routed all list-like `spark_size` through
Arrow's `length` kernel, which returns `Int64` for `LargeList` and then needs a
`cast_with_options(..., Int32, safe: false)` (and a `to_vec()` when patching
null slots). That left an extra Int64 length array and Int32 cast on the
LargeList path, and the bench showed it: `LargeList (10% null)` was only ~1.3x
faster than main after #5233, while the pure `List` path was ~12x.
Skip the length kernel entirely for `LargeList`. Subtract adjacent i64
offsets straight into `Int32` lengths, then apply the shared null → `-1`
rewrite. No intermediate `Int64Array`, no cast allocation.
`CometSize.convert` still wraps `size` in `CASE WHEN isnotnull(child)`, so
the production shape is the no-null path (693 ns here, within ~2x of the pure
`List` no-null path at 392 ns).
Overflow on LargeList now surfaces as `size(): list length exceeds i32::MAX`
(same as the scalar path) instead of the Arrow `cast_with_options(safe: false)`
error. Spark arrays are Int-indexed and cannot exceed `i32::MAX` elements per
row, so the overflow branch is unreachable from a Spark plan today; it is kept
as a defensive guard for non-Spark producers.
## What changes are included in this PR?
- `LargeList` gets its own path `spark_size_large_list_from_offsets`,
splitting off the shared `spark_size_list_like` (which now handles only `List`
and `FixedSizeList`).
- Hot path (when the full offset span fits in `i32`):
`offsets.windows(2).map(|w| (w[1] - w[0]) as i32).collect()`. Sound because
Arrow offsets are monotonically non-decreasing, so no per-row length can exceed
the full span.
- Checked fallback `spark_size_large_list_lengths_checked` for the rare case
where the offset span exceeds `i32::MAX`. Uses `i32::try_from` per row and
errors on overflow (matches the old `safe: false` cast contract). Skips
`try_from` on null rows since the caller overwrites them with `-1` anyway.
- Extract the shared null → `-1` rewrite into `ints_with_nulls_as_neg_one`
so `List` / `FixedSizeList` and `LargeList` cannot drift.
- Share the overflow error message via a `SIZE_OVERFLOW_MSG` constant used
by both the array and scalar paths.
- Bench: parameterise `create_large_list_array` on `with_nulls`, add
`spark_size: LargeList of long arrays` and `spark_size: LargeList, no nulls`
shapes to match the `List` coverage.
- Docs: add a row to
`docs/source/contributor-guide/expression-audits/collection_funcs.md` (date,
PR, technique, speedup, benchmark file) per `optimizing_expressions.md`.
### Benchmark (`array_size`, 8192 rows, on top of #5233)
| shape | main (post-#5233) | this PR | change |
|---|---|---|---|
| LargeList (10% null) | 7.13 µs | 1.08 µs | **6.6x** |
| LargeList of long arrays (10% null) | new shape | 1.06 µs | — |
| LargeList, no nulls (production path) | new shape | 693 ns | — |
Numbers from criterion comparison against a `pre-5272` baseline saved on
main (macOS). `List` and `FixedSizeList` benches are unchanged as expected
(their code path did not move).
## How are these changes tested?
- Existing `spark_size` unit tests plus three new ones:
- `test_spark_size_sliced_large_list_array`: pins slicing behavior
(mirrors the `List` slice test).
- `test_spark_size_large_list_length_overflow`: single-row `i32::MAX + 1`
length errors, exercising the checked fallback directly via a valid
`OffsetBuffer`.
- `test_spark_size_large_list_checked_null_row_skips_overflow`: null rows
with an overflowing offset delta must not error (they get rewritten to `-1` by
the caller, same as the fast path).
- `cargo test -p datafusion-comet-spark-expr --lib -- spark_size` (14
passed).
- `cargo clippy -p datafusion-comet-spark-expr --all-targets -- -D warnings`.
## Are there any user-facing changes?
No. Values remain bit-identical for in-range lengths; null still returns
`-1`. The only observable change is the overflow error message text (see
Rationale), which is not reachable from Spark's Int-bounded `Size`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
