Skip to content

perf: use aligned pointer reads for SparkUnsafeRow field accessors#3670

Merged
andygrove merged 4 commits intoapache:mainfrom
andygrove:perf/aligned-reads-spark-unsafe-row
Mar 12, 2026
Merged

perf: use aligned pointer reads for SparkUnsafeRow field accessors#3670
andygrove merged 4 commits intoapache:mainfrom
andygrove:perf/aligned-reads-spark-unsafe-row

Conversation

@andygrove
Copy link
Member

Which issue does this PR close?

Performance optimization for native shuffle row-to-columnar conversion.

Rationale for this change

The SparkUnsafeObject trait previously used a one-size-fits-all approach for reading primitive values: creating a byte slice via from_raw_parts, then calling from_le_bytes(slice.try_into().unwrap()). This incurs unnecessary overhead from slice creation, try_into, and unwrap on every field access.

The two implementors of SparkUnsafeObject have fundamentally different alignment guarantees:

  • SparkUnsafeRow: All field offsets are always 8-byte aligned. The JVM guarantees 8-byte alignment on the base address, bitset_width is a multiple of 8, and each field slot is 8 bytes. This means aligned ptr::read() is safe and optimal.
  • SparkUnsafeArray: The array base address may be unaligned when nested within a row's variable-length region (accessed via arbitrary byte offset), so ptr::read_unaligned() is required for correctness.

What changes are included in this PR?

  • Move primitive accessor method implementations (get_int, get_long, get_float, get_double, etc.) out of the trait defaults and into each concrete impl block via a impl_primitive_accessors! macro parameterized on the read method (read vs read_unaligned).
  • SparkUnsafeRow uses ptr::read() (aligned) — avoids the from_le_bytes + slice overhead.
  • SparkUnsafeArray uses ptr::read_unaligned() — correct for potentially unaligned data.
  • Switch is_null_at and set_not_null_at in SparkUnsafeRow from read_unaligned/write_unaligned to aligned read/write, since the null bitset words are always at 8-byte aligned offsets within the row.

How are these changes tested?

Existing tests cover these code paths. The change is purely an optimization of pointer read methods — no behavioral change. cargo clippy and cargo check pass cleanly.

SparkUnsafeRow field offsets are always 8-byte aligned (the JVM
guarantees 8-byte alignment on the base address, bitset_width is a
multiple of 8, and each field slot is 8 bytes). This means we can
safely use ptr::read() instead of the from_le_bytes(slice) pattern
for all typed accesses, avoiding slice creation and try_into overhead.

Move primitive accessor implementations out of the SparkUnsafeObject
trait defaults and into each concrete impl via a macro parameterized
on the read method:
- SparkUnsafeRow uses ptr::read() (aligned)
- SparkUnsafeArray uses ptr::read_unaligned() (may be unaligned when
  nested in a row's variable-length region)

Also switch is_null_at/set_not_null_at in SparkUnsafeRow from
read_unaligned/write_unaligned to aligned read/write, since the null
bitset is always at 8-byte aligned offsets within the row.
@andygrove andygrove marked this pull request as ready for review March 11, 2026 19:36
The test_append_null_struct_field_to_struct_builder test used a plain
[u8; 16] stack buffer with no alignment guarantee. Since is_null_at
performs aligned i64 reads, Miri flags this as undefined behavior when
the buffer lands at a non-8-byte-aligned address.

Wrap the buffer in a #[repr(align(8))] struct to match the alignment
that real Spark UnsafeRow data always has from JVM memory.
@andygrove andygrove force-pushed the perf/aligned-reads-spark-unsafe-row branch from 729f044 to 8bc5761 Compare March 11, 2026 20:04
Copy link
Contributor

@mbutrovich mbutrovich left a comment

Choose a reason for hiding this comment

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

Thanks @andygrove! Like I mentioned before, this seems to compile down the same so not expecting performance differences, but it's more readable/maintainable.

@andygrove
Copy link
Member Author

Thanks for the review @mbutrovich

@andygrove andygrove merged commit a05c568 into apache:main Mar 12, 2026
206 of 213 checks passed
@andygrove andygrove deleted the perf/aligned-reads-spark-unsafe-row branch March 12, 2026 13:51
@martin-g
Copy link
Member

I was just looking into this PR.
A note: ptr::read()/read_unaligned() read in native order. Spark Unsafe uses Little Endian. So, the new code won't work on Big Endian machines.

@mbutrovich
Copy link
Contributor

I was just looking into this PR. A note: ptr::read()/read_unaligned() read in native order. Spark Unsafe uses Little Endian. So, the new code won't work on Big Endian machines.

AFAIK arrow-rs' position on big endian support has not changed (i.e., it does not support big endian). apache/arrow-rs#6917 (comment)

It's possible we've acquired some code that supports both targets, but we likely only need to consider little endian.

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.

3 participants