alamb commented on code in PR #9297:
URL: https://github.com/apache/arrow-rs/pull/9297#discussion_r2790626465
##########
arrow-buffer/src/buffer/boolean.rs:
##########
@@ -191,67 +189,48 @@ impl BooleanBuffer {
where
F: FnMut(u64) -> u64,
{
- // try fast path for aligned input
- if offset_in_bits & 0x7 == 0 {
- // align to byte boundary
- let aligned = &src.as_ref()[offset_in_bits / 8..];
- if let Some(result) =
- Self::try_from_aligned_bitwise_unary_op(aligned, len_in_bits,
&mut op)
- {
- return result;
+ let end = offset_in_bits + len_in_bits;
+ // Align start and end to 64 bit (8 byte) boundaries if possible to
allow using the
+ // optimized code path as much as possible.
+ let aligned_offset = offset_in_bits & !63;
+ let aligned_end_bytes = bit_util::ceil(end, 64) * 8;
+ let src_len = src.as_ref().len();
+ let slice_end = aligned_end_bytes.min(src_len);
+
+ let aligned_start = &src.as_ref()[aligned_offset / 8..slice_end];
+
+ let (prefix, aligned_u64s, suffix) = unsafe {
aligned_start.as_ref().align_to::<u64>() };
+ match (prefix, suffix) {
+ ([], []) => {
+ // the buffer is word (64 bit) aligned, so use optimized Vec
code.
+ let result_u64s: Vec<u64> = aligned_u64s.iter().map(|l|
op(*l)).collect();
+ return BooleanBuffer::new(result_u64s.into(), offset_in_bits %
64, len_in_bits);
}
- }
-
- let chunks = BitChunks::new(src.as_ref(), offset_in_bits, len_in_bits);
- let mut result = MutableBuffer::with_capacity(chunks.num_u64s() * 8);
- for chunk in chunks.iter() {
- // SAFETY: reserved enough capacity above, (exactly num_u64s()
- // items) and we assume `BitChunks` correctly reports upper bound
- unsafe {
- result.push_unchecked(op(chunk));
+ ([], suffix) => {
+ let suffix = read_u64(suffix);
+ let result_u64s: Vec<u64> = aligned_u64s
+ .iter()
+ .cloned()
+ .chain(std::iter::once(suffix))
+ .map(&mut op)
+ .collect();
+ return BooleanBuffer::new(result_u64s.into(), offset_in_bits %
64, len_in_bits);
}
- }
- if chunks.remainder_len() > 0 {
- debug_assert!(result.capacity() >= result.len() + 8); // should
not reallocate
- // SAFETY: reserved enough capacity above, (exactly num_u64s()
- // items) and we assume `BitChunks` correctly reports upper bound
- unsafe {
- result.push_unchecked(op(chunks.remainder_bits()));
- }
- // Just pushed one u64, which may have trailing zeros
- result.truncate(chunks.num_bytes());
+ _ => {}
}
- BooleanBuffer {
- buffer: Buffer::from(result),
- bit_offset: 0,
- bit_len: len_in_bits,
- }
- }
+ // align to byte boundaries
+ // Use unaligned code path, handle remainder bytes
+ let chunks = aligned_start.chunks_exact(8);
+ let remainder = chunks.remainder();
+ let iter = chunks.map(|c| u64::from_le_bytes(c.try_into().unwrap()));
+ let vec_u64s: Vec<u64> = if remainder.is_empty() {
+ iter.map(&mut op).collect()
+ } else {
+ iter.chain(Some(read_u64(remainder))).map(&mut op).collect()
+ };
- /// Fast path for [`Self::from_bitwise_unary_op`] when input is aligned to
- /// 8-byte (64-bit) boundaries
- ///
- /// Returns None if the fast path cannot be taken
- fn try_from_aligned_bitwise_unary_op<F>(
- src: &[u8],
- len_in_bits: usize,
- op: &mut F,
- ) -> Option<Self>
- where
- F: FnMut(u64) -> u64,
- {
- // Safety: all valid bytes are valid u64s
- let (prefix, aligned_u6us, suffix) = unsafe { src.align_to::<u64>() };
- if !(prefix.is_empty() && suffix.is_empty()) {
- // Couldn't make this case any faster than the default path, see
- // https://github.com/apache/arrow-rs/pull/8996/changes#r2620022082
- return None;
- }
- // the buffer is word (64 bit) aligned, so use optimized Vec code.
- let result_u64s: Vec<u64> = aligned_u6us.iter().map(|l|
op(*l)).collect();
- let buffer = Buffer::from(result_u64s);
- Some(BooleanBuffer::new(buffer, 0, len_in_bits))
+ BooleanBuffer::new(vec_u64s.into(), offset_in_bits % 64, len_in_bits)
Review Comment:
I'll mark this PR as an API change to account for this
--
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]
