sunchao commented on code in PR #4672:
URL: https://github.com/apache/datafusion-comet/pull/4672#discussion_r4103284194
##########
native/shuffle/src/spark_unsafe/list.rs:
##########
@@ -38,38 +47,58 @@ use datafusion_comet_jni_bridge::errors::CometError;
/// - `null_bitset_ptr()` returns a pointer to `ceil(num_elements/64)` i64
words
/// - These invariants are guaranteed by the SparkUnsafeArray layout from the
JVM
macro_rules! impl_append_to_builder {
- ($method_name:ident, $builder_type:ty, $element_type:ty) => {
+ ($method_name:ident, $builder_type:ty, $element_type:ty, $arrow_type:ty)
=> {
pub(crate) fn $method_name<const NULLABLE: bool>(&self, builder: &mut
$builder_type) {
let num_elements = self.num_elements;
if num_elements == 0 {
return;
}
+ // Note: alignment is not guaranteed - that is why do this
+ // This runtime check is needed. Look at `unsafe_object.rs:49` for
more info
+ let ptr = self.element_offset as *const $element_type;
+ let aligned = (ptr as
usize).is_multiple_of(std::mem::align_of::<$element_type>());
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 {
- // 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 {
- builder.append_null();
- } else {
- // SAFETY: ptr is within element data bounds
- builder.append_value(unsafe { ptr.read_unaligned() });
+ if aligned {
+ // Raw values
+ let values = unsafe { std::slice::from_raw_parts(ptr,
num_elements) };
+
+ // Note: in Spark bitmap is padded to 8 byte
word-boundaries
+ // In Arrow we just use the needed number of whole bytes
without padding
+ let null_mask_len = num_elements.div_ceil(8);
+ // Spark and Arrow both use little-endian bit ordering
within each byte,
+ // so casting the i64 null words to *const u8 preserves
correct bit indexing.
+ let null_mask = unsafe {
+ std::slice::from_raw_parts::<u8>(null_words as *const
u8, null_mask_len)
+ };
+ // We need to perform this flip due to the null bitmap
Spark vs Arrow incompatibility
+ // In `Spark` we have 1 set in bitmap meaning that element
IS NULL
+ // In `Arrow` we have 1 set in bitmap meaning that element
IS VALID (non-null)
+ let flipped: Vec<u8> = null_mask.iter().map(|n|
!n).collect();
+ // Constructing null-buffer
+ let validity =
+
NullBuffer::new(BooleanBuffer::new(Buffer::from(flipped), 0, num_elements));
+
+ let arr = PrimitiveArray::<$arrow_type>::new(
+ ScalarBuffer::from(Buffer::from_slice_ref(values)),
Review Comment:
[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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
