This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new f1ba4352bb fix(arrow-buffer): preserve bits outside the requested
range in in-place bitwise ops (#10444)
f1ba4352bb is described below
commit f1ba4352bb583c47feba8a008516b1be21edf7f8
Author: Huaijin <[email protected]>
AuthorDate: Tue Aug 4 02:24:15 2026 +0800
fix(arrow-buffer): preserve bits outside the requested range in in-place
bitwise ops (#10444)
# Which issue does this PR close?
No dedicated issue — found while implementing #10425, whose optimization
is the first caller to depend on the behaviour fixed here. Happy to file
one if you'd prefer it tracked separately.
# Rationale for this change
`apply_bitwise_binary_op` and `apply_bitwise_unary_op` are built around
the invariant that only the bits in `offset_in_bits..offset_in_bits +
len_in_bits` are modified — the helpers implementing it
(`align_to_byte`, `set_remainder_bits`,
`handle_mutable_buffer_remainder_unary`) all say so in their doc
comments. Two paths didn't honour it:
**1. `align_to_byte` ignored `len_in_bits`**, writing every bit from
`bit_offset` to the byte boundary. If the range started and ended inside
the same non-byte-aligned byte, the trailing bits were overwritten with
`op` applied to padding:
```rust
let mut left = vec![0b11111111u8, 0b11111111u8];
apply_bitwise_binary_op(&mut left, 1, &[0u8, 0u8], 0, 1, |a, b| a & b);
// expected [0b11111101, 0b11111111], got [0b00000001, 0b11111111]
```
**2. `set_remainder_bits` zeroed the boundary byte's out-of-range
bits**, because it read that byte into the *low* bits of a `u64` and
masked it with `!((1 << remainder_len) - 1)`. Once the remainder spanned
more than one byte, the mask selected nothing. This hit any length
leaving a remainder of 9..63 bits that isn't a multiple of 8.
No existing caller was affected, which is why it went unnoticed:
`BooleanBufferBuilder::append_packed_range` copies with `|_a, b| b` into
capacity `advance()` just zeroed, and the `BooleanBuffer`/`BooleanArray`
in-place paths run only on uniquely owned buffers and re-wrap with the
original offset/length. So the out-of-range bits were either already
zero or unobservable.
# What changes are included in this PR?
- `align_to_byte` takes `len_in_bits` and masks its write to
`bit_offset..bit_offset + min(8 - bit_offset, len_in_bits)`.
- `set_remainder_bits` shifts the boundary byte to the position it
actually occupies in the word before masking.
- Both public functions now document the preservation guarantee.
# Are these changes tested?
Yes. The gap existed because the two shared test helpers only asserted
bits *inside* the operated range. They now also assert every bit outside
it is byte-for-byte identical before and after, which retroactively
covers all their call sites — that alone caught bug 2 in 11 pre-existing
tests. Added `test_ops_ending_inside_the_first_partial_byte` (sweeps
every offset/len pair inside one partial byte, for AND/OR/XOR and NOT)
plus minimal regression tests for the two reproducers.
`cargo test -p arrow-buffer` passes (342 unit + 54 doc), as do
`arrow-array`, `arrow` and `arrow-select`. Clippy and fmt clean.
# Are there any user-facing changes?
A behaviour change in two public functions, in the direction of the
documented intent: bits outside the requested range are no longer
clobbered. No signature changes, and no existing caller could observe
the old behaviour. The guarantee is now explicit in the rustdoc.
---
arrow-buffer/src/util/bit_util.rs | 152 +++++++++++++++++++++++++++++++++++---
1 file changed, 142 insertions(+), 10 deletions(-)
diff --git a/arrow-buffer/src/util/bit_util.rs
b/arrow-buffer/src/util/bit_util.rs
index 920ecba6d6..7530670554 100644
--- a/arrow-buffer/src/util/bit_util.rs
+++ b/arrow-buffer/src/util/bit_util.rs
@@ -168,6 +168,10 @@ pub(crate) fn read_up_to_byte_from_offset(
/// * `len_in_bits` - Number of bits to process
/// * `op` - Binary operation to apply (e.g., `|a, b| a & b`). Applied a word
at a time
///
+/// Only the bits in `left_offset_in_bits..left_offset_in_bits + len_in_bits`
are
+/// modified. Bits of `left` outside that range are left unchanged, including
the
+/// bits sharing a byte with either end of the range.
+///
/// # Example: Modify entire buffer
/// ```
/// # use arrow_buffer::MutableBuffer;
@@ -247,6 +251,7 @@ pub fn apply_bitwise_binary_op<F>(
// Hope it gets inlined
&mut |left| op(left, right_first_byte as u64),
left_offset_in_bits,
+ bits_to_next_byte,
);
}
@@ -280,6 +285,10 @@ pub fn apply_bitwise_binary_op<F>(
/// * `len_in_bits` - Number of bits to process
/// * `op` - Unary operation to apply (e.g., `|a| !a`). Applied a word at a
time
///
+/// Only the bits in `offset_in_bits..offset_in_bits + len_in_bits` are
modified.
+/// Bits outside that range are left unchanged, including the bits sharing a
byte
+/// with either end of the range.
+///
/// # Example: Modify entire buffer
/// ```
/// # use arrow_buffer::MutableBuffer;
@@ -325,7 +334,7 @@ pub fn apply_bitwise_unary_op<F>(
if is_mutable_buffer_byte_aligned {
byte_aligned_bitwise_unary_op_helper(buffer, offset_in_bits,
len_in_bits, op);
} else {
- align_to_byte(buffer, &mut op, offset_in_bits);
+ align_to_byte(buffer, &mut op, offset_in_bits, len_in_bits);
// If we are not byte aligned we will read the first few bits
let bits_to_next_byte = 8 - left_bit_offset;
@@ -449,13 +458,23 @@ fn byte_aligned_bitwise_unary_op_helper<F>(
/// * `op` - Unary operation to apply
/// * `buffer` - The mutable buffer to modify
/// * `offset_in_bits` - Starting bit offset (not byte-aligned)
-fn align_to_byte<F>(buffer: &mut [u8], op: &mut F, offset_in_bits: usize)
-where
+/// * `remaining_len_in_bits` - Number of bits still to process starting at
`offset_in_bits`.
+/// When this is smaller than the number of bits left in the byte, the
trailing bits of
+/// the byte are left untouched.
+fn align_to_byte<F>(
+ buffer: &mut [u8],
+ op: &mut F,
+ offset_in_bits: usize,
+ remaining_len_in_bits: usize,
+) where
F: FnMut(u64) -> u64,
{
let byte_offset = offset_in_bits / 8;
let bit_offset = offset_in_bits % 8;
+ // Byte aligned offsets must take the byte aligned path instead
+ debug_assert_ne!(bit_offset, 0, "offset_in_bits must not be byte aligned");
+
// 1. read the first byte from the buffer
let first_byte: u8 = buffer[byte_offset];
@@ -468,12 +487,16 @@ where
// 4. Shift back the result to the original position
let result_first_byte = result_first_byte << bit_offset;
- // 5. Mask the bits that are outside the relevant bits in the byte
- // so the bits until bit_offset are 1 and the rest are 0
- let mask_for_first_bit_offset = (1 << bit_offset) - 1;
+ // 5. Mask in only the bits the caller asked to process, i.e. the bits in
+ // `bit_offset..bit_offset + bits_in_this_byte`. The request may end
before the
+ // byte boundary, in which case the trailing bits must be preserved as
well.
+ //
+ // `bit_offset` is in `1..=7` per the assert above, so
`bits_in_this_byte` is at
+ // most 7 and `bits_in_this_byte + bit_offset <= 8`, keeping the mask
within a `u8`.
+ let bits_in_this_byte = (8 - bit_offset).min(remaining_len_in_bits);
+ let write_mask = ((1u8 << bits_in_this_byte) - 1) << bit_offset;
- let result_first_byte =
- (first_byte & mask_for_first_bit_offset) | (result_first_byte &
!mask_for_first_bit_offset);
+ let result_first_byte = (first_byte & !write_mask) | (result_first_byte &
write_mask);
// 6. write back the result to the buffer
buffer[byte_offset] = result_first_byte;
@@ -719,7 +742,9 @@ fn set_remainder_bits(start_remainder_mut_slice: &mut [u8],
rem: u64, remainder_
// Unwrap as we already validated the slice is not empty
.unwrap();
- let current = *current as u64;
+ // Shift the boundary byte to the position it occupies within `rem`,
otherwise
+ // its bits would be compared against the wrong end of the mask below
+ let current = (*current as u64) << ((start_remainder_mut_slice.len() -
1) * 8);
// Mask where the bits that are inside the remainder are 1
// and the bits outside the remainder are 0
@@ -740,7 +765,7 @@ fn set_remainder_bits(start_remainder_mut_slice: &mut [u8],
rem: u64, remainder_
// Write back the result to the mutable slice
{
- let remainder_bytes = self::ceil(remainder_len, 8);
+ let remainder_bytes = start_remainder_mut_slice.len();
// we are counting starting from the least significant bit, so
to_le_bytes should be correct
let rem = &rem.to_le_bytes()[0..remainder_bytes];
@@ -1105,6 +1130,8 @@ mod tests {
.map(|(l, r)| expected_op(*l, *r))
.collect();
+ let before = left_buffer.as_slice().to_vec();
+
apply_bitwise_binary_op(
left_buffer.as_slice_mut(),
left_offset_in_bits,
@@ -1122,6 +1149,41 @@ mod tests {
"Failed with left_offset={}, right_offset={}, len={}",
left_offset_in_bits, right_offset_in_bits, len_in_bits
);
+
+ assert_bits_outside_range_preserved(
+ &before,
+ left_buffer.as_slice(),
+ left_offset_in_bits,
+ len_in_bits,
+ &format!(
+ "left_offset={}, right_offset={}, len={}",
+ left_offset_in_bits, right_offset_in_bits, len_in_bits
+ ),
+ );
+ }
+
+ /// Asserts that every bit outside `offset_in_bits..offset_in_bits +
len_in_bits`
+ /// is identical in `before` and `after`.
+ fn assert_bits_outside_range_preserved(
+ before: &[u8],
+ after: &[u8],
+ offset_in_bits: usize,
+ len_in_bits: usize,
+ context: &str,
+ ) {
+ assert_eq!(before.len(), after.len());
+ for i in 0..before.len() * 8 {
+ if i >= offset_in_bits && i < offset_in_bits + len_in_bits {
+ continue;
+ }
+ assert_eq!(
+ get_bit(before, i),
+ get_bit(after, i),
+ "bit {} outside the requested range was modified ({})",
+ i,
+ context
+ );
+ }
}
/// Verifies that a unary operation applied to a buffer using u64 chunks
@@ -1146,6 +1208,8 @@ mod tests {
.map(|b| expected_op(*b))
.collect();
+ let before = buffer.as_slice().to_vec();
+
apply_bitwise_unary_op(buffer.as_slice_mut(), offset_in_bits,
len_in_bits, op);
let result: Vec<bool> =
@@ -1156,6 +1220,14 @@ mod tests {
"Failed with offset={}, len={}",
offset_in_bits, len_in_bits
);
+
+ assert_bits_outside_range_preserved(
+ &before,
+ buffer.as_slice(),
+ offset_in_bits,
+ len_in_bits,
+ &format!("offset={}, len={}", offset_in_bits, len_in_bits),
+ );
}
// Helper to create test data of specific length
@@ -1424,6 +1496,66 @@ mod tests {
);
}
+ /// Ranges that start and end inside the same non-byte-aligned byte must
not
+ /// touch the trailing bits of that byte.
+ #[test]
+ fn test_ops_ending_inside_the_first_partial_byte() {
+ let (left, right) = create_test_data(32);
+ for offset in 1..8 {
+ // Inclusive so the range ending exactly on the byte boundary is
covered too
+ for len in 1..=(8 - offset) {
+ test_all_binary_ops(&left, &right, offset, offset, len);
+ test_all_binary_ops(&left, &right, offset, (offset + 3) % 8,
len);
+ test_mutable_buffer_unary_op_helper(&left, offset, len, |a|
!a, |a| !a);
+ }
+ }
+ }
+
+ #[test]
+ fn test_and_within_first_partial_byte_preserves_trailing_bits() {
+ let mut left = vec![0b11111111u8, 0b11111111u8];
+ let right = vec![0b00000000u8, 0b00000000u8];
+ // AND a single bit at bit offset 1: only bit 1 may be cleared
+ apply_bitwise_binary_op(&mut left, 1, &right, 0, 1, |a, b| a & b);
+ assert_eq!(left, vec![0b11111101u8, 0b11111111u8]);
+ }
+
+ #[test]
+ fn test_not_within_first_partial_byte_preserves_trailing_bits() {
+ let mut buffer = vec![0b00000000u8];
+ // NOT two bits at bit offset 3: only bits 3 and 4 may be flipped
+ apply_bitwise_unary_op(&mut buffer, 3, 2, |a| !a);
+ assert_eq!(buffer, vec![0b00011000u8]);
+ }
+
+ /// When the remainder spans more than one byte, the byte holding the end
of the
+ /// range is the *last* byte of the remainder, not the first. Its bits
above the
+ /// remainder must survive.
+ #[test]
+ fn test_or_with_multi_byte_remainder_preserves_boundary_bits() {
+ let mut left = vec![0b00000000u8, 0b00000000u8, 0b11110000u8];
+ let right = vec![0b11111111u8, 0b11111111u8, 0b11111111u8];
+ // OR over 20 bits: bits 20..24 of `left` are outside the range and
must stay set
+ apply_bitwise_binary_op(&mut left, 0, &right, 0, 20, |a, b| a | b);
+ assert_eq!(
+ left,
+ vec![0b11111111u8, 0b11111111u8, 0b11111111u8],
+ "the boundary byte lost its out-of-range bits"
+ );
+ }
+
+ #[test]
+ fn test_not_with_multi_byte_remainder_preserves_boundary_bits() {
+ let mut buffer = vec![0b00000000u8, 0b00000000u8, 0b11111111u8];
+ // NOT over 20 bits: only bits 16..20 of the last byte may be flipped
+ apply_bitwise_unary_op(&mut buffer, 0, 20, |a| !a);
+ assert_eq!(
+ buffer,
+ vec![0b11111111u8, 0b11111111u8, 0b11110000u8],
+ "the boundary byte lost its out-of-range bits"
+ );
+ }
+
#[test]
fn test_bitwise_binary_op_offset_out_of_bounds() {
let input = vec![0b10101010u8, 0b01010101u8];
