Skip to content

perf: compute spark_size list lengths with Arrow length kernel - #5233

Merged
andygrove merged 11 commits into
apache:mainfrom
0lai0:refactor-5099-vector-length-kernel
Aug 5, 2026
Merged

perf: compute spark_size list lengths with Arrow length kernel#5233
andygrove merged 11 commits into
apache:mainfrom
0lai0:refactor-5099-vector-length-kernel

Conversation

@0lai0

@0lai0 0lai0 commented Aug 3, 2026

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Closes #5099

Rationale for this change

spark_size reimplemented List/LargeList/FixedSizeList sizing with a per-row builder loop that duplicates Arrow's length kernel.
Reuse the kernel and apply Spark's null to -1 semantics afterward

CometSize.convert wraps size in CASE WHEN isnotnull(child), so in a real Comet plan spark_size mostly sees null-free batches. That production shape is ~12x faster here, not the ~6x on 10%-null inputs.

What changes are included in this PR?

  • List / LargeList / FixedSizeList: use arrow::compute::kernels::length::length
  • LargeList Int64 → Int32 via cast_with_options(safe: false) so overflow errors instead of becoming -1
  • Null → -1 via into_parts + (!nulls).set_indices() (avoid zip / MutableArrayData, which regressed ~2x; also drop .expect from the kernel)
  • Fast path: null_count() == 0 returns the length kernel output as-is
  • Scalar path: read value_length instead of value(0).len()
  • Docs: drop PR-history ## size from array_funcs.md; keep offset-buffer row and add kernel-reuse row (~10x) in optimizing_expressions.md
  • Bench: add no-null List and LargeList shapes to benches/array_size.rs

NOTE:

Benchmark (array_size, 8192 rows)

shape main this PR change
list of short arrays (10% null) 9.30 µs 1.63 µs 5.7x
list of long arrays (10% null) 9.93 µs 1.61 µs 6.2x
list, no nulls (production path) 6.67 µs 0.56 µs 12x
LargeList (10% null) 9.70 µs 7.54 µs 1.3x

How are these changes tested?

  • Existing size unit tests (11 passed), including new: no-null List fast path, sliced ListArray, scalar LargeList, scalar FixedSizeList
  • cargo clippy -p datafusion-comet-spark-expr --all-targets -- -D warnings

@andygrove andygrove left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for picking this up. The kernel swap is the right call, the correctness looks solid, and it is a win on every shape I measured. I checked the two things most likely to break when moving from an index loop to buffer-level ops. Sliced input works, I wrote a throwaway test with a ListArray sliced to offset 2 and the offset handling plus the values() / nulls().iter() index alignment agree. Null semantics are preserved too, length() clones the input null buffer and the patch loop rewrites exactly those slots.

CI has not run on this yet, both workflow suites are sitting at action_required. I ran the relevant checks locally on 7575580 and they are clean:

  • cargo test -p datafusion-comet-spark-expr size, 8 passed
  • cargo clippy -p datafusion-comet-spark-expr --all-targets -- -D warnings, no warnings

I do have a few things I would like to see addressed before merge. The main one is that I think this PR is selling itself short.

The benchmark measures a path Comet does not take

CometSize.convert in spark/src/main/scala/org/apache/comet/serde/arrays.scala wraps the UDF in CASE WHEN isnotnull(child) THEN size(child) ELSE <legacy literal> END. DataFusion's CaseExpr::case_when_no_expr calls filter_record_batch on the WHEN predicate before it evaluates the THEN branch. So in a real Comet plan spark_size_array only ever receives a null-free array and the -1 rewrite never runs. Every shape in benches/array_size.rs has 10% nulls.

Could you add a no-null ListArray shape? I added one along with a LargeList shape and ran the base commit against this branch:

shape main this PR change
list of short arrays (10% null) 9.30 µs 3.89 µs 2.4x
list of long arrays (10% null) 9.93 µs 3.83 µs 2.6x
list, no nulls 6.67 µs 0.67 µs 10x
LargeList (10% null) 9.70 µs 9.79 µs no change

The path that actually runs in a query is 10x faster, not 2.2x. That is a much better number to be putting in the docs.

The null-patch loop dominates, and it is cheap to fix

In spark_size_list_like, nulls.iter().enumerate() walks every bit in the buffer, and int_lengths.values().to_vec() copies the whole values buffer. Only the null slots need touching, so iterating just the unset indices and taking ownership of the buffer instead of copying is a lot cheaper. Something along these lines:

let (_, values, nulls) = int_lengths.clone().into_parts();
let Some(nulls) = nulls else {
    return Ok(Arc::new(Int32Array::new(values, None)));
};
let mut values = values.to_vec();
for i in (!nulls.inner()).set_indices() {
    values[i] = -1;
}
Ok(Arc::new(Int32Array::from(values)))

I benched this and it takes the 10%-null list shapes from 3.89 µs down to 1.56 µs, another 2.5x. It also drops the .expect("null_count > 0 implies a null buffer"), which is worth doing on its own since we try to keep panics out of the expression kernels.

LargeList does not benefit yet

The LargeList path now allocates an Int64 array, casts it to Int32 into a second allocation, then to_vecs into a third. That is why it comes out flat at 9.70 µs against 9.79 µs. Would you add a LargeList shape to benches/array_size.rs so this is visible? With the set_indices change above it drops to 7.29 µs, so it does become a win, just not from the kernel swap on its own. Restructuring it to build Int32 directly rather than Int64-then-cast is a bigger change and I am happy for that to be separate work if you would rather file an issue for it.

Map

The comment explains that length does not accept MapArray, which is true, but MapArray::offsets() is right there and the same windows(2) computation the kernel does would work. Is there a reason to leave Map on the per-row loop? If you would rather keep this PR focused I am fine with that, but could you file an issue so it does not get lost?

Docs

In docs/source/contributor-guide/expression-audits/array_funcs.md, the new ## size entry reads as PR history rather than audit content. The tuning date, the issue link, the zip / MutableArrayData regression note, and the speedup figure are all things that belong in the PR description. The rest of that file is per-Spark-version behavioral audit notes describing what the expression does today. Could this be dropped? The optimizing_expressions.md row already captures the technique for future contributors.

In docs/source/contributor-guide/optimizing_expressions.md, this replaces the "Read from the offset buffer directly" row, but that technique is still in use here for Map and in the scalar path, and it is a distinct general technique from reusing an Arrow kernel. Could you add the new row rather than swapping it out? And once the benchmark picks up a no-null shape, the speedup figure should reflect that number.

Tests

test_spark_size_array_no_nulls is a good addition. Two more would be worth having. A sliced ListArray case, since optimizing_expressions.md calls out slicing as the classic trap when moving to buffer-level ops, and pinning it protects against a future edit to this function. And a scalar case for ScalarValue::LargeList and ScalarValue::FixedSizeList, since both changed here and neither has coverage.

@0lai0

0lai0 commented Aug 5, 2026

Copy link
Copy Markdown
Contributor Author

Open a issue #5266, kept out of #5233 because the technique is different (offset-buffer windowing, not the length kernel) and needs its own benchmark. I’ll fix it as soon as possible.

@andygrove

Copy link
Copy Markdown
Member

Thanks for the revisions. Every item from the last round is addressed, and the benchmark story is much better now that the no-null shape is in.

I ran the checks locally on e28ea8c:

  • cargo test -p datafusion-comet-spark-expr size, 11 passed
  • cargo clippy -p datafusion-comet-spark-expr --all-targets -- -D warnings, clean
  • cargo bench --bench array_size

The benchmark numbers reproduce on my machine, against the main numbers I measured last time:

shape main this PR change
list of short arrays (10% null) 9.30 µs 1.63 µs 5.7x
list of long arrays (10% null) 9.93 µs 1.61 µs 6.2x
list, no nulls 6.67 µs 0.56 µs 12x
LargeList (10% null) 9.70 µs 7.54 µs 1.3x

I also went back through the two things that worried me about moving to buffer-level ops, and both hold. I read the length kernel in arrow-string 58.4.0 to confirm it covers all three types the dispatch now routes to it, and that FixedSizeList carries the input null buffer through. The set_indices patch is offset-correct because Not on a &BooleanBuffer rebuilds at offset 0 with the sliced length, which is exactly what test_spark_size_sliced_list_array pins. And the safe: false cast on the LargeList path visits null slots too, but a null list row's offsets still give a small in-range length, so it cannot raise a spurious overflow.

Two things I would like to see before merge.

A tracking issue for the LargeList path. Thanks for opening #5266 for Map. LargeList still needs one. It is at 1.3x while List gets 12x, and the reason is visible in the code: the kernel allocates an Int64 array, cast_with_options allocates a second, and to_vec allocates a third. Building Int32 straight from the offsets would close that. I do not want this to hold the PR up, but without an issue it will not get picked up.

The comment at size.rs:164. It says into_parts moves the values buffer without copying and that the work is O(null_count) rather than O(n). Line 171 is values.to_vec(), which is a full O(n) copy, and int_lengths.clone() bumps the buffer refcount while lengths is still alive so there is no way to reclaim it as written. The set_indices loop itself is genuinely O(null_count), but the function is not. Could you reword it? optimizing_expressions.md now points contributors here as the reference for this technique, so I would rather the comment not promise more than the code delivers.

@0lai0

0lai0 commented Aug 5, 2026

Copy link
Copy Markdown
Contributor Author

LargeList still uses Int64 length + cast, follow-up in #5272(build Int32 from offsets)

@0lai0

0lai0 commented Aug 5, 2026

Copy link
Copy Markdown
Contributor Author

Thanks @andygrove for review, PR updated
Modify the documentation and open an issue to track LargeList

@andygrove andygrove left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM pending CI. Thanks @0lai0

@andygrove
andygrove merged commit 1feef7a into apache:main Aug 5, 2026
70 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

size: compute list sizes with the arrow length kernel

2 participants