From d5a3604a456796411a37d08f9e17036859e680b5 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sun, 16 Nov 2025 08:44:14 -0500 Subject: [PATCH 01/15] Add `Buffer::from_bitwise_unary` and `Buffer::from_bitwise_binary` methods, deprecate old methods --- arrow-arith/src/boolean.rs | 12 +-- arrow-buffer/src/buffer/immutable.rs | 155 +++++++++++++++++++++++++-- arrow-buffer/src/buffer/ops.rs | 64 ++++------- arrow-select/src/nullif.rs | 7 +- 4 files changed, 176 insertions(+), 62 deletions(-) diff --git a/arrow-arith/src/boolean.rs b/arrow-arith/src/boolean.rs index d94df49de256..a5faa2b4a77d 100644 --- a/arrow-arith/src/boolean.rs +++ b/arrow-arith/src/boolean.rs @@ -23,8 +23,8 @@ //! [here](https://doc.rust-lang.org/stable/core/arch/) for more information. use arrow_array::*; -use arrow_buffer::buffer::{bitwise_bin_op_helper, bitwise_quaternary_op_helper}; -use arrow_buffer::{BooleanBuffer, NullBuffer, buffer_bin_and_not}; +use arrow_buffer::buffer::bitwise_quaternary_op_helper; +use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer, buffer_bin_and_not}; use arrow_schema::ArrowError; /// Logical 'and' boolean values with Kleene logic @@ -74,7 +74,7 @@ pub fn and_kleene(left: &BooleanArray, right: &BooleanArray) -> Result Result { // Same as above - Some(bitwise_bin_op_helper( + Some(Buffer::from_bitwise_binary_op( right_null_buffer.buffer(), right_null_buffer.offset(), left_values.inner(), @@ -169,7 +169,7 @@ pub fn or_kleene(left: &BooleanArray, right: &BooleanArray) -> Result Result { // Same as above - Some(bitwise_bin_op_helper( + Some(Buffer::from_bitwise_binary_op( right_nulls.buffer(), right_nulls.offset(), left_values.inner(), diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs index 55a4621540c8..188753f161f3 100644 --- a/arrow-buffer/src/buffer/immutable.rs +++ b/arrow-buffer/src/buffer/immutable.rs @@ -22,13 +22,12 @@ use std::sync::Arc; use crate::BufferBuilder; use crate::alloc::{Allocation, Deallocation}; -use crate::util::bit_chunk_iterator::{BitChunks, UnalignedBitChunk}; -use crate::{bit_util, bytes::Bytes, native::ArrowNativeType}; - +use crate::bit_util::ceil; #[cfg(feature = "pool")] use crate::pool::MemoryPool; +use crate::util::bit_chunk_iterator::{BitChunks, UnalignedBitChunk}; +use crate::{bit_util, bytes::Bytes, native::ArrowNativeType}; -use super::ops::bitwise_unary_op_helper; use super::{MutableBuffer, ScalarBuffer}; /// A contiguous memory region that can be shared with other buffers and across @@ -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( + 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); + + buffer.into() + } + + /// Create a new [`Buffer`] by applying the bitwise operation to `op` to an input buffer. + /// + /// 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_binary_op`] for binary operations on a single input buffer. + /// - [`apply_bitwise_unary_op`](bit_util::apply_bitwise_unary_op) for in-place unary bitwise operations + /// + /// # Example: Create new [`Buffer`] from bitwise `NOT` of an input [`Buffer`] + /// ``` + /// # use arrow_buffer::Buffer; + /// let input = Buffer::from(&[0b11001100u8, 0b10111010u8]); // 2 bytes = 16 bits + /// // NOT of the first 12 bits + /// let result = Buffer::from_bitwise_unary_op( + /// &input, 0, 12, |a| !a + /// ); + /// assert_eq!(result.as_slice(), &[0b00110011u8, 0b11110101u8]); + /// ``` + /// + /// # Example: Create a new [`Buffer`] copying a bit slice from in input slice + /// ``` + /// # use arrow_buffer::Buffer; + /// let input = [0b11001100u8, 0b10111010u8]; + /// // // Copy bits 4..16 from input + /// let result = Buffer::from_bitwise_unary_op( + /// &input, 4, 12, |a| a + /// ); + /// assert_eq!(result.as_slice(), &[0b10101100u8, 0b00001011u8], "[{:08b}, {:08b}]", result.as_slice()[0], result.as_slice()[1]); + pub fn from_bitwise_unary_op( + left: impl AsRef<[u8]>, + offset_in_bits: usize, + len_in_bits: usize, + mut op: F, + ) -> Buffer + where + F: FnMut(u64) -> u64, + { + // reserve capacity and set length so we can get a typed view of u64 chunks + let mut result = + MutableBuffer::new(ceil(len_in_bits, 8)).with_bitset(len_in_bits / 64 * 8, false); + + let left_chunks = BitChunks::new(left.as_ref(), offset_in_bits, len_in_bits); + + let result_chunks = result.typed_data_mut::().iter_mut(); + + result_chunks + .zip(left_chunks.iter()) + .for_each(|(res, left)| { + *res = op(left); + }); + + let remainder_bytes = ceil(left_chunks.remainder_len(), 8); + let rem = op(left_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]; + result.extend_from_slice(rem); + + result.into() + } + /// Returns the offset, in bytes, of `Self::ptr` to `Self::data` /// /// self.ptr and self.data can be different after slicing or advancing the buffer. @@ -344,10 +487,10 @@ impl Buffer { return self.slice_with_length(offset / 8, bit_util::ceil(len, 8)); } - bitwise_unary_op_helper(self, offset, len, |a| a) + Self::from_bitwise_unary_op(self, offset, len, |a| a) } - /// Returns a `BitChunks` instance which can be used to iterate over this buffers bits + /// Returns a `BitChunks` instance which can be used to iterate over this buffer's bits /// in larger chunks and starting at arbitrary bit offsets. /// Note that both `offset` and `length` are measured in bits. pub fn bit_chunks(&self, offset: usize, len: usize) -> BitChunks<'_> { diff --git a/arrow-buffer/src/buffer/ops.rs b/arrow-buffer/src/buffer/ops.rs index c69e5c6deb10..a3e7e06aa56f 100644 --- a/arrow-buffer/src/buffer/ops.rs +++ b/arrow-buffer/src/buffer/ops.rs @@ -60,69 +60,41 @@ where /// Apply a bitwise operation `op` to two inputs and return the result as a Buffer. /// The inputs are treated as bitmaps, meaning that offsets and length are specified in number of bits. +#[deprecated(since = "57.1.0", note = "use Buffer::from_bitwise_binary_op instead")] pub fn bitwise_bin_op_helper( left: &Buffer, left_offset_in_bits: usize, right: &Buffer, right_offset_in_bits: usize, len_in_bits: usize, - mut op: F, + op: F, ) -> Buffer where F: FnMut(u64, u64) -> u64, { - let left_chunks = left.bit_chunks(left_offset_in_bits, len_in_bits); - let right_chunks = right.bit_chunks(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); - - buffer.into() + Buffer::from_bitwise_binary_op( + left, + left_offset_in_bits, + right, + right_offset_in_bits, + len_in_bits, + op, + ) } /// Apply a bitwise operation `op` to one input and return the result as a Buffer. /// The input is treated as a bitmap, meaning that offset and length are specified in number of bits. +#[deprecated(since = "57.1.0", note = "use Buffer::from_bitwise_unary_op instead")] pub fn bitwise_unary_op_helper( left: &Buffer, offset_in_bits: usize, len_in_bits: usize, - mut op: F, + op: F, ) -> Buffer where F: FnMut(u64) -> u64, { - // reserve capacity and set length so we can get a typed view of u64 chunks - let mut result = - MutableBuffer::new(ceil(len_in_bits, 8)).with_bitset(len_in_bits / 64 * 8, false); - - let left_chunks = left.bit_chunks(offset_in_bits, len_in_bits); - - let result_chunks = result.typed_data_mut::().iter_mut(); - - result_chunks - .zip(left_chunks.iter()) - .for_each(|(res, left)| { - *res = op(left); - }); - - let remainder_bytes = ceil(left_chunks.remainder_len(), 8); - let rem = op(left_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]; - result.extend_from_slice(rem); - - result.into() + Buffer::from_bitwise_unary_op(left, offset_in_bits, len_in_bits, op) } /// Apply a bitwise and to two inputs and return the result as a Buffer. @@ -134,7 +106,7 @@ pub fn buffer_bin_and( right_offset_in_bits: usize, len_in_bits: usize, ) -> Buffer { - bitwise_bin_op_helper( + Buffer::from_bitwise_binary_op( left, left_offset_in_bits, right, @@ -153,7 +125,7 @@ pub fn buffer_bin_or( right_offset_in_bits: usize, len_in_bits: usize, ) -> Buffer { - bitwise_bin_op_helper( + Buffer::from_bitwise_binary_op( left, left_offset_in_bits, right, @@ -172,7 +144,7 @@ pub fn buffer_bin_xor( right_offset_in_bits: usize, len_in_bits: usize, ) -> Buffer { - bitwise_bin_op_helper( + Buffer::from_bitwise_binary_op( left, left_offset_in_bits, right, @@ -191,7 +163,7 @@ pub fn buffer_bin_and_not( right_offset_in_bits: usize, len_in_bits: usize, ) -> Buffer { - bitwise_bin_op_helper( + Buffer::from_bitwise_binary_op( left, left_offset_in_bits, right, @@ -204,5 +176,5 @@ pub fn buffer_bin_and_not( /// Apply a bitwise not to one input and return the result as a Buffer. /// The input is treated as a bitmap, meaning that offset and length are specified in number of bits. pub fn buffer_unary_not(left: &Buffer, offset_in_bits: usize, len_in_bits: usize) -> Buffer { - bitwise_unary_op_helper(left, offset_in_bits, len_in_bits, |a| !a) + Buffer::from_bitwise_unary_op(left, offset_in_bits, len_in_bits, |a| !a) } diff --git a/arrow-select/src/nullif.rs b/arrow-select/src/nullif.rs index 8e3cc7d56c71..8d97457cf7be 100644 --- a/arrow-select/src/nullif.rs +++ b/arrow-select/src/nullif.rs @@ -18,8 +18,7 @@ //! Implements the `nullif` function for Arrow arrays. use arrow_array::{Array, ArrayRef, BooleanArray, make_array}; -use arrow_buffer::buffer::{bitwise_bin_op_helper, bitwise_unary_op_helper}; -use arrow_buffer::{BooleanBuffer, NullBuffer}; +use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer}; use arrow_schema::{ArrowError, DataType}; /// Returns a new array with the same values and the validity bit to false where @@ -75,7 +74,7 @@ pub fn nullif(left: &dyn Array, right: &BooleanArray) -> Result { let mut valid_count = 0; - let b = bitwise_bin_op_helper( + let b = Buffer::from_bitwise_binary_op( left.buffer(), left.offset(), right.inner(), @@ -91,7 +90,7 @@ pub fn nullif(left: &dyn Array, right: &BooleanArray) -> Result { let mut null_count = 0; - let buffer = bitwise_unary_op_helper(right.inner(), right.offset(), len, |b| { + let buffer = Buffer::from_bitwise_unary_op(right.inner(), right.offset(), len, |b| { let t = !b; null_count += t.count_zeros() as usize; t From af31aebf77f221ff112937e736ce87d6befdf434 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sat, 13 Dec 2025 09:57:06 -0500 Subject: [PATCH 02/15] Reduce allocations --- arrow-buffer/src/buffer/immutable.rs | 40 +++++++++++---------- arrow-buffer/src/util/bit_chunk_iterator.rs | 2 +- 2 files changed, 23 insertions(+), 19 deletions(-) diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs index c01345c92cf4..7d8a30e4fa89 100644 --- a/arrow-buffer/src/buffer/immutable.rs +++ b/arrow-buffer/src/buffer/immutable.rs @@ -116,9 +116,8 @@ impl Buffer { /// 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. + /// This function is much faster than applying the operation bit by bit as + /// it processes input buffers in chunks of 64 bits (8 bytes) at a time /// /// # Notes: /// * `op` takes two `u64` inputs and produces one `u64` output, @@ -167,31 +166,36 @@ impl Buffer { where F: FnMut(u64, u64) -> u64, { + // each chunk is 64 bits 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 + let num_u64s = if left_chunks.remainder_len() > 0 { + left_chunks.chunk_len() + 1 + } else { + left_chunks.chunk_len() + }; + + let mut result = Vec::with_capacity(num_u64s * size_of::()); + result.extend(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); + .map(|(left, right)| op(left, right)) + ); + if left_chunks.remainder_len() > 0 { + debug_assert_eq!(result.capacity(), result.len() + 8); + result.push(op(left_chunks.remainder_bits(), right_chunks.remainder_bits())); + } - buffer.into() + // Note the result is a vec of u64, so it may have more bytes than the remainder, + // so we need to slice it down to the correct length + Buffer::from(result).slice_with_length(0, ceil(len_in_bits, 8)) } /// Create a new [`Buffer`] by applying the bitwise operation to `op` to an input buffer. /// - /// 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. + /// This function is much faster than applying the operation bit by bit as + /// it processes input buffers in chunks of 64 bits (8 bytes) at a time /// /// # Notes: /// * `op` takes two `u64` inputs and produces one `u64` output, diff --git a/arrow-buffer/src/util/bit_chunk_iterator.rs b/arrow-buffer/src/util/bit_chunk_iterator.rs index e11383f6f3db..1a9899a3f829 100644 --- a/arrow-buffer/src/util/bit_chunk_iterator.rs +++ b/arrow-buffer/src/util/bit_chunk_iterator.rs @@ -259,7 +259,7 @@ impl<'a> BitChunks<'a> { self.remainder_len } - /// Returns the number of chunks + /// Returns the number of `u64` chunks #[inline] pub const fn chunk_len(&self) -> usize { self.chunk_len From 70b113297a78eae2b7c8372f1e45c451a66baaa9 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sat, 13 Dec 2025 10:07:59 -0500 Subject: [PATCH 03/15] Reduce allocations --- arrow-buffer/src/buffer/immutable.rs | 59 +++++++++------------ arrow-buffer/src/util/bit_chunk_iterator.rs | 20 +++++++ 2 files changed, 44 insertions(+), 35 deletions(-) diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs index 7d8a30e4fa89..56c03cf3fe61 100644 --- a/arrow-buffer/src/buffer/immutable.rs +++ b/arrow-buffer/src/buffer/immutable.rs @@ -22,7 +22,6 @@ use std::sync::Arc; use crate::BufferBuilder; use crate::alloc::{Allocation, Deallocation}; -use crate::bit_util::ceil; #[cfg(feature = "pool")] use crate::pool::MemoryPool; use crate::util::bit_chunk_iterator::{BitChunks, UnalignedBitChunk}; @@ -170,26 +169,24 @@ impl Buffer { 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 num_u64s = if left_chunks.remainder_len() > 0 { - left_chunks.chunk_len() + 1 - } else { - left_chunks.chunk_len() - }; - - let mut result = Vec::with_capacity(num_u64s * size_of::()); - result.extend(left_chunks - .iter() - .zip(right_chunks.iter()) - .map(|(left, right)| op(left, right)) + let mut result = Vec::with_capacity(left_chunks.num_u64s() * 8); + result.extend( + left_chunks + .iter() + .zip(right_chunks.iter()) + .map(|(left, right)| op(left, right)), ); if left_chunks.remainder_len() > 0 { - debug_assert_eq!(result.capacity(), result.len() + 8); - result.push(op(left_chunks.remainder_bits(), right_chunks.remainder_bits())); + debug_assert_eq!(result.capacity(), result.len() + 8); // should not reallocate + result.push(op( + left_chunks.remainder_bits(), + right_chunks.remainder_bits(), + )); } - // Note the result is a vec of u64, so it may have more bytes than the remainder, - // so we need to slice it down to the correct length - Buffer::from(result).slice_with_length(0, ceil(len_in_bits, 8)) + // Result a vec of u64, so it may have trailing zero bytes, so we need + // to slice it down to the correct length + Buffer::from(result).slice_with_length(0, left_chunks.num_bytes()) } /// Create a new [`Buffer`] by applying the bitwise operation to `op` to an input buffer. @@ -239,27 +236,19 @@ impl Buffer { where F: FnMut(u64) -> u64, { - // reserve capacity and set length so we can get a typed view of u64 chunks - let mut result = - MutableBuffer::new(ceil(len_in_bits, 8)).with_bitset(len_in_bits / 64 * 8, false); - + // each chunk is 64 bits let left_chunks = BitChunks::new(left.as_ref(), offset_in_bits, len_in_bits); - let result_chunks = result.typed_data_mut::().iter_mut(); - - result_chunks - .zip(left_chunks.iter()) - .for_each(|(res, left)| { - *res = op(left); - }); - - let remainder_bytes = ceil(left_chunks.remainder_len(), 8); - let rem = op(left_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]; - result.extend_from_slice(rem); + let mut result = Vec::with_capacity(left_chunks.num_u64s() * 8); + result.extend(left_chunks.iter().map(&mut op)); + if left_chunks.remainder_len() > 0 { + debug_assert_eq!(result.capacity(), result.len() + 8); // should not reallocate + result.push(op(left_chunks.remainder_bits())); + } - result.into() + // Result a vec of u64, so it may have trailing zero bytes, so we need + // to slice it down to the correct length + Buffer::from(result).slice_with_length(0, left_chunks.num_bytes()) } /// Returns the offset, in bytes, of `Self::ptr` to `Self::data` diff --git a/arrow-buffer/src/util/bit_chunk_iterator.rs b/arrow-buffer/src/util/bit_chunk_iterator.rs index 1a9899a3f829..d2028cd20048 100644 --- a/arrow-buffer/src/util/bit_chunk_iterator.rs +++ b/arrow-buffer/src/util/bit_chunk_iterator.rs @@ -265,6 +265,26 @@ impl<'a> BitChunks<'a> { self.chunk_len } + /// Return the number of `u64` that are needed to represent all bits + /// (including remainder) + /// + /// This is the size of a + #[inline] + pub fn num_u64s(&self) -> usize { + if self.remainder_len == 0 { + self.chunk_len + } else { + self.chunk_len + 1 + } + } + + /// Return the number of bytes that are needed to represent all bits + /// (including remainder) + #[inline] + pub fn num_bytes(&self) -> usize { + ceil(self.chunk_len * 64 + self.remainder_len, 8) + } + /// Returns the bitmask of remaining bits #[inline] pub fn remainder_bits(&self) -> u64 { From cf17667e5ea09b76167433dc7a22e4ed61db7014 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sat, 13 Dec 2025 10:21:37 -0500 Subject: [PATCH 04/15] fix --- arrow-buffer/src/buffer/immutable.rs | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs index 56c03cf3fe61..fc0b2f7607db 100644 --- a/arrow-buffer/src/buffer/immutable.rs +++ b/arrow-buffer/src/buffer/immutable.rs @@ -20,11 +20,11 @@ use std::fmt::Debug; use std::ptr::NonNull; use std::sync::Arc; -use crate::BufferBuilder; use crate::alloc::{Allocation, Deallocation}; #[cfg(feature = "pool")] use crate::pool::MemoryPool; use crate::util::bit_chunk_iterator::{BitChunks, UnalignedBitChunk}; +use crate::BufferBuilder; use crate::{bit_util, bytes::Bytes, native::ArrowNativeType}; use super::{MutableBuffer, ScalarBuffer}; @@ -169,7 +169,7 @@ impl Buffer { 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 mut result = Vec::with_capacity(left_chunks.num_u64s() * 8); + let mut result = Vec::with_capacity(left_chunks.num_u64s()); result.extend( left_chunks .iter() @@ -177,7 +177,7 @@ impl Buffer { .map(|(left, right)| op(left, right)), ); if left_chunks.remainder_len() > 0 { - debug_assert_eq!(result.capacity(), result.len() + 8); // should not reallocate + debug_assert_eq!(result.capacity(), result.len() + 1); // should not reallocate result.push(op( left_chunks.remainder_bits(), right_chunks.remainder_bits(), @@ -238,11 +238,10 @@ impl Buffer { { // each chunk is 64 bits let left_chunks = BitChunks::new(left.as_ref(), offset_in_bits, len_in_bits); - - let mut result = Vec::with_capacity(left_chunks.num_u64s() * 8); + let mut result = Vec::with_capacity(left_chunks.num_u64s()); result.extend(left_chunks.iter().map(&mut op)); if left_chunks.remainder_len() > 0 { - debug_assert_eq!(result.capacity(), result.len() + 8); // should not reallocate + debug_assert_eq!(result.capacity(), result.len() + 1); // should not reallocate result.push(op(left_chunks.remainder_bits())); } From 819210e0881353e2a9f37f98e4c4138289ee254e Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sat, 13 Dec 2025 10:21:55 -0500 Subject: [PATCH 05/15] fix --- arrow-buffer/src/buffer/immutable.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs index fc0b2f7607db..46d4cdd64612 100644 --- a/arrow-buffer/src/buffer/immutable.rs +++ b/arrow-buffer/src/buffer/immutable.rs @@ -20,11 +20,11 @@ use std::fmt::Debug; use std::ptr::NonNull; use std::sync::Arc; +use crate::BufferBuilder; use crate::alloc::{Allocation, Deallocation}; #[cfg(feature = "pool")] use crate::pool::MemoryPool; use crate::util::bit_chunk_iterator::{BitChunks, UnalignedBitChunk}; -use crate::BufferBuilder; use crate::{bit_util, bytes::Bytes, native::ArrowNativeType}; use super::{MutableBuffer, ScalarBuffer}; From c6a2e40b3a77b4fc1f43df9319e52531857c6ea4 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sun, 14 Dec 2025 10:02:33 -0500 Subject: [PATCH 06/15] sprinkle unsafe --- arrow-buffer/src/buffer/immutable.rs | 41 ++++++++++++++++------------ 1 file changed, 24 insertions(+), 17 deletions(-) diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs index 46d4cdd64612..f0a02ce6186d 100644 --- a/arrow-buffer/src/buffer/immutable.rs +++ b/arrow-buffer/src/buffer/immutable.rs @@ -169,24 +169,26 @@ impl Buffer { 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 mut result = Vec::with_capacity(left_chunks.num_u64s()); - result.extend( - left_chunks - .iter() - .zip(right_chunks.iter()) - .map(|(left, right)| op(left, right)), - ); + let mut result = MutableBuffer::with_capacity(left_chunks.num_u64s() * 8); + + for (left, right) in left_chunks.iter().zip(right_chunks.iter()) { + // SAFETY: we have reserved enough capacity above + unsafe { + result.push_unchecked(op(left, right)); + } + } if left_chunks.remainder_len() > 0 { - debug_assert_eq!(result.capacity(), result.len() + 1); // should not reallocate + debug_assert!(result.capacity() > result.len() + 8); // should not reallocate result.push(op( left_chunks.remainder_bits(), right_chunks.remainder_bits(), )); + // Just pushed one u64, which may have have trailing zeros, + // so truncate back to the correct length + result.truncate(left_chunks.num_bytes()); } - // Result a vec of u64, so it may have trailing zero bytes, so we need - // to slice it down to the correct length - Buffer::from(result).slice_with_length(0, left_chunks.num_bytes()) + Buffer::from(result) } /// Create a new [`Buffer`] by applying the bitwise operation to `op` to an input buffer. @@ -238,16 +240,21 @@ impl Buffer { { // each chunk is 64 bits let left_chunks = BitChunks::new(left.as_ref(), offset_in_bits, len_in_bits); - let mut result = Vec::with_capacity(left_chunks.num_u64s()); - result.extend(left_chunks.iter().map(&mut op)); + let mut result = MutableBuffer::with_capacity(left_chunks.num_u64s() * 8); + for left in left_chunks.iter() { + // SAFETY: we have reserved enough capacity above + unsafe { + result.push_unchecked(op(left)); + } + } if left_chunks.remainder_len() > 0 { - debug_assert_eq!(result.capacity(), result.len() + 1); // should not reallocate + debug_assert!(result.capacity() > result.len() + 8); // should not reallocate result.push(op(left_chunks.remainder_bits())); + // Just pushed one u64, which may have have trailing zeros, + result.truncate(left_chunks.num_bytes()); } - // Result a vec of u64, so it may have trailing zero bytes, so we need - // to slice it down to the correct length - Buffer::from(result).slice_with_length(0, left_chunks.num_bytes()) + Buffer::from(result) } /// Returns the offset, in bytes, of `Self::ptr` to `Self::data` From 1592fcfde2705c7c7768b70b1004b426cc610012 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sun, 14 Dec 2025 15:03:35 -0500 Subject: [PATCH 07/15] clenaup --- arrow-arith/src/boolean.rs | 25 ++-- arrow-buffer/src/buffer/boolean.rs | 163 ++++++++++++++++++++++++++- arrow-buffer/src/buffer/immutable.rs | 147 +----------------------- arrow-buffer/src/buffer/ops.rs | 30 ++--- arrow-select/src/nullif.rs | 7 +- 5 files changed, 197 insertions(+), 175 deletions(-) diff --git a/arrow-arith/src/boolean.rs b/arrow-arith/src/boolean.rs index a5faa2b4a77d..4898fea3d249 100644 --- a/arrow-arith/src/boolean.rs +++ b/arrow-arith/src/boolean.rs @@ -24,7 +24,7 @@ use arrow_array::*; use arrow_buffer::buffer::bitwise_quaternary_op_helper; -use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer, buffer_bin_and_not}; +use arrow_buffer::{BooleanBuffer, NullBuffer, buffer_bin_and_not}; use arrow_schema::ArrowError; /// Logical 'and' boolean values with Kleene logic @@ -74,7 +74,7 @@ pub fn and_kleene(left: &BooleanArray, right: &BooleanArray) -> Result Result { // Same as above - Some(Buffer::from_bitwise_binary_op( + Some(BooleanBuffer::from_bitwise_binary_op( right_null_buffer.buffer(), right_null_buffer.offset(), left_values.inner(), @@ -100,7 +100,7 @@ pub fn and_kleene(left: &BooleanArray, right: &BooleanArray) -> Result Result Result Result { // Same as above - Some(Buffer::from_bitwise_binary_op( + Some(BooleanBuffer::from_bitwise_binary_op( right_nulls.buffer(), right_nulls.offset(), left_values.inner(), @@ -195,7 +197,7 @@ pub fn or_kleene(left: &BooleanArray, right: &BooleanArray) -> Result Result( + 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, + ) -> Self + where + F: FnMut(u64, u64) -> u64, + { + // each chunk is 64 bits + 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 mut result = MutableBuffer::with_capacity(left_chunks.num_u64s() * 8); + + for (left, right) in left_chunks.iter().zip(right_chunks.iter()) { + // SAFETY: we have reserved enough capacity above, and we are + // pushing exactly num_u64s() items and `BitChunks` correctly + // reports its upper bound + unsafe { + result.push_unchecked(op(left, right)); + } + } + if left_chunks.remainder_len() > 0 { + debug_assert!(result.capacity() > result.len() + 8); // should not reallocate + result.push(op( + left_chunks.remainder_bits(), + right_chunks.remainder_bits(), + )); + // Just pushed one u64, which may have trailing zeros, + // so truncate back to the correct length + result.truncate(left_chunks.num_bytes()); + } + + BooleanBuffer { + buffer: Buffer::from(result), + offset: 0, + len: len_in_bits, + } + } + + /// Create a new [`Buffer`] by applying the bitwise operation to `op` to an input buffer. + /// + /// This function is much faster than applying the operation bit by bit as + /// it processes input buffers in chunks of 64 bits (8 bytes) at a time + /// + /// # 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_binary_op`] for binary operations on a single input buffer. + /// - [`apply_bitwise_unary_op`](bit_util::apply_bitwise_unary_op) for in-place unary bitwise operations + /// + /// # Example: Create new [`Buffer`] from bitwise `NOT` of an input [`Buffer`] + /// ``` + /// # use arrow_buffer::BooleanBuffer; + /// let input = [0b11001100u8, 0b10111010u8]; // 2 bytes = 16 bits + /// // NOT of the first 12 bits + /// let result = BooleanBuffer::from_bitwise_unary_op( + /// &input, 0, 12, |a| !a + /// ); + /// assert_eq!(result.inner().as_slice(), &[0b00110011u8, 0b11110101u8]); + /// ``` + /// + /// # Example: Create a new [`Buffer`] copying a bit slice from in input slice + /// ``` + /// # use arrow_buffer::BooleanBuffer; + /// let input = [0b11001100u8, 0b10111010u8]; + /// // // Copy bits 4..16 from input + /// let result = BooleanBuffer::from_bitwise_unary_op( + /// &input, 4, 12, |a| a + /// ); + /// assert_eq!(result.inner().as_slice(), &[0b10101100u8, 0b00001011u8]); + pub fn from_bitwise_unary_op( + left: impl AsRef<[u8]>, + offset_in_bits: usize, + len_in_bits: usize, + mut op: F, + ) -> Self + where + F: FnMut(u64) -> u64, + { + // each chunk is 64 bits + let left_chunks = BitChunks::new(left.as_ref(), offset_in_bits, len_in_bits); + let mut result = MutableBuffer::with_capacity(left_chunks.num_u64s() * 8); + for left in left_chunks.iter() { + // SAFETY: we have reserved enough capacity above, and we are + // pushing exactly num_u64s() items and `BitChunks` correctly + // reports its upper bound + unsafe { + result.push_unchecked(op(left)); + } + } + if left_chunks.remainder_len() > 0 { + debug_assert!(result.capacity() > result.len() + 8); // should not reallocate + result.push(op(left_chunks.remainder_bits())); + // Just pushed one u64, which may have have trailing zeros, + result.truncate(left_chunks.num_bytes()); + } + + BooleanBuffer { + buffer: Buffer::from(result), + offset: 0, + len: len_in_bits, + } + } + /// Invokes `f` with indexes `0..len` collecting the boolean results into a new `BooleanBuffer` pub fn collect_bool bool>(len: usize, f: F) -> Self { let buffer = MutableBuffer::collect_bool(len, f); @@ -188,6 +347,8 @@ impl BooleanBuffer { } /// Returns the inner [`Buffer`] + /// + /// Note this does not account for offset and length of this [`BooleanBuffer`] #[inline] pub fn inner(&self) -> &Buffer { &self.buffer diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs index f0a02ce6186d..3fa0750a4f93 100644 --- a/arrow-buffer/src/buffer/immutable.rs +++ b/arrow-buffer/src/buffer/immutable.rs @@ -20,7 +20,7 @@ use std::fmt::Debug; use std::ptr::NonNull; use std::sync::Arc; -use crate::BufferBuilder; +use crate::{BooleanBuffer, BufferBuilder}; use crate::alloc::{Allocation, Deallocation}; #[cfg(feature = "pool")] use crate::pool::MemoryPool; @@ -113,149 +113,6 @@ impl Buffer { Self::from(bytes) } - /// Create a new [`Buffer`] by applying the bitwise operation `op` to two input buffers. - /// - /// This function is much faster than applying the operation bit by bit as - /// it processes input buffers in chunks of 64 bits (8 bytes) at a time - /// - /// # 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( - 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, - { - // each chunk is 64 bits - 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 mut result = MutableBuffer::with_capacity(left_chunks.num_u64s() * 8); - - for (left, right) in left_chunks.iter().zip(right_chunks.iter()) { - // SAFETY: we have reserved enough capacity above - unsafe { - result.push_unchecked(op(left, right)); - } - } - if left_chunks.remainder_len() > 0 { - debug_assert!(result.capacity() > result.len() + 8); // should not reallocate - result.push(op( - left_chunks.remainder_bits(), - right_chunks.remainder_bits(), - )); - // Just pushed one u64, which may have have trailing zeros, - // so truncate back to the correct length - result.truncate(left_chunks.num_bytes()); - } - - Buffer::from(result) - } - - /// Create a new [`Buffer`] by applying the bitwise operation to `op` to an input buffer. - /// - /// This function is much faster than applying the operation bit by bit as - /// it processes input buffers in chunks of 64 bits (8 bytes) at a time - /// - /// # 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_binary_op`] for binary operations on a single input buffer. - /// - [`apply_bitwise_unary_op`](bit_util::apply_bitwise_unary_op) for in-place unary bitwise operations - /// - /// # Example: Create new [`Buffer`] from bitwise `NOT` of an input [`Buffer`] - /// ``` - /// # use arrow_buffer::Buffer; - /// let input = Buffer::from(&[0b11001100u8, 0b10111010u8]); // 2 bytes = 16 bits - /// // NOT of the first 12 bits - /// let result = Buffer::from_bitwise_unary_op( - /// &input, 0, 12, |a| !a - /// ); - /// assert_eq!(result.as_slice(), &[0b00110011u8, 0b11110101u8]); - /// ``` - /// - /// # Example: Create a new [`Buffer`] copying a bit slice from in input slice - /// ``` - /// # use arrow_buffer::Buffer; - /// let input = [0b11001100u8, 0b10111010u8]; - /// // // Copy bits 4..16 from input - /// let result = Buffer::from_bitwise_unary_op( - /// &input, 4, 12, |a| a - /// ); - /// assert_eq!(result.as_slice(), &[0b10101100u8, 0b00001011u8], "[{:08b}, {:08b}]", result.as_slice()[0], result.as_slice()[1]); - pub fn from_bitwise_unary_op( - left: impl AsRef<[u8]>, - offset_in_bits: usize, - len_in_bits: usize, - mut op: F, - ) -> Buffer - where - F: FnMut(u64) -> u64, - { - // each chunk is 64 bits - let left_chunks = BitChunks::new(left.as_ref(), offset_in_bits, len_in_bits); - let mut result = MutableBuffer::with_capacity(left_chunks.num_u64s() * 8); - for left in left_chunks.iter() { - // SAFETY: we have reserved enough capacity above - unsafe { - result.push_unchecked(op(left)); - } - } - if left_chunks.remainder_len() > 0 { - debug_assert!(result.capacity() > result.len() + 8); // should not reallocate - result.push(op(left_chunks.remainder_bits())); - // Just pushed one u64, which may have have trailing zeros, - result.truncate(left_chunks.num_bytes()); - } - - Buffer::from(result) - } /// Returns the offset, in bytes, of `Self::ptr` to `Self::data` /// @@ -486,7 +343,7 @@ impl Buffer { return self.slice_with_length(offset / 8, bit_util::ceil(len, 8)); } - Self::from_bitwise_unary_op(self, offset, len, |a| a) + BooleanBuffer::from_bitwise_unary_op(self, offset, len, |a| a).into_inner() } /// Returns a `BitChunks` instance which can be used to iterate over this buffer's bits diff --git a/arrow-buffer/src/buffer/ops.rs b/arrow-buffer/src/buffer/ops.rs index a3e7e06aa56f..a7d26d321ad7 100644 --- a/arrow-buffer/src/buffer/ops.rs +++ b/arrow-buffer/src/buffer/ops.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use crate::BooleanBuffer; use super::{Buffer, MutableBuffer}; use crate::util::bit_util::ceil; @@ -60,7 +61,7 @@ where /// Apply a bitwise operation `op` to two inputs and return the result as a Buffer. /// The inputs are treated as bitmaps, meaning that offsets and length are specified in number of bits. -#[deprecated(since = "57.1.0", note = "use Buffer::from_bitwise_binary_op instead")] +#[deprecated(since = "57.1.0", note = "use BooleanBuffer::from_bitwise_binary_op instead")] pub fn bitwise_bin_op_helper( left: &Buffer, left_offset_in_bits: usize, @@ -72,19 +73,19 @@ pub fn bitwise_bin_op_helper( where F: FnMut(u64, u64) -> u64, { - Buffer::from_bitwise_binary_op( + BooleanBuffer::from_bitwise_binary_op( left, left_offset_in_bits, right, right_offset_in_bits, len_in_bits, op, - ) + ).into_inner() } /// Apply a bitwise operation `op` to one input and return the result as a Buffer. /// The input is treated as a bitmap, meaning that offset and length are specified in number of bits. -#[deprecated(since = "57.1.0", note = "use Buffer::from_bitwise_unary_op instead")] +#[deprecated(since = "57.1.0", note = "use BooleanBuffer::from_bitwise_unary_op instead")] pub fn bitwise_unary_op_helper( left: &Buffer, offset_in_bits: usize, @@ -94,7 +95,8 @@ pub fn bitwise_unary_op_helper( where F: FnMut(u64) -> u64, { - Buffer::from_bitwise_unary_op(left, offset_in_bits, len_in_bits, op) + BooleanBuffer::from_bitwise_unary_op(left, offset_in_bits, len_in_bits, op).into_inner() + } /// Apply a bitwise and to two inputs and return the result as a Buffer. @@ -106,14 +108,14 @@ pub fn buffer_bin_and( right_offset_in_bits: usize, len_in_bits: usize, ) -> Buffer { - Buffer::from_bitwise_binary_op( + BooleanBuffer::from_bitwise_binary_op( left, left_offset_in_bits, right, right_offset_in_bits, len_in_bits, |a, b| a & b, - ) + ).into_inner() } /// Apply a bitwise or to two inputs and return the result as a Buffer. @@ -125,14 +127,14 @@ pub fn buffer_bin_or( right_offset_in_bits: usize, len_in_bits: usize, ) -> Buffer { - Buffer::from_bitwise_binary_op( + BooleanBuffer::from_bitwise_binary_op( left, left_offset_in_bits, right, right_offset_in_bits, len_in_bits, |a, b| a | b, - ) + ).into_inner() } /// Apply a bitwise xor to two inputs and return the result as a Buffer. @@ -144,14 +146,14 @@ pub fn buffer_bin_xor( right_offset_in_bits: usize, len_in_bits: usize, ) -> Buffer { - Buffer::from_bitwise_binary_op( + BooleanBuffer::from_bitwise_binary_op( left, left_offset_in_bits, right, right_offset_in_bits, len_in_bits, |a, b| a ^ b, - ) + ).into_inner() } /// Apply a bitwise and_not to two inputs and return the result as a Buffer. @@ -163,18 +165,18 @@ pub fn buffer_bin_and_not( right_offset_in_bits: usize, len_in_bits: usize, ) -> Buffer { - Buffer::from_bitwise_binary_op( + BooleanBuffer::from_bitwise_binary_op( left, left_offset_in_bits, right, right_offset_in_bits, len_in_bits, |a, b| a & !b, - ) + ).into_inner() } /// Apply a bitwise not to one input and return the result as a Buffer. /// The input is treated as a bitmap, meaning that offset and length are specified in number of bits. pub fn buffer_unary_not(left: &Buffer, offset_in_bits: usize, len_in_bits: usize) -> Buffer { - Buffer::from_bitwise_unary_op(left, offset_in_bits, len_in_bits, |a| !a) + BooleanBuffer::from_bitwise_unary_op(left, offset_in_bits, len_in_bits, |a| !a).into_inner() } diff --git a/arrow-select/src/nullif.rs b/arrow-select/src/nullif.rs index 8d97457cf7be..e8c1fe650717 100644 --- a/arrow-select/src/nullif.rs +++ b/arrow-select/src/nullif.rs @@ -18,7 +18,7 @@ //! Implements the `nullif` function for Arrow arrays. use arrow_array::{Array, ArrayRef, BooleanArray, make_array}; -use arrow_buffer::{BooleanBuffer, Buffer, NullBuffer}; +use arrow_buffer::{BooleanBuffer, NullBuffer}; use arrow_schema::{ArrowError, DataType}; /// Returns a new array with the same values and the validity bit to false where @@ -74,7 +74,7 @@ pub fn nullif(left: &dyn Array, right: &BooleanArray) -> Result { let mut valid_count = 0; - let b = Buffer::from_bitwise_binary_op( + let b = BooleanBuffer::from_bitwise_binary_op( left.buffer(), left.offset(), right.inner(), @@ -90,7 +90,7 @@ pub fn nullif(left: &dyn Array, right: &BooleanArray) -> Result { let mut null_count = 0; - let buffer = Buffer::from_bitwise_unary_op(right.inner(), right.offset(), len, |b| { + let buffer = BooleanBuffer::from_bitwise_unary_op(right.inner(), right.offset(), len, |b| { let t = !b; null_count += t.count_zeros() as usize; t @@ -99,7 +99,6 @@ pub fn nullif(left: &dyn Array, right: &BooleanArray) -> Result Date: Sun, 14 Dec 2025 15:26:16 -0500 Subject: [PATCH 08/15] special case aligned data --- arrow-buffer/src/buffer/boolean.rs | 93 ++++++++++++++++++++++++++++++ 1 file changed, 93 insertions(+) diff --git a/arrow-buffer/src/buffer/boolean.rs b/arrow-buffer/src/buffer/boolean.rs index 5c7481d11a07..e7060cff7dff 100644 --- a/arrow-buffer/src/buffer/boolean.rs +++ b/arrow-buffer/src/buffer/boolean.rs @@ -145,6 +145,18 @@ impl BooleanBuffer { where F: FnMut(u64, u64) -> u64, { + // Fast path for aligned inputs + if left_offset_in_bits % 8 == 0 && right_offset_in_bits % 8 == 0 { + if let Some(result) = Self::try_from_aligned_bitwise_binary_op( + &left.as_ref()[left_offset_in_bits / 8..], // aligned to byte boundary + &right.as_ref()[right_offset_in_bits / 8..], + len_in_bits, + &mut op, + ) { + return result; + } + } + // each chunk is 64 bits 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); @@ -177,6 +189,46 @@ impl BooleanBuffer { } } + /// Like [`from_bitwise_binary_op`] but optimized for the case where the + /// inputs are aligned to byte boundaries + /// + /// Returns `None` if the inputs are not fully u64 aligned + fn try_from_aligned_bitwise_binary_op( + left: &[u8], + right: &[u8], + len_in_bits: usize, + op: &mut F, + ) -> Option + where + F: FnMut(u64, u64) -> u64, + { + unsafe { + // safety: all valid bytes are valid u64s + let (left_prefix, left_u64s, left_suffix) = left.align_to::(); + let (right_prefix, right_u64s, right_suffix) = right.align_to::(); + // if there is no prefix or suffix, both buffers are aligned and we can do the operation directly + // on u64s + // TODO also handle non empty suffixes by processing them separately + if left_prefix.is_empty() + && right_prefix.is_empty() + && left_suffix.is_empty() + && right_suffix.is_empty() + { + let result_u64s = left_u64s + .iter() + .zip(right_u64s.iter()) + .map(|(l, r)| op(*l, *r)) + .collect::>(); + Some(BooleanBuffer::new(Buffer::from(result_u64s), 0, len_in_bits)) + } + else { + None + } + } + + } + + /// Create a new [`Buffer`] by applying the bitwise operation to `op` to an input buffer. /// /// This function is much faster than applying the operation bit by bit as @@ -224,6 +276,17 @@ impl BooleanBuffer { where F: FnMut(u64) -> u64, { + // try fast path for aligned input + if offset_in_bits % 8 == 0 { + if let Some(result) = Self::try_from_aligned_bitwise_unary_op( + &left.as_ref()[offset_in_bits / 8..], // align to byte boundary + len_in_bits, + &mut op, + ) { + return result; + } + } + // each chunk is 64 bits let left_chunks = BitChunks::new(left.as_ref(), offset_in_bits, len_in_bits); let mut result = MutableBuffer::with_capacity(left_chunks.num_u64s() * 8); @@ -249,6 +312,36 @@ impl BooleanBuffer { } } + /// Like [`from_bitwise_unary_op`] but optimized for the case where the + /// input is aligned to byte boundaries + fn try_from_aligned_bitwise_unary_op( + left: &[u8], + len_in_bits: usize, + op: &mut F, + ) -> Option + where + F: FnMut(u64) -> u64, + { + unsafe { + // safety: all valid bytes are valid u64s + let (left_prefix, left_u64s, left_suffix) = left.align_to::(); + // if there is no prefix or suffix, the buffer is aligned and we can do the operation directly + // on u64s + // TODO also handle non empty suffixes by processing them separately + if left_prefix.is_empty() && left_suffix.is_empty() { + let result_u64s = left_u64s + .iter() + .map(|l| op(*l)) + .collect::>(); + Some(BooleanBuffer::new(Buffer::from(result_u64s), 0, len_in_bits)) + } + else { + None + } + } + + } + /// Invokes `f` with indexes `0..len` collecting the boolean results into a new `BooleanBuffer` pub fn collect_bool bool>(len: usize, f: F) -> Self { let buffer = MutableBuffer::collect_bool(len, f); From 82bc7aab99c696b1bb1686a7baa2a51b1bfacd41 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sun, 14 Dec 2025 15:27:47 -0500 Subject: [PATCH 09/15] cleanup --- arrow-buffer/src/buffer/boolean.rs | 38 ++++++++++++++-------------- arrow-buffer/src/buffer/immutable.rs | 3 +-- arrow-buffer/src/buffer/ops.rs | 28 +++++++++++++------- arrow-select/src/nullif.rs | 11 ++++---- 4 files changed, 45 insertions(+), 35 deletions(-) diff --git a/arrow-buffer/src/buffer/boolean.rs b/arrow-buffer/src/buffer/boolean.rs index e7060cff7dff..4fbcbd41c171 100644 --- a/arrow-buffer/src/buffer/boolean.rs +++ b/arrow-buffer/src/buffer/boolean.rs @@ -172,7 +172,7 @@ impl BooleanBuffer { } } if left_chunks.remainder_len() > 0 { - debug_assert!(result.capacity() > result.len() + 8); // should not reallocate + debug_assert!(result.capacity() >= result.len() + 8); // should not reallocate result.push(op( left_chunks.remainder_bits(), right_chunks.remainder_bits(), @@ -215,20 +215,21 @@ impl BooleanBuffer { && right_suffix.is_empty() { let result_u64s = left_u64s - .iter() - .zip(right_u64s.iter()) - .map(|(l, r)| op(*l, *r)) - .collect::>(); - Some(BooleanBuffer::new(Buffer::from(result_u64s), 0, len_in_bits)) - } - else { + .iter() + .zip(right_u64s.iter()) + .map(|(l, r)| op(*l, *r)) + .collect::>(); + Some(BooleanBuffer::new( + Buffer::from(result_u64s), + 0, + len_in_bits, + )) + } else { None } } - } - /// Create a new [`Buffer`] by applying the bitwise operation to `op` to an input buffer. /// /// This function is much faster than applying the operation bit by bit as @@ -299,7 +300,7 @@ impl BooleanBuffer { } } if left_chunks.remainder_len() > 0 { - debug_assert!(result.capacity() > result.len() + 8); // should not reallocate + debug_assert!(result.capacity() >= result.len() + 8); // should not reallocate result.push(op(left_chunks.remainder_bits())); // Just pushed one u64, which may have have trailing zeros, result.truncate(left_chunks.num_bytes()); @@ -329,17 +330,16 @@ impl BooleanBuffer { // on u64s // TODO also handle non empty suffixes by processing them separately if left_prefix.is_empty() && left_suffix.is_empty() { - let result_u64s = left_u64s - .iter() - .map(|l| op(*l)) - .collect::>(); - Some(BooleanBuffer::new(Buffer::from(result_u64s), 0, len_in_bits)) - } - else { + let result_u64s = left_u64s.iter().map(|l| op(*l)).collect::>(); + Some(BooleanBuffer::new( + Buffer::from(result_u64s), + 0, + len_in_bits, + )) + } else { None } } - } /// Invokes `f` with indexes `0..len` collecting the boolean results into a new `BooleanBuffer` diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs index 3fa0750a4f93..971c590ada6b 100644 --- a/arrow-buffer/src/buffer/immutable.rs +++ b/arrow-buffer/src/buffer/immutable.rs @@ -20,11 +20,11 @@ use std::fmt::Debug; use std::ptr::NonNull; use std::sync::Arc; -use crate::{BooleanBuffer, BufferBuilder}; use crate::alloc::{Allocation, Deallocation}; #[cfg(feature = "pool")] use crate::pool::MemoryPool; use crate::util::bit_chunk_iterator::{BitChunks, UnalignedBitChunk}; +use crate::{BooleanBuffer, BufferBuilder}; use crate::{bit_util, bytes::Bytes, native::ArrowNativeType}; use super::{MutableBuffer, ScalarBuffer}; @@ -113,7 +113,6 @@ impl Buffer { Self::from(bytes) } - /// Returns the offset, in bytes, of `Self::ptr` to `Self::data` /// /// self.ptr and self.data can be different after slicing or advancing the buffer. diff --git a/arrow-buffer/src/buffer/ops.rs b/arrow-buffer/src/buffer/ops.rs index a7d26d321ad7..b10c1ed52e4c 100644 --- a/arrow-buffer/src/buffer/ops.rs +++ b/arrow-buffer/src/buffer/ops.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -use crate::BooleanBuffer; use super::{Buffer, MutableBuffer}; +use crate::BooleanBuffer; use crate::util::bit_util::ceil; /// Apply a bitwise operation `op` to four inputs and return the result as a Buffer. @@ -61,7 +61,10 @@ where /// Apply a bitwise operation `op` to two inputs and return the result as a Buffer. /// The inputs are treated as bitmaps, meaning that offsets and length are specified in number of bits. -#[deprecated(since = "57.1.0", note = "use BooleanBuffer::from_bitwise_binary_op instead")] +#[deprecated( + since = "57.1.0", + note = "use BooleanBuffer::from_bitwise_binary_op instead" +)] pub fn bitwise_bin_op_helper( left: &Buffer, left_offset_in_bits: usize, @@ -80,12 +83,16 @@ where right_offset_in_bits, len_in_bits, op, - ).into_inner() + ) + .into_inner() } /// Apply a bitwise operation `op` to one input and return the result as a Buffer. /// The input is treated as a bitmap, meaning that offset and length are specified in number of bits. -#[deprecated(since = "57.1.0", note = "use BooleanBuffer::from_bitwise_unary_op instead")] +#[deprecated( + since = "57.1.0", + note = "use BooleanBuffer::from_bitwise_unary_op instead" +)] pub fn bitwise_unary_op_helper( left: &Buffer, offset_in_bits: usize, @@ -96,7 +103,6 @@ where F: FnMut(u64) -> u64, { BooleanBuffer::from_bitwise_unary_op(left, offset_in_bits, len_in_bits, op).into_inner() - } /// Apply a bitwise and to two inputs and return the result as a Buffer. @@ -115,7 +121,8 @@ pub fn buffer_bin_and( right_offset_in_bits, len_in_bits, |a, b| a & b, - ).into_inner() + ) + .into_inner() } /// Apply a bitwise or to two inputs and return the result as a Buffer. @@ -134,7 +141,8 @@ pub fn buffer_bin_or( right_offset_in_bits, len_in_bits, |a, b| a | b, - ).into_inner() + ) + .into_inner() } /// Apply a bitwise xor to two inputs and return the result as a Buffer. @@ -153,7 +161,8 @@ pub fn buffer_bin_xor( right_offset_in_bits, len_in_bits, |a, b| a ^ b, - ).into_inner() + ) + .into_inner() } /// Apply a bitwise and_not to two inputs and return the result as a Buffer. @@ -172,7 +181,8 @@ pub fn buffer_bin_and_not( right_offset_in_bits, len_in_bits, |a, b| a & !b, - ).into_inner() + ) + .into_inner() } /// Apply a bitwise not to one input and return the result as a Buffer. diff --git a/arrow-select/src/nullif.rs b/arrow-select/src/nullif.rs index e8c1fe650717..f05ca3804cce 100644 --- a/arrow-select/src/nullif.rs +++ b/arrow-select/src/nullif.rs @@ -90,11 +90,12 @@ pub fn nullif(left: &dyn Array, right: &BooleanArray) -> Result { let mut null_count = 0; - let buffer = BooleanBuffer::from_bitwise_unary_op(right.inner(), right.offset(), len, |b| { - let t = !b; - null_count += t.count_zeros() as usize; - t - }); + let buffer = + BooleanBuffer::from_bitwise_unary_op(right.inner(), right.offset(), len, |b| { + let t = !b; + null_count += t.count_zeros() as usize; + t + }); (buffer, null_count) } }; From 3b12a465861391df70b998d9a22de6a4973276eb Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sun, 14 Dec 2025 15:40:59 -0500 Subject: [PATCH 10/15] fix docs --- arrow-buffer/src/buffer/boolean.rs | 31 ++++++++++++++++-------------- 1 file changed, 17 insertions(+), 14 deletions(-) diff --git a/arrow-buffer/src/buffer/boolean.rs b/arrow-buffer/src/buffer/boolean.rs index 4fbcbd41c171..a684340e0737 100644 --- a/arrow-buffer/src/buffer/boolean.rs +++ b/arrow-buffer/src/buffer/boolean.rs @@ -93,22 +93,24 @@ impl BooleanBuffer { } } - /// Create a new [`Buffer`] by applying the bitwise operation `op` to two input buffers. + /// Create a new [`BooleanBuffer`] by applying the bitwise operation `op` to + /// two input buffers. /// /// This function is much faster than applying the operation bit by bit as /// it processes input buffers in chunks of 64 bits (8 bytes) at a time /// /// # 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.** + /// operating on 64 bits at a time. + /// * `op` 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 + /// * The output always has zero offset. /// /// # See Also - /// - [`Buffer::from_bitwise_unary_op`] for unary operations on a single input buffer. + /// - [`BooleanBuffer::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 @@ -203,7 +205,7 @@ impl BooleanBuffer { F: FnMut(u64, u64) -> u64, { unsafe { - // safety: all valid bytes are valid u64s + // safety: all bytes are valid u64s let (left_prefix, left_u64s, left_suffix) = left.align_to::(); let (right_prefix, right_u64s, right_suffix) = right.align_to::(); // if there is no prefix or suffix, both buffers are aligned and we can do the operation directly @@ -230,22 +232,23 @@ impl BooleanBuffer { } } - /// Create a new [`Buffer`] by applying the bitwise operation to `op` to an input buffer. + /// Create a new [`BooleanBuffer`] by applying the bitwise operation to `op` to an input buffer. /// /// This function is much faster than applying the operation bit by bit as /// it processes input buffers in chunks of 64 bits (8 bytes) at a time /// /// # Notes: - /// * `op` takes two `u64` inputs and produces one `u64` output, - /// operating on 64 bits at a time. **It must only apply bitwise operations + /// * `op` takes a single `u64` inputs and produces one `u64` output + /// operating on 64 bits at a time. + /// * `op` 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.** + /// 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_binary_op`] for binary operations on a single input buffer. + /// - [`BooleanBuffer::from_bitwise_binary_op`] for binary operations on a single input buffer. /// - [`apply_bitwise_unary_op`](bit_util::apply_bitwise_unary_op) for in-place unary bitwise operations /// /// # Example: Create new [`Buffer`] from bitwise `NOT` of an input [`Buffer`] @@ -259,7 +262,7 @@ impl BooleanBuffer { /// assert_eq!(result.inner().as_slice(), &[0b00110011u8, 0b11110101u8]); /// ``` /// - /// # Example: Create a new [`Buffer`] copying a bit slice from in input slice + /// # Example: Create a new [`BooleanBuffer`] copying a bit slice from in input slice /// ``` /// # use arrow_buffer::BooleanBuffer; /// let input = [0b11001100u8, 0b10111010u8]; @@ -441,7 +444,7 @@ impl BooleanBuffer { /// Returns the inner [`Buffer`] /// - /// Note this does not account for offset and length of this [`BooleanBuffer`] + /// Note: this does not account for offset and length of this [`BooleanBuffer`] #[inline] pub fn inner(&self) -> &Buffer { &self.buffer From 02513cd9f5472e1794ff350301d884645dd54379 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sun, 14 Dec 2025 16:15:36 -0500 Subject: [PATCH 11/15] Add tests --- arrow-buffer/src/buffer/boolean.rs | 37 ++++++++++++++++++++++++++---- 1 file changed, 32 insertions(+), 5 deletions(-) diff --git a/arrow-buffer/src/buffer/boolean.rs b/arrow-buffer/src/buffer/boolean.rs index a684340e0737..59dc119e3acb 100644 --- a/arrow-buffer/src/buffer/boolean.rs +++ b/arrow-buffer/src/buffer/boolean.rs @@ -113,11 +113,11 @@ impl BooleanBuffer { /// - [`BooleanBuffer::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 + /// # Example: Create new [`BooleanBuffer`] from bitwise `AND` of two [`Buffer`]s /// ``` - /// # use arrow_buffer::BooleanBuffer; - /// let left = [0b11001100u8, 0b10111010u8]; // 2 bytes = 16 bits - /// let right = [0b10101010u8, 0b11011100u8, 0b11110000u8]; // 3 bytes = 24 bits + /// # use arrow_buffer::{Buffer, BooleanBuffer}; + /// let left = Buffer::from(vec![0b11001100u8, 0b10111010u8]); // 2 bytes = 16 bits + /// let right = Buffer::from(vec![0b10101010u8, 0b11011100u8, 0b11110000u8]); // 3 bytes = 24 bits /// // AND of the first 12 bits /// let result = BooleanBuffer::from_bitwise_binary_op( /// &left, 0, &right, 0, 12, |a, b| a & b @@ -125,7 +125,7 @@ impl BooleanBuffer { /// assert_eq!(result.inner().as_slice(), &[0b10001000u8, 0b00001000u8]); /// ``` /// - /// # Example: Create new [`Buffer`] from bitwise `OR` of two byte slices + /// # Example: Create new [`BooleanBuffer`] from bitwise `OR` of two byte slices /// ``` /// # use arrow_buffer::BooleanBuffer; /// let left = [0b11001100u8, 0b10111010u8]; @@ -694,4 +694,31 @@ mod tests { assert_eq!(buf.values().len(), 1); assert!(buf.value(0)); } + + #[test] + fn test_from_bitwise_unary_op() { + // Use 1024 boolean values so that at least some of the tests cover multiple u64 chunks and + // perfect alignment + let input_bools = (0..1024).map(|_| rand::random::()).collect::>(); + let input_buffer = BooleanBuffer::from(&input_bools[..]); + + // Note ensure we test offsets over 100 to cover multiple u64 chunks + for offset in 0..1024 { + let result = BooleanBuffer::from_bitwise_unary_op( + input_buffer.values(), + offset, + input_buffer.len() - offset, + |a| !a, + ); + let expected = input_bools[offset..] + .iter() + .map(|b| !*b) + .collect::(); + assert_eq!(result, expected); + } + } + + #[test] + fn test_from_bitwise_binary_op() { + } } From ce09b384585a822f68be6f86cc5430b0e8661784 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sun, 14 Dec 2025 16:17:39 -0500 Subject: [PATCH 12/15] more test --- arrow-buffer/src/buffer/boolean.rs | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/arrow-buffer/src/buffer/boolean.rs b/arrow-buffer/src/buffer/boolean.rs index 59dc119e3acb..773c59433d00 100644 --- a/arrow-buffer/src/buffer/boolean.rs +++ b/arrow-buffer/src/buffer/boolean.rs @@ -716,6 +716,23 @@ mod tests { .collect::(); assert_eq!(result, expected); } + + // Also test when the input doesn't cover the entire buffer + for offset in 0..512 { + let len = 512 - offset; // fixed length less than total + let result = BooleanBuffer::from_bitwise_unary_op( + input_buffer.values(), + offset, + len, + |a| !a, + ); + let expected = input_bools[offset..] + .iter() + .take(len) + .map(|b| !*b) + .collect::(); + assert_eq!(result, expected); + } } #[test] From cf0262ec039fadc9894133d2853b737ae25b7b43 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sun, 14 Dec 2025 16:25:00 -0500 Subject: [PATCH 13/15] more tests --- arrow-buffer/src/buffer/boolean.rs | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/arrow-buffer/src/buffer/boolean.rs b/arrow-buffer/src/buffer/boolean.rs index 773c59433d00..dd7fe66f7edb 100644 --- a/arrow-buffer/src/buffer/boolean.rs +++ b/arrow-buffer/src/buffer/boolean.rs @@ -737,5 +737,35 @@ mod tests { #[test] fn test_from_bitwise_binary_op() { + // pick random boolean inputs + let input_bools_left = (0..1024).map(|_| rand::random::()).collect::>(); + let input_bools_right = (0..1024).map(|_| rand::random::()).collect::>(); + let input_buffer_left = BooleanBuffer::from(&input_bools_left[..]); + let input_buffer_right = BooleanBuffer::from(&input_bools_right[..]); + + for left_offset in 0..200 { + for right_offset in [0, 4, 5, 17, 33, 24, 45, 64, 65, 100, 200] { + for len_offset in [0, 1, 44, 100, 256, 300, 512] { + let len = 1024 - len_offset - left_offset.max(right_offset); // ensure we don't go out of bounds + // compute with AND + let result = BooleanBuffer::from_bitwise_binary_op( + input_buffer_left.values(), + left_offset, + input_buffer_right.values(), + right_offset, + len, + |a, b| a & b, + ); + // compute directly from bools + let expected = input_bools_left[left_offset..] + .iter() + .zip(&input_bools_right[right_offset..]) + .take(len) + .map(|(a, b)| *a & *b) + .collect::(); + assert_eq!(result, expected); + } + } + } } } From d050bb62d48c5c832757a8dcbc69bd9ff00c0c39 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sun, 14 Dec 2025 16:25:35 -0500 Subject: [PATCH 14/15] fmt + clippy --- arrow-buffer/src/buffer/boolean.rs | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/arrow-buffer/src/buffer/boolean.rs b/arrow-buffer/src/buffer/boolean.rs index dd7fe66f7edb..a9ff78b335e2 100644 --- a/arrow-buffer/src/buffer/boolean.rs +++ b/arrow-buffer/src/buffer/boolean.rs @@ -699,7 +699,9 @@ mod tests { fn test_from_bitwise_unary_op() { // Use 1024 boolean values so that at least some of the tests cover multiple u64 chunks and // perfect alignment - let input_bools = (0..1024).map(|_| rand::random::()).collect::>(); + let input_bools = (0..1024) + .map(|_| rand::random::()) + .collect::>(); let input_buffer = BooleanBuffer::from(&input_bools[..]); // Note ensure we test offsets over 100 to cover multiple u64 chunks @@ -720,12 +722,8 @@ mod tests { // Also test when the input doesn't cover the entire buffer for offset in 0..512 { let len = 512 - offset; // fixed length less than total - let result = BooleanBuffer::from_bitwise_unary_op( - input_buffer.values(), - offset, - len, - |a| !a, - ); + let result = + BooleanBuffer::from_bitwise_unary_op(input_buffer.values(), offset, len, |a| !a); let expected = input_bools[offset..] .iter() .take(len) @@ -738,8 +736,12 @@ mod tests { #[test] fn test_from_bitwise_binary_op() { // pick random boolean inputs - let input_bools_left = (0..1024).map(|_| rand::random::()).collect::>(); - let input_bools_right = (0..1024).map(|_| rand::random::()).collect::>(); + let input_bools_left = (0..1024) + .map(|_| rand::random::()) + .collect::>(); + let input_bools_right = (0..1024) + .map(|_| rand::random::()) + .collect::>(); let input_buffer_left = BooleanBuffer::from(&input_bools_left[..]); let input_buffer_right = BooleanBuffer::from(&input_bools_right[..]); From 86f7dad85d3958713380124a8b99cf7b4d81b7a4 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Sun, 14 Dec 2025 16:38:22 -0500 Subject: [PATCH 15/15] fix docs --- arrow-buffer/src/buffer/boolean.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/arrow-buffer/src/buffer/boolean.rs b/arrow-buffer/src/buffer/boolean.rs index a9ff78b335e2..83ca792d2510 100644 --- a/arrow-buffer/src/buffer/boolean.rs +++ b/arrow-buffer/src/buffer/boolean.rs @@ -191,7 +191,7 @@ impl BooleanBuffer { } } - /// Like [`from_bitwise_binary_op`] but optimized for the case where the + /// Like [`Self::from_bitwise_binary_op`] but optimized for the case where the /// inputs are aligned to byte boundaries /// /// Returns `None` if the inputs are not fully u64 aligned @@ -316,7 +316,7 @@ impl BooleanBuffer { } } - /// Like [`from_bitwise_unary_op`] but optimized for the case where the + /// Like [`Self::from_bitwise_unary_op`] but optimized for the case where the /// input is aligned to byte boundaries fn try_from_aligned_bitwise_unary_op( left: &[u8],