Skip to content

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

Open
0lai0 wants to merge 4 commits into
apache:mainfrom
0lai0:refactor-5099-vector-length-kernel
Open

perf: compute spark_size list lengths with Arrow length kernel#5233
0lai0 wants to merge 4 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

What changes are included in this PR?

  • List / LargeList / FixedSizeList: use arrow::compute::kernels::length::length
  • LargeList Int64 to Int32 via cast_with_options(safe: false) so overflow errors instead of becoming -1
  • Null to -1 by patching the values buffer (avoid zip / MutableArrayData, which regressed ~2x)
  • Scalar path: read value_length instead of value(0).len()
  • Map unchanged (length does not support MapArray)
  • Docs: size entry in array_funcs audit + optimizing_expressions table

Benchmark (array_size, 8192 rows, ~10% nulls)

shape before after change
list of short arrays 7.53 µs 3.36 µs −55.7% (~2.2x)
list of long arrays 7.65 µs 3.30 µs −55.8% (~2.3x)

How are these changes tested?

Existing size unit tests plus a new no-null List case covering the null_count == 0 fast path

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

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