alamb commented on code in PR #8854:
URL: https://github.com/apache/arrow-rs/pull/8854#discussion_r2541996812
##########
arrow-buffer/src/buffer/immutable.rs:
##########
@@ -115,6 +114,150 @@ impl Buffer {
Self::from(bytes)
}
+ /// Create a new [`Buffer`] by applying the bitwise operation `op` to two
input buffers.
+ ///
+ /// This function is highly optimized for bitwise operations on large
+ /// bitmaps by processing input buffers in chunks of 64 bits (8 bytes) at a
+ /// time, and thus is much faster than applying the operation bit by bit.
+ ///
+ /// # Notes:
+ /// * `op` takes two `u64` inputs and produces one `u64` output,
+ /// operating on 64 bits at a time. **It must only apply bitwise
operations
+ /// on the relevant bits, as the input `u64` may contain irrelevant bits
+ /// and may be processed differently on different endian architectures.**
+ /// * The inputs are treated as bitmaps, meaning that offsets and length
+ /// are specified in number of bits.
+ /// * The output always has zero offset
+ ///
+ /// # See Also
+ /// - [`Buffer::from_bitwise_unary_op`] for unary operations on a single
input buffer.
+ /// - [`apply_bitwise_binary_op`](bit_util::apply_bitwise_binary_op) for
in-place binary bitwise operations
+ ///
+ /// # Example: Create new [`Buffer`] from bitwise `AND` of two [`Buffer`]s
+ /// ```
+ /// # use arrow_buffer::Buffer;
+ /// let left = Buffer::from(&[0b11001100u8, 0b10111010u8]); // 2 bytes =
16 bits
+ /// let right = Buffer::from(&[0b10101010u8, 0b11011100u8, 0b11110000u8]);
// 3 bytes = 24 bits
+ /// // AND of the first 12 bits
+ /// let result = Buffer::from_bitwise_binary_op(
+ /// &left, 0, &right, 0, 12, |a, b| a & b
+ /// );
+ /// assert_eq!(result.as_slice(), &[0b10001000u8, 0b00001000u8]);
+ /// ```
+ ///
+ /// # Example: Create new [`Buffer`] from bitwise `OR` of two byte slices
+ /// ```
+ /// # use arrow_buffer::Buffer;
+ /// let left = [0b11001100u8, 0b10111010u8];
+ /// let right = [0b10101010u8, 0b11011100u8];
+ /// // OR of bits 4..16 from left and bits 0..12 from right
+ /// let result = Buffer::from_bitwise_binary_op(
+ /// &left, 4, &right, 0, 12, |a, b| a | b
+ /// );
+ /// assert_eq!(result.as_slice(), &[0b10101110u8, 0b00001111u8]);
+ /// ```
+ pub fn from_bitwise_binary_op<F>(
+ left: impl AsRef<[u8]>,
+ left_offset_in_bits: usize,
+ right: impl AsRef<[u8]>,
+ right_offset_in_bits: usize,
+ len_in_bits: usize,
+ mut op: F,
+ ) -> Buffer
+ where
+ F: FnMut(u64, u64) -> u64,
+ {
+ let left_chunks = BitChunks::new(left.as_ref(), left_offset_in_bits,
len_in_bits);
+ let right_chunks = BitChunks::new(right.as_ref(),
right_offset_in_bits, len_in_bits);
+
+ let chunks = left_chunks
+ .iter()
+ .zip(right_chunks.iter())
+ .map(|(left, right)| op(left, right));
+ // Soundness: `BitChunks` is a `BitChunks` iterator which
+ // correctly reports its upper bound
+ let mut buffer = unsafe { MutableBuffer::from_trusted_len_iter(chunks)
};
+
+ let remainder_bytes = ceil(left_chunks.remainder_len(), 8);
+ let rem = op(left_chunks.remainder_bits(),
right_chunks.remainder_bits());
+ // we are counting its starting from the least significant bit, to
to_le_bytes should be correct
+ let rem = &rem.to_le_bytes()[0..remainder_bytes];
+ buffer.extend_from_slice(rem);
Review Comment:
I also tried making a version of MutableBuffer::from_trusted_len_iter that
also added additional and it didn't seem to help either (perhaps because the
benchmarks happen to avoid reallocation 🤔 )
```rust
/// Like [`from_trusted_len_iter`] but can add additional capacity at
the end
/// in case the caller wants to add more data after the initial iterator.
#[inline]
pub unsafe fn from_trusted_len_iter_with_additional_capacity<T:
ArrowNativeType, I: Iterator<Item = T>>(
iterator: I,
additional_capacity: usize,
) -> Self {
let item_size = std::mem::size_of::<T>();
let (_, upper) = iterator.size_hint();
let upper = upper.expect("from_trusted_len_iter requires an upper
limit");
let len = upper * item_size;
let mut buffer = MutableBuffer::new(len + additional_capacity);
let mut dst = buffer.data.as_ptr();
for item in iterator {
// note how there is no reserve here (compared with
`extend_from_iter`)
let src = item.to_byte_slice().as_ptr();
unsafe { std::ptr::copy_nonoverlapping(src, dst, item_size) };
dst = unsafe { dst.add(item_size) };
}
assert_eq!(
unsafe { dst.offset_from(buffer.data.as_ptr()) } as usize,
len,
"Trusted iterator length was not accurately reported"
);
buffer.len = len;
buffer
}
```
##########
arrow-buffer/src/buffer/immutable.rs:
##########
@@ -115,6 +114,150 @@ impl Buffer {
Self::from(bytes)
}
+ /// Create a new [`Buffer`] by applying the bitwise operation `op` to two
input buffers.
+ ///
+ /// This function is highly optimized for bitwise operations on large
+ /// bitmaps by processing input buffers in chunks of 64 bits (8 bytes) at a
+ /// time, and thus is much faster than applying the operation bit by bit.
+ ///
+ /// # Notes:
+ /// * `op` takes two `u64` inputs and produces one `u64` output,
+ /// operating on 64 bits at a time. **It must only apply bitwise
operations
+ /// on the relevant bits, as the input `u64` may contain irrelevant bits
+ /// and may be processed differently on different endian architectures.**
+ /// * The inputs are treated as bitmaps, meaning that offsets and length
+ /// are specified in number of bits.
+ /// * The output always has zero offset
+ ///
+ /// # See Also
+ /// - [`Buffer::from_bitwise_unary_op`] for unary operations on a single
input buffer.
+ /// - [`apply_bitwise_binary_op`](bit_util::apply_bitwise_binary_op) for
in-place binary bitwise operations
+ ///
+ /// # Example: Create new [`Buffer`] from bitwise `AND` of two [`Buffer`]s
+ /// ```
+ /// # use arrow_buffer::Buffer;
+ /// let left = Buffer::from(&[0b11001100u8, 0b10111010u8]); // 2 bytes =
16 bits
+ /// let right = Buffer::from(&[0b10101010u8, 0b11011100u8, 0b11110000u8]);
// 3 bytes = 24 bits
+ /// // AND of the first 12 bits
+ /// let result = Buffer::from_bitwise_binary_op(
+ /// &left, 0, &right, 0, 12, |a, b| a & b
+ /// );
+ /// assert_eq!(result.as_slice(), &[0b10001000u8, 0b00001000u8]);
+ /// ```
+ ///
+ /// # Example: Create new [`Buffer`] from bitwise `OR` of two byte slices
+ /// ```
+ /// # use arrow_buffer::Buffer;
+ /// let left = [0b11001100u8, 0b10111010u8];
+ /// let right = [0b10101010u8, 0b11011100u8];
+ /// // OR of bits 4..16 from left and bits 0..12 from right
+ /// let result = Buffer::from_bitwise_binary_op(
+ /// &left, 4, &right, 0, 12, |a, b| a | b
+ /// );
+ /// assert_eq!(result.as_slice(), &[0b10101110u8, 0b00001111u8]);
+ /// ```
+ pub fn from_bitwise_binary_op<F>(
+ left: impl AsRef<[u8]>,
+ left_offset_in_bits: usize,
+ right: impl AsRef<[u8]>,
+ right_offset_in_bits: usize,
+ len_in_bits: usize,
+ mut op: F,
+ ) -> Buffer
+ where
+ F: FnMut(u64, u64) -> u64,
+ {
+ let left_chunks = BitChunks::new(left.as_ref(), left_offset_in_bits,
len_in_bits);
+ let right_chunks = BitChunks::new(right.as_ref(),
right_offset_in_bits, len_in_bits);
+
+ let chunks = left_chunks
+ .iter()
+ .zip(right_chunks.iter())
+ .map(|(left, right)| op(left, right));
+ // Soundness: `BitChunks` is a `BitChunks` iterator which
+ // correctly reports its upper bound
+ let mut buffer = unsafe { MutableBuffer::from_trusted_len_iter(chunks)
};
+
+ let remainder_bytes = ceil(left_chunks.remainder_len(), 8);
+ let rem = op(left_chunks.remainder_bits(),
right_chunks.remainder_bits());
+ // we are counting its starting from the least significant bit, to
to_le_bytes should be correct
+ let rem = &rem.to_le_bytes()[0..remainder_bytes];
+ buffer.extend_from_slice(rem);
Review Comment:
🤔 I tried this
```rust
pub fn from_bitwise_binary_op<F>(
left: impl AsRef<[u8]>,
left_offset_in_bits: usize,
right: impl AsRef<[u8]>,
right_offset_in_bits: usize,
len_in_bits: usize,
mut op: F,
) -> Buffer
where
F: FnMut(u64, u64) -> u64,
{
let left_chunks = BitChunks::new(left.as_ref(), left_offset_in_bits,
len_in_bits);
let right_chunks = BitChunks::new(right.as_ref(),
right_offset_in_bits, len_in_bits);
let remainder_bytes = ceil(left_chunks.remainder_len(), 8);
// if it evenly divides into u64 chunks
let buffer = if remainder_bytes == 0 {
let chunks = left_chunks
.iter()
.zip(right_chunks.iter())
.map(|(left, right)| op(left, right));
// Soundness: `BitChunks` is a `BitChunks` iterator which
// correctly reports its upper bound
unsafe { MutableBuffer::from_trusted_len_iter(chunks) }
} else {
// Compute last u64 here so that we can reserve exact capacity
let rem = op(left_chunks.remainder_bits(),
right_chunks.remainder_bits());
let chunks = left_chunks
.iter()
.zip(right_chunks.iter())
.map(|(left, right)| op(left, right))
.chain(std::iter::once(rem));
// Soundness: `BitChunks` is a `BitChunks` iterator which
// correctly reports its upper bound, and so is the `chain`
iterator
let mut buffer = unsafe {
MutableBuffer::from_trusted_len_iter(chunks) };
// Adjust the length down if last u64 is not fully used
let extra_bytes = 8 - remainder_bytes;
buffer.truncate(buffer.len() - extra_bytes);
buffer
};
buffer.into()
}
```
But it seems to be slower.
--
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]
