feat: use aligned slice access during bulk append in SparkUnsafeArray

[sandugood]

Which issue does this PR close?

Closes #4626

Rationale for this change

Currently, there is a bottleneck in performance during bulk append in the SparkUnsafeArray implementation (in macro and respective Builder types for bool, date32 and timestamp). If the array is NULLABLE there is a hotspot:

• we check each corresponding element at the index for nullability using Self::is_null_in_bitset() which is suboptimal (see the benchmarking results below)

What changes are included in this PR?

With this PR we change the flow of execution for nullable arrays:

• reading raw bytes from the null bitmap and flipping them according to Spark vs Arrow logic
• constructing BooleanBuffer from the bytes and then creating a NullBuffer
• creating a PrimitiveArray (with type specified in the macro_rules! and inside the according methods in the SparkUnsafeArray for several dtypes) from slice data
• appending this array to the current builder

How are these changes tested?

All of the basic tests pass.
Here is the benchmarking result (copied from the comment below).
Result from with_nulls version of benches (no_nulls wasnt affected with this PR)

array_type	main (current)	incoming PR	time_reduce
i32	29.274μs	2.4256μs	-91.71%
i64	27.260μs	4.7887μs	-82.43%
f64	29.352μs	4.6621μs	-84.11%
date32	26.863μs	2.5345μs	-89.89%
timestamp	45.818μs	4.3659μs	-90.47%

[sandugood]

When running CI pipeline, got few UB.
It was caused presumably by the bit_offset and Spark <-> Arrow null bitmask tinkering.

Switched to a less optimal way, which still provided boost in performance. By materializing &[bool] and passing it to builder's append_values()

array_type	main (current)	incoming PR	time_reduce
i32	30.030μs	21.250μs	-29.23%
i64	27.572μs	21.996μs	-20.22%
f64	30.059μs	21.889μs	-27.17%
date32	28.693μs	21.041μs	-26.66%
timestamp	46.832μs	21.888μs	-53.26%

[sandugood]

Could somebody potentially trigger the CI run? Thanks in advance

[andygrove]

Thanks for picking this up. The shape of the change is exactly what #4626 was asking for: keep the runtime alignment check as load-bearing, and replace the per-element null branch with PrimitiveBuilder::append_values(values, &is_valid). The benchmark numbers look great.

A few things to look at, the first one is a correctness concern.

1. Boolean nullable path may have UB. In append_booleans_to_builder, the new nullable path casts *const u8 to *const bool and iterates:

   let slice = unsafe { std::slice::from_raw_parts(self.element_offset as *const bool, num_elements) };
   for (idx, &value) in slice.iter().enumerate() { ... }

   Rust requires every bool value to be exactly 0 or 1. For null slots, the underlying byte is JVM-uninitialized memory, so the let &value materialization in the iterator pulls out a bool from arbitrary bytes which is UB even when the value is then ignored. The non-nullable branch a few lines down correctly uses *const u8 plus value != 0. Could you mirror that pattern in the nullable branch? cargo +nightly miri test would give a definitive answer, and it might be related to the UB you mentioned hitting earlier.

2. Boolean nullable path still branches per-element. Once the cast is fixed, the loop still does a per-element null check, which is the shape perf: use aligned slice access in SparkUnsafeArray bulk append #4626 specifically called out as the bigger win for booleans. BooleanBuilder::append_values(values: &[bool], is_valid: Option<&[bool]>) would absorb the validity in bulk and remove the branch. Worth taking here if straightforward.

3. Vec<bool> allocation per call. The simpler Vec<bool> form is fine for landing, but perf: use aligned slice access in SparkUnsafeArray bulk append #4626 suggested feeding a BooleanBuffer directly to the builder to skip the bool materialization. If you'd rather defer that, a follow-up issue with the link to the previous discussion would be helpful. At minimum a Vec::with_capacity(num_elements) instead of collect would avoid one growth.

4. Stale assert message in append_dates_to_builder. The new debug_assert! message reads "append_timestamps: element_offset is null". Copy-paste from the timestamp path, should be append_dates.

5. Comment phrasing on the alignment fallback. Could you reword the // Note: alignment is not guaranteed - that is why do this comments to point at the existing explanation at list.rs:105 and say why (nested-array variable-length region)? The prior PR was bounced because the description left the runtime check looking redundant, and a clear comment here is the easiest way to keep that from happening again.

6. PR description. The benchmark numbers in your comment look like the right shape. Could you pull them into the PR description body and mention they came from cargo bench --bench array_element_append?

I'll approve the CI workflows so the matrix actually exercises the change. Out of curiosity, what was the UB you hit earlier? Worth knowing whether item 1 above was the cause or there's a separate hazard.

---

Disclosure: this review was assisted by an AI tool (Anthropic Claude via Claude Code). I read the diff, the linked issue, and the prior closed PR's review feedback before forming the comments above, and I take responsibility for them.

[sandugood]

Thank you for the review @andygrove

I have changed my approach and bechmarked it.
My current approach is about using arrow internal types and buffers (essentially NullBuffer, from the underlying BooleanBufer, which is constructed from the bitmap, although with bits flipped), without materializing and looping through Vec<> which was still an improvement, but not a drastic one.

Posting the benchmarking results, which I ran using cargo bench --bench array_element_append. Posting result from with_nulls version of benches (no_nulls wasnt affected with this PR)

array_type	main (current)	incoming PR	time_reduce
i32	29.274μs	2.4256μs	-91.71%
i64	27.260μs	4.7887μs	-82.43%
f64	29.352μs	4.6621μs	-84.11%
date32	26.863μs	2.5345μs	-89.89%
timestamp	45.818μs	4.3659μs	-90.47%

cc @mbutrovich, if you could take a quick look at this, please. Thanks!

[andygrove]

Thanks for the rework @sandugood, and for pulling the benchmark numbers and the cargo bench invocation into the description. The primitive, date, and timestamp paths are implemented the way #4626 was after. Reading the values as a slice and building a NullBuffer from the flipped Spark bitmap is the right shape, and the speedups are great.

One thing I want to call out as a nice catch: threading the timezone through and calling .with_timezone_opt(timezone) is actually load-bearing, not just defensive. PrimitiveBuilder::append_array asserts data_type equality (primitive_builder.rs:291), so without it the timestamp path would panic whenever the column carries a timezone. Good that you handled it.

A few items before this is ready:

1. The boolean nullable path is UB, and this revision moved it the wrong way. This is item 1 from the earlier review. The new nullable branch does:

let slice = unsafe { std::slice::from_raw_parts(self.element_offset as *const bool, num_elements) };
for (idx, &value) in slice.iter().enumerate() { ... }

Rust requires every bool to be exactly 0 or 1. For null slots the JVM byte is not guaranteed to be 0/1, and for (idx, &value) materializes the bool out of each byte even for slots that get skipped, which is UB. The previous version of this branch read *const u8 and did *ptr != 0, and the non-nullable branch right below still does the safe thing:

let values = unsafe { std::slice::from_raw_parts(self.element_offset as *const u8, num_elements) };
for &value in values { builder.append_value(value != 0); }

Could you mirror that *const u8 pattern in the nullable branch (builder.append_value(value != 0))? Since Spark stores booleans as one byte each rather than bit-packed, a per-element loop here is reasonable, so the u8 cast is really all that is needed. This is very likely the UB you hit in CI earlier, since there is a miri workflow. Worth confirming against a nullable boolean array with a null slot once the cast is in.

2. Copy-paste leftover in append_dates_to_builder. The debug_assert! message reads "append_timestamps: element_offset is null" and the comment above says contiguous i64 data, but this is the dates path. Both should say append_dates and i32.

3. A couple of targeted tests would help. The bitmap flip and the partial-byte handling are easy to get subtly wrong. Would you be open to adding unit tests for nullable arrays where num_elements is not a multiple of 8 with a null in the final byte (say 7, 9, 17), plus all-null and all-valid cases? That locks in the trailing-bit behavior, which currently relies on the BooleanBuffer length trimming the flipped padding bits.

4. Formatting. The process_primitive_lists macro in row.rs looks like the indentation drifted. A quick cargo fmt should clean it up.

5. Minor. The byte reinterpretation of the null word assumes little-endian bit ordering. That matches Arrow and every platform we target, so it is fine, but a one-line comment next to the null_words as *const u8 cast noting the little-endian assumption would help the next reader.

The boolean path is the only blocker. Everything else is small. Once that is a u8 read I think this is in good shape.

Disclosure: this review was assisted by an AI tool (Anthropic Claude via Claude Code). I read the diff, the linked issue, and the prior review thread before forming these comments, and I take responsibility for them.

[sandugood]

Thank you again for your review @andygrove
Addressed the comments.

[sandugood]

Could somebody, please, retrigger the CI run? I think it was due to a .md mismatch that was fixed on the latest main

[sandugood]

Should I rebase it or just simply merge the curr main?

[andygrove]

Note on this review: this was generated by an LLM (Claude Code) at my request while I worked through a review backlog. I have not verified the individual findings myself. Please treat everything below as suggestions to evaluate rather than as authoritative review feedback, and push back on anything that is wrong or already handled.

Replacing the per-element is_null_in_bitset probe with a bulk byte-level flip plus a NullBuffer is the right idea, and keeping the unaligned pointer path as a fallback rather than assuming alignment is important given the unsafe_object.rs note.

Four things.

Does append_array care about the builder's timezone?

The timezone is now threaded through append_timestamps_to_builder and applied with with_timezone_opt before builder.append_array(&arr). That suggests append_array compares data types and would reject a mismatch. Is that right? If so, is there a test with a non-UTC session timezone, so we know the threading actually matches what the builder was constructed with? If append_array does not check, then with_timezone_opt is doing nothing and the extra parameter can go.

The unaligned path is now untested

There are two implementations per type that must produce identical results, and the new tests construct buffers through make_i32_array, which I would expect to be aligned. So the aligned bulk path is covered and the per-element fallback is not, even though it is the one that runs when the JVM hands over a misaligned buffer.

Could one test force the misaligned case, for example by allocating with one extra byte of padding and offsetting the pointer? Even one test per type asserting the two paths agree on the same logical data would be worth having, since a divergence between them would be a wrong answer rather than a crash.

Trailing bits in the flipped mask

flipped inverts whole bytes, so for a num_elements that is not a multiple of 8 the padding bits in the last byte become 1. BooleanBuffer::new(buffer, 0, num_elements) bounds the logical length so those bits should be ignored, and null_count should be computed only over the logical range. I convinced myself this is fine, but it is the kind of thing worth a test with a deliberately awkward length such as 3, 9, or 65, mixing nulls into the final partial byte.

Benchmark numbers should be in the description

The description says the results are "copied from the comment below", but the body itself has no table. Comments get buried and do not survive into the commit history. Could you inline the before and after medians for each type and shape, and state explicitly that the no_nulls shapes are unaffected? That last point matters because it tells a reader the win is scoped to nullable arrays.

[sunchao]

Summary

• Prior state and problem: Nullable primitive arrays performed a null check and append for every element, limiting throughput for large arrays.
• Design approach: Convert Spark’s null bitmap into Arrow validity and append an intermediate PrimitiveArray for aligned inputs.
• Correctness / compatibility analysis: Spark’s relevant array-layout methods are identical across the supported versions examined. Focused validation passed for bitmap boundaries, forced misalignment, repeated appends, booleans, and timestamp timezones. The previously reported boolean concern is resolved.
• Key design decisions: Runtime alignment checks and scalar fallbacks remain. Timestamp metadata is preserved to satisfy append_array’s type check.
• Implementation sketch: The primitive macro, date helper, and timestamp helper construct validity and value buffers. Both timestamp dispatch sites forward the timezone.
• Behavioral changes worth calling out: Large arrays improve, but temporary allocations materially regress short-array conversion. One introduced P2 finding follows.
• Suggested improvements: Keep an allocation-free path for short arrays while retaining bulk conversion for larger inputs.

Reviewed the entire diff from 2fd3960dfee464b2da5a98c0e754c049d720002b to 9b4510590fd675710f251466f0fc24a00abb8c46. The PR is not a draft. Existing discussion was read before forming findings. No substantiated existing P1/P2 blockers remain.

Routed skills: review-comet-pr, review-comet-shuffle-pr, review-comet-ffi-pr, and review-comet-memory-pr. This checkout lacks .ai/skills/, so the repository copies in checkout 6161-2bf433338e77 were used.

Exact-head CI: 61 successful checks and 9 skipped, with no failures. Linux shuffle tests cover Spark 3.4–4.2, and macOS shuffle tests cover Spark 4.0. Native CI reports 673 passed and 5 skipped, including all five added tests. Its merge commit has the same tree as the reviewed head. The benchmark check was skipped.

Validation limits: The local full-crate test attempt stopped on uncached anyhow 1.0.103. A disposable release harness using the exact base/head append-method bodies and locked Arrow 58.3.0 passed nine correctness tests and reproduced the performance finding. This harness does not exercise JNI or complete shuffle queries. No local Miri run was performed. Project files remain unchanged.

[sunchao]

[P2] Could we retain an allocation-free path for short arrays? Every aligned nullable append now allocates temporary bitmap/value buffers, and append_array copies the values again into the builder. append_list_column_batch invokes this with <true> for every list, including lists without null elements. Using the exact base/head methods with Arrow 58.3.0, a release benchmark including ListBuilder construction, downcasting, and list appends measured singleton array<int> conversion at 19.34 ns/row before versus 212.48 ns/row after, about 11× slower. An allocation counter independently measured 0 versus 4,000 allocations for 1,000 singleton appends into a preallocated builder. This materially increases conversion CPU for JVM shuffle workloads containing many short arrays. Preserve the scalar loop below a measured size cutoff, or reuse storage to avoid these per-list allocations. The same approach applies to the date and timestamp helpers.

Evidence: Disposable harness: /tmp/comet-4672-append-check. The extracted append-method bodies were verified against both supplied SHAs. Run cargo run --release --offline --manifest-path /tmp/comet-4672-append-check/Cargo.toml --bin list_bench and the same command with --bin allocations. The list benchmark uses 8,192 rows per batch, 64 distinct aligned input arrays, initial values-builder capacity 32, and medians of nine alternating base/head measurements. Singleton arrays without nulls measured 19.34→212.48 ns/row, and singleton null arrays measured 20.86→199.24 ns/row. Shorter-path allocation overhead is confirmed by the deterministic 0→4 allocations per singleton append. A separate 10,000-element benchmark confirmed the intended large-array improvement, so the finding concerns the unconditional use of this path for short arrays.

[andygrove]

@sandugood It looks like this PR is no longer being updated, so I plan on branching from this and creating a new PR to address the final issues. I hope that is ok. Thanks for the contribution!

[andygrove]

@sandugood thanks for all the work on this. It's been quiet for a few weeks, so I've picked it up in #6824 to get it over the line. I kept your commits on that branch, so you'll be credited as co-author when it merges.

On top of your changes, it merges main, since #6604 reworked the same macro. It also takes the bitmap path only for arrays of 64 or more elements. Building the validity and value buffers costs four allocations per array, which made shorter arrays with nulls slower than the per-element loop (sunchao's review above has the details). Your review on #6824 would be very welcome.