From a0dc5f801de9e00dca4a2a56941f624ab905a7da Mon Sep 17 00:00:00 2001 From: Huaijin Date: Mon, 27 Jul 2026 00:04:35 +0800 Subject: [PATCH 1/2] fix(arrow-buffer): preserve bits outside the requested range in in-place bitwise ops `apply_bitwise_binary_op` and `apply_bitwise_unary_op` are documented internally as modifying only the bits in `offset_in_bits..offset_in_bits + len_in_bits` -- see the doc comments on `align_to_byte`, `set_remainder_bits` and `handle_mutable_buffer_remainder_unary`. Two paths did not honour that invariant: * `align_to_byte` wrote every bit from `bit_offset` to the byte boundary, ignoring `len_in_bits`. When the requested range started and ended inside the same non-byte-aligned byte, the trailing bits of that byte were overwritten with the result of `op` applied to padding. For example `apply_bitwise_binary_op(left, 1, right, 0, 1, |a, b| a & b)` cleared bits 2..8. * `set_remainder_bits` read the boundary byte into the low bits of a `u64` and masked it with `!((1 << remainder_len) - 1)`. Whenever the remainder spanned more than one byte that mask selected nothing, so the out-of-range bits of the boundary byte were dropped to zero instead of being preserved. No existing caller was affected. `BooleanBufferBuilder::append_packed_range` copies with `|_a, b| b` into freshly zeroed capacity, and the `BooleanBuffer`/`BooleanArray` in-place paths only run on uniquely owned buffers and re-wrap the result with the original offset and length, so the out-of-range bits were either already zero or unobservable. The existing tests only asserted the bits inside the operated range. The two shared test helpers now also assert that every bit outside the range is identical before and after, which covers all of their call sites, plus targeted tests for ranges contained in a single partial byte. The guarantee is now stated on both public functions, since callers that read the surrounding bits back depend on it. --- arrow-buffer/src/util/bit_util.rs | 111 +++++++++++++++++++++++++++--- 1 file changed, 101 insertions(+), 10 deletions(-) diff --git a/arrow-buffer/src/util/bit_util.rs b/arrow-buffer/src/util/bit_util.rs index 920ecba6d60c..eabcf3b99da8 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( // 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( /// * `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( 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,7 +458,10 @@ fn byte_aligned_bitwise_unary_op_helper( /// * `op` - Unary operation to apply /// * `buffer` - The mutable buffer to modify /// * `offset_in_bits` - Starting bit offset (not byte-aligned) -fn align_to_byte(buffer: &mut [u8], op: &mut F, offset_in_bits: usize) +/// * `len_in_bits` - Total number of bits the caller wants to process. 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(buffer: &mut [u8], op: &mut F, offset_in_bits: usize, len_in_bits: usize) where F: FnMut(u64) -> u64, { @@ -468,12 +480,13 @@ 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. + let bits_in_this_byte = (8 - bit_offset).min(len_in_bits); + let write_mask = ((((1u16 << bits_in_this_byte) - 1) << bit_offset) & 0xFF) as u8; - 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; @@ -704,6 +717,8 @@ fn set_remainder_bits(start_remainder_mut_slice: &mut [u8], rem: u64, remainder_ "start_remainder_mut_slice length must be equal to ceil(remainder_len, 8)" ); + let remainder_bytes = self::ceil(remainder_len, 8); + // Need to update the remainder bytes in the mutable buffer // but not override the bits outside the remainder @@ -719,7 +734,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) << ((remainder_bytes - 1) * 8); // Mask where the bits that are inside the remainder are 1 // and the bits outside the remainder are 0 @@ -740,8 +757,6 @@ 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); - // 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 +1120,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 +1139,39 @@ 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 {i} outside the requested range was modified ({context})" + ); + } } /// Verifies that a unary operation applied to a buffer using u64 chunks @@ -1146,6 +1196,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 = @@ -1156,6 +1208,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 +1484,37 @@ 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 { + 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]); + } + #[test] fn test_bitwise_binary_op_offset_out_of_bounds() { let input = vec![0b10101010u8, 0b01010101u8]; From 4fbbb018640b3d52e6268259c3a51b7a7a724cf3 Mon Sep 17 00:00:00 2001 From: Huaijin Date: Mon, 27 Jul 2026 12:53:43 +0800 Subject: [PATCH 2/2] fix(arrow-buffer): address review feedback on out-of-range bit preservation - Rename `align_to_byte`'s new parameter to `remaining_len_in_bits`; the binary caller passes the already-clamped `bits_to_next_byte` while the unary caller passes the full length, so "total bits to process" only described one of them. - Add `debug_assert_ne!(bit_offset, 0)` to `align_to_byte`, making the documented "not byte-aligned" precondition checked, and drop the dead `& 0xFF` from the write mask now that the bound is explicit. - Use the slice length instead of recomputing `ceil(remainder_len, 8)` in `set_remainder_bits`; the assertion above already proves them equal. - Cover `len == 8 - offset` in `test_ops_ending_inside_the_first_partial_byte`; the exclusive range skipped the byte-boundary case and never ran at all for `offset == 7`. - Add targeted regression tests for the multi-byte remainder bug. It was only covered indirectly through the shared helpers, so weakening them would have let it regress silently. --- arrow-buffer/src/util/bit_util.rs | 65 +++++++++++++++++++++++++------ 1 file changed, 53 insertions(+), 12 deletions(-) diff --git a/arrow-buffer/src/util/bit_util.rs b/arrow-buffer/src/util/bit_util.rs index eabcf3b99da8..e1bf14bf2c34 100644 --- a/arrow-buffer/src/util/bit_util.rs +++ b/arrow-buffer/src/util/bit_util.rs @@ -458,16 +458,23 @@ fn byte_aligned_bitwise_unary_op_helper( /// * `op` - Unary operation to apply /// * `buffer` - The mutable buffer to modify /// * `offset_in_bits` - Starting bit offset (not byte-aligned) -/// * `len_in_bits` - Total number of bits the caller wants to process. 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(buffer: &mut [u8], op: &mut F, offset_in_bits: usize, len_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( + 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]; @@ -483,8 +490,11 @@ where // 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. - let bits_in_this_byte = (8 - bit_offset).min(len_in_bits); - let write_mask = ((((1u16 << bits_in_this_byte) - 1) << bit_offset) & 0xFF) as u8; + // + // `bits_in_this_byte + bit_offset <= 8`, so the mask always fits in a `u8`. The + // shift is done in `u16` only so the expression stays correct for a full byte. + let bits_in_this_byte = (8 - bit_offset).min(remaining_len_in_bits); + let write_mask = (((1u16 << bits_in_this_byte) - 1) << bit_offset) as u8; let result_first_byte = (first_byte & !write_mask) | (result_first_byte & write_mask); @@ -717,8 +727,6 @@ fn set_remainder_bits(start_remainder_mut_slice: &mut [u8], rem: u64, remainder_ "start_remainder_mut_slice length must be equal to ceil(remainder_len, 8)" ); - let remainder_bytes = self::ceil(remainder_len, 8); - // Need to update the remainder bytes in the mutable buffer // but not override the bits outside the remainder @@ -736,7 +744,7 @@ fn set_remainder_bits(start_remainder_mut_slice: &mut [u8], rem: u64, remainder_ // 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) << ((remainder_bytes - 1) * 8); + 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 @@ -757,6 +765,8 @@ fn set_remainder_bits(start_remainder_mut_slice: &mut [u8], rem: u64, remainder_ // Write back the result to the mutable slice { + 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]; @@ -1169,7 +1179,9 @@ mod tests { assert_eq!( get_bit(before, i), get_bit(after, i), - "bit {i} outside the requested range was modified ({context})" + "bit {} outside the requested range was modified ({})", + i, + context ); } } @@ -1490,7 +1502,8 @@ mod tests { fn test_ops_ending_inside_the_first_partial_byte() { let (left, right) = create_test_data(32); for offset in 1..8 { - for len in 1..(8 - offset) { + // 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); @@ -1515,6 +1528,34 @@ mod tests { 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];