From 7f13df2dfdcc09fd005b5d0b38afb057eec40b43 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 10 Mar 2026 18:11:52 -0600 Subject: [PATCH 1/2] perf: use aligned slice access in SparkUnsafeArray nullable bulk append SparkUnsafeArray guarantees natural alignment for element data: the header is 8 + ceil(n/64)*8 bytes (always 8-byte aligned), and elements are at element_size stride from the aligned base. This means we can create a slice upfront and use indexed access instead of per-element read_unaligned() with manual pointer arithmetic. Changes: - Nullable paths in impl_append_to_builder macro, append_booleans, append_timestamps, and append_dates now create a slice upfront and iterate with enumerate, eliminating read_unaligned and ptr.add(1) - Non-nullable paths simplified: removed runtime alignment checks since alignment is guaranteed, always use append_slice bulk copy - Added debug_assert for alignment invariant --- .../execution/shuffle/spark_unsafe/list.rs | 154 +++++++----------- 1 file changed, 57 insertions(+), 97 deletions(-) diff --git a/native/core/src/execution/shuffle/spark_unsafe/list.rs b/native/core/src/execution/shuffle/spark_unsafe/list.rs index 72610d2d82..5b65107542 100644 --- a/native/core/src/execution/shuffle/spark_unsafe/list.rs +++ b/native/core/src/execution/shuffle/spark_unsafe/list.rs @@ -46,40 +46,31 @@ macro_rules! impl_append_to_builder { return; } + // SAFETY: element_offset points to contiguous, naturally aligned data of + // length num_elements. Alignment is guaranteed by SparkUnsafeArray layout: + // header is 8 + ceil(n/64)*8 bytes (always 8-byte aligned), and elements + // are at element_size stride from an aligned base. + debug_assert!(self.element_offset != 0, "element_offset is null"); + let ptr = self.element_offset as *const $element_type; + debug_assert!( + (ptr as usize).is_multiple_of(std::mem::align_of::<$element_type>()), + "element_offset is not properly aligned" + ); + let values = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; + if NULLABLE { - let mut ptr = self.element_offset as *const $element_type; let null_words = self.null_bitset_ptr(); debug_assert!(!null_words.is_null(), "null_bitset_ptr is null"); - debug_assert!(!ptr.is_null(), "element_offset pointer is null"); - for idx in 0..num_elements { + for (idx, &val) in values.iter().enumerate() { // SAFETY: null_words has ceil(num_elements/64) words, idx < num_elements - let is_null = unsafe { Self::is_null_in_bitset(null_words, idx) }; - - if is_null { + if unsafe { Self::is_null_in_bitset(null_words, idx) } { builder.append_null(); } else { - // SAFETY: ptr is within element data bounds - builder.append_value(unsafe { ptr.read_unaligned() }); + builder.append_value(val); } - // SAFETY: ptr stays within bounds, iterating num_elements times - ptr = unsafe { ptr.add(1) }; } } else { - // SAFETY: element_offset points to contiguous data of length num_elements - debug_assert!(self.element_offset != 0, "element_offset is null"); - let ptr = self.element_offset as *const $element_type; - // Use bulk copy when data is properly aligned, fall back to - // per-element unaligned reads otherwise - if (ptr as usize).is_multiple_of(std::mem::align_of::<$element_type>()) { - let slice = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; - builder.append_slice(slice); - } else { - let mut ptr = ptr; - for _ in 0..num_elements { - builder.append_value(unsafe { ptr.read_unaligned() }); - ptr = unsafe { ptr.add(1) }; - } - } + builder.append_slice(values); } } }; @@ -191,11 +182,13 @@ impl SparkUnsafeArray { return; } - let mut ptr = self.element_offset as *const u8; + // SAFETY: element_offset points to contiguous u8 data of length num_elements. debug_assert!( - !ptr.is_null(), + self.element_offset != 0, "append_booleans: element_offset pointer is null" ); + let values = + unsafe { std::slice::from_raw_parts(self.element_offset as *const u8, num_elements) }; if NULLABLE { let null_words = self.null_bitset_ptr(); @@ -203,24 +196,17 @@ impl SparkUnsafeArray { !null_words.is_null(), "append_booleans: null_bitset_ptr is null" ); - for idx in 0..num_elements { + for (idx, &val) in values.iter().enumerate() { // SAFETY: null_words has ceil(num_elements/64) words, idx < num_elements - let is_null = unsafe { Self::is_null_in_bitset(null_words, idx) }; - - if is_null { + if unsafe { Self::is_null_in_bitset(null_words, idx) } { builder.append_null(); } else { - // SAFETY: ptr is within element data bounds - builder.append_value(unsafe { *ptr != 0 }); + builder.append_value(val != 0); } - // SAFETY: ptr stays within bounds, iterating num_elements times - ptr = unsafe { ptr.add(1) }; } } else { - for _ in 0..num_elements { - // SAFETY: ptr is within element data bounds - builder.append_value(unsafe { *ptr != 0 }); - ptr = unsafe { ptr.add(1) }; + for &val in values { + builder.append_value(val != 0); } } } @@ -235,47 +221,34 @@ impl SparkUnsafeArray { return; } + // SAFETY: element_offset points to contiguous, naturally aligned i64 data. + debug_assert!( + self.element_offset != 0, + "append_timestamps: element_offset is null" + ); + let ptr = self.element_offset as *const i64; + debug_assert!( + (ptr as usize).is_multiple_of(std::mem::align_of::()), + "append_timestamps: element_offset is not properly aligned" + ); + let values = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; + if NULLABLE { - let mut ptr = self.element_offset as *const i64; let null_words = self.null_bitset_ptr(); debug_assert!( !null_words.is_null(), "append_timestamps: null_bitset_ptr is null" ); - debug_assert!( - !ptr.is_null(), - "append_timestamps: element_offset pointer is null" - ); - for idx in 0..num_elements { + for (idx, &val) in values.iter().enumerate() { // SAFETY: null_words has ceil(num_elements/64) words, idx < num_elements - let is_null = unsafe { Self::is_null_in_bitset(null_words, idx) }; - - if is_null { + if unsafe { Self::is_null_in_bitset(null_words, idx) } { builder.append_null(); } else { - // SAFETY: ptr is within element data bounds - builder.append_value(unsafe { ptr.read_unaligned() }); + builder.append_value(val); } - // SAFETY: ptr stays within bounds, iterating num_elements times - ptr = unsafe { ptr.add(1) }; } } else { - // SAFETY: element_offset points to contiguous i64 data of length num_elements - debug_assert!( - self.element_offset != 0, - "append_timestamps: element_offset is null" - ); - let ptr = self.element_offset as *const i64; - if (ptr as usize).is_multiple_of(std::mem::align_of::()) { - let slice = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; - builder.append_slice(slice); - } else { - let mut ptr = ptr; - for _ in 0..num_elements { - builder.append_value(unsafe { ptr.read_unaligned() }); - ptr = unsafe { ptr.add(1) }; - } - } + builder.append_slice(values); } } @@ -289,47 +262,34 @@ impl SparkUnsafeArray { return; } + // SAFETY: element_offset points to contiguous, naturally aligned i32 data. + debug_assert!( + self.element_offset != 0, + "append_dates: element_offset is null" + ); + let ptr = self.element_offset as *const i32; + debug_assert!( + (ptr as usize).is_multiple_of(std::mem::align_of::()), + "append_dates: element_offset is not properly aligned" + ); + let values = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; + if NULLABLE { - let mut ptr = self.element_offset as *const i32; let null_words = self.null_bitset_ptr(); debug_assert!( !null_words.is_null(), "append_dates: null_bitset_ptr is null" ); - debug_assert!( - !ptr.is_null(), - "append_dates: element_offset pointer is null" - ); - for idx in 0..num_elements { + for (idx, &val) in values.iter().enumerate() { // SAFETY: null_words has ceil(num_elements/64) words, idx < num_elements - let is_null = unsafe { Self::is_null_in_bitset(null_words, idx) }; - - if is_null { + if unsafe { Self::is_null_in_bitset(null_words, idx) } { builder.append_null(); } else { - // SAFETY: ptr is within element data bounds - builder.append_value(unsafe { ptr.read_unaligned() }); + builder.append_value(val); } - // SAFETY: ptr stays within bounds, iterating num_elements times - ptr = unsafe { ptr.add(1) }; } } else { - // SAFETY: element_offset points to contiguous i32 data of length num_elements - debug_assert!( - self.element_offset != 0, - "append_dates: element_offset is null" - ); - let ptr = self.element_offset as *const i32; - if (ptr as usize).is_multiple_of(std::mem::align_of::()) { - let slice = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; - builder.append_slice(slice); - } else { - let mut ptr = ptr; - for _ in 0..num_elements { - builder.append_value(unsafe { ptr.read_unaligned() }); - ptr = unsafe { ptr.add(1) }; - } - } + builder.append_slice(values); } } } From 471ca5767e607117863fbdc1033fdf03b2f5bc49 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 10 Mar 2026 18:59:03 -0600 Subject: [PATCH 2/2] fix: restore alignment check for SparkUnsafeArray bulk append SparkUnsafeArray elements may not be naturally aligned (e.g., i64 at 4-byte offset). Restore runtime alignment check: use slice-based access when aligned, fall back to per-element read_unaligned when not. The nullable path now has the same aligned fast path as non-nullable, which was the original goal of this PR. --- .../execution/shuffle/spark_unsafe/list.rs | 125 ++++++++++++------ 1 file changed, 83 insertions(+), 42 deletions(-) diff --git a/native/core/src/execution/shuffle/spark_unsafe/list.rs b/native/core/src/execution/shuffle/spark_unsafe/list.rs index 5b65107542..d194df739c 100644 --- a/native/core/src/execution/shuffle/spark_unsafe/list.rs +++ b/native/core/src/execution/shuffle/spark_unsafe/list.rs @@ -46,31 +46,44 @@ macro_rules! impl_append_to_builder { return; } - // SAFETY: element_offset points to contiguous, naturally aligned data of - // length num_elements. Alignment is guaranteed by SparkUnsafeArray layout: - // header is 8 + ceil(n/64)*8 bytes (always 8-byte aligned), and elements - // are at element_size stride from an aligned base. + // SAFETY: element_offset points to contiguous element data of length num_elements. debug_assert!(self.element_offset != 0, "element_offset is null"); let ptr = self.element_offset as *const $element_type; - debug_assert!( - (ptr as usize).is_multiple_of(std::mem::align_of::<$element_type>()), - "element_offset is not properly aligned" - ); - let values = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; + let aligned = + (ptr as usize).is_multiple_of(std::mem::align_of::<$element_type>()); if NULLABLE { let null_words = self.null_bitset_ptr(); debug_assert!(!null_words.is_null(), "null_bitset_ptr is null"); - for (idx, &val) in values.iter().enumerate() { - // SAFETY: null_words has ceil(num_elements/64) words, idx < num_elements - if unsafe { Self::is_null_in_bitset(null_words, idx) } { - builder.append_null(); - } else { - builder.append_value(val); + if aligned { + let values = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; + for (idx, &val) in values.iter().enumerate() { + if unsafe { Self::is_null_in_bitset(null_words, idx) } { + builder.append_null(); + } else { + builder.append_value(val); + } + } + } else { + let mut ptr = ptr; + for idx in 0..num_elements { + if unsafe { Self::is_null_in_bitset(null_words, idx) } { + builder.append_null(); + } else { + builder.append_value(unsafe { ptr.read_unaligned() }); + } + ptr = unsafe { ptr.add(1) }; } } - } else { + } else if aligned { + let values = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; builder.append_slice(values); + } else { + let mut ptr = ptr; + for _ in 0..num_elements { + builder.append_value(unsafe { ptr.read_unaligned() }); + ptr = unsafe { ptr.add(1) }; + } } } }; @@ -221,17 +234,12 @@ impl SparkUnsafeArray { return; } - // SAFETY: element_offset points to contiguous, naturally aligned i64 data. debug_assert!( self.element_offset != 0, "append_timestamps: element_offset is null" ); let ptr = self.element_offset as *const i64; - debug_assert!( - (ptr as usize).is_multiple_of(std::mem::align_of::()), - "append_timestamps: element_offset is not properly aligned" - ); - let values = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; + let aligned = (ptr as usize).is_multiple_of(std::mem::align_of::()); if NULLABLE { let null_words = self.null_bitset_ptr(); @@ -239,16 +247,35 @@ impl SparkUnsafeArray { !null_words.is_null(), "append_timestamps: null_bitset_ptr is null" ); - for (idx, &val) in values.iter().enumerate() { - // SAFETY: null_words has ceil(num_elements/64) words, idx < num_elements - if unsafe { Self::is_null_in_bitset(null_words, idx) } { - builder.append_null(); - } else { - builder.append_value(val); + if aligned { + let values = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; + for (idx, &val) in values.iter().enumerate() { + if unsafe { Self::is_null_in_bitset(null_words, idx) } { + builder.append_null(); + } else { + builder.append_value(val); + } + } + } else { + let mut ptr = ptr; + for idx in 0..num_elements { + if unsafe { Self::is_null_in_bitset(null_words, idx) } { + builder.append_null(); + } else { + builder.append_value(unsafe { ptr.read_unaligned() }); + } + ptr = unsafe { ptr.add(1) }; } } - } else { + } else if aligned { + let values = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; builder.append_slice(values); + } else { + let mut ptr = ptr; + for _ in 0..num_elements { + builder.append_value(unsafe { ptr.read_unaligned() }); + ptr = unsafe { ptr.add(1) }; + } } } @@ -262,17 +289,12 @@ impl SparkUnsafeArray { return; } - // SAFETY: element_offset points to contiguous, naturally aligned i32 data. debug_assert!( self.element_offset != 0, "append_dates: element_offset is null" ); let ptr = self.element_offset as *const i32; - debug_assert!( - (ptr as usize).is_multiple_of(std::mem::align_of::()), - "append_dates: element_offset is not properly aligned" - ); - let values = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; + let aligned = (ptr as usize).is_multiple_of(std::mem::align_of::()); if NULLABLE { let null_words = self.null_bitset_ptr(); @@ -280,16 +302,35 @@ impl SparkUnsafeArray { !null_words.is_null(), "append_dates: null_bitset_ptr is null" ); - for (idx, &val) in values.iter().enumerate() { - // SAFETY: null_words has ceil(num_elements/64) words, idx < num_elements - if unsafe { Self::is_null_in_bitset(null_words, idx) } { - builder.append_null(); - } else { - builder.append_value(val); + if aligned { + let values = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; + for (idx, &val) in values.iter().enumerate() { + if unsafe { Self::is_null_in_bitset(null_words, idx) } { + builder.append_null(); + } else { + builder.append_value(val); + } + } + } else { + let mut ptr = ptr; + for idx in 0..num_elements { + if unsafe { Self::is_null_in_bitset(null_words, idx) } { + builder.append_null(); + } else { + builder.append_value(unsafe { ptr.read_unaligned() }); + } + ptr = unsafe { ptr.add(1) }; } } - } else { + } else if aligned { + let values = unsafe { std::slice::from_raw_parts(ptr, num_elements) }; builder.append_slice(values); + } else { + let mut ptr = ptr; + for _ in 0..num_elements { + builder.append_value(unsafe { ptr.read_unaligned() }); + ptr = unsafe { ptr.add(1) }; + } } } }