diff --git a/arrow-buffer/src/util/bit_util.rs b/arrow-buffer/src/util/bit_util.rs index 920ecba6d60c..e1bf14bf2c34 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,13 +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) -fn align_to_byte(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( + 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. + // + // `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 & 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 = @@ -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]; diff --git a/parquet/src/arrow/arrow_reader/selection/algebra.rs b/parquet/src/arrow/arrow_reader/selection/algebra.rs index 9c25afbd3b15..022051805b87 100644 --- a/parquet/src/arrow/arrow_reader/selection/algebra.rs +++ b/parquet/src/arrow/arrow_reader/selection/algebra.rs @@ -23,7 +23,7 @@ //! [`BooleanBuffer`] masks. use super::{MaskRunIter, RowSelection, RowSelectionInner, RowSelector}; -use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder}; +use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, MutableBuffer, bit_util}; use std::cmp::Ordering; use std::iter::Peekable; @@ -271,18 +271,7 @@ pub(super) fn intersect_masks(l: &BooleanBuffer, r: &BooleanBuffer) -> BooleanBu if l.len() == r.len() { return l & r; } - let common = l.len().min(r.len()); - let head = &l.slice(0, common) & &r.slice(0, common); - let (longer, longer_len) = if l.len() > r.len() { - (l, l.len()) - } else { - (r, r.len()) - }; - let tail = longer.slice(common, longer_len - common); - let mut builder = BooleanBufferBuilder::new(longer_len); - builder.append_buffer(&head); - builder.append_buffer(&tail); - builder.finish() + combine_uneven_masks(l, r, |a, b| a & b) } /// Bitwise OR of two mask-backed selections. Longer side's tail passes through. @@ -290,18 +279,47 @@ pub(super) fn union_masks(l: &BooleanBuffer, r: &BooleanBuffer) -> BooleanBuffer if l.len() == r.len() { return l | r; } - let common = l.len().min(r.len()); - let head = &l.slice(0, common) | &r.slice(0, common); - let (longer, longer_len) = if l.len() > r.len() { - (l, l.len()) - } else { - (r, r.len()) - }; - let tail = longer.slice(common, longer_len - common); - let mut builder = BooleanBufferBuilder::new(longer_len); - builder.append_buffer(&head); - builder.append_buffer(&tail); - builder.finish() + combine_uneven_masks(l, r, |a, b| a | b) +} + +/// Combines two masks of differing lengths with the bitwise operation `op`, +/// passing the longer side's tail through unchanged. +/// +/// The longer mask is copied once into a [`MutableBuffer`] and `op` is then +/// applied in place over the common prefix. This avoids materialising the +/// prefix into its own buffer and copying both prefix and tail again through a +/// [`BooleanBufferBuilder`]. +/// +/// Neither the mask offsets nor the prefix length are assumed to be byte +/// aligned: the copy keeps the longer mask's sub-byte offset so it stays a +/// plain byte copy, and the offset is carried over to the returned buffer. +fn combine_uneven_masks(l: &BooleanBuffer, r: &BooleanBuffer, op: F) -> BooleanBuffer +where + F: FnMut(u64, u64) -> u64, +{ + let (longer, shorter) = if l.len() > r.len() { (l, r) } else { (r, l) }; + let common = shorter.len(); + if common == 0 { + return longer.clone(); + } + + let bit_offset = longer.offset() % 8; + let start = longer.offset() / 8; + let end = bit_util::ceil(longer.offset() + longer.len(), 8); + let bytes = &longer.values()[start..end]; + let mut buffer = MutableBuffer::new(bytes.len()); + buffer.extend_from_slice(bytes); + + bit_util::apply_bitwise_binary_op( + buffer.as_slice_mut(), + bit_offset, + shorter.values(), + shorter.offset(), + common, + op, + ); + + BooleanBuffer::new(buffer.into(), bit_offset, longer.len()) } /// Applies `other` to the selected rows of `mask`, preserving the original row domain. @@ -825,4 +843,89 @@ mod tests { let bits: Vec = (0..5).map(|i| r_mask.value(i)).collect(); assert_eq!(bits, vec![true, true, false, false, true]); } + + /// Expected result of combining two masks of possibly differing lengths: + /// `op` over the common prefix, then the longer side's tail unchanged. + fn expected_uneven(l: &[bool], r: &[bool], op: fn(bool, bool) -> bool) -> Vec { + let common = l.len().min(r.len()); + let longer = if l.len() > r.len() { l } else { r }; + (0..common) + .map(|i| op(l[i], r[i])) + .chain(longer[common..].iter().copied()) + .collect() + } + + fn assert_mask_eq(actual: &BooleanBuffer, expected: &[bool], context: &str) { + assert_eq!(actual.len(), expected.len(), "{context}: length"); + let actual: Vec = actual.iter().collect(); + assert_eq!(actual, expected, "{context}"); + } + + #[test] + fn test_uneven_masks_with_offsets() { + // Cover offsets and prefix lengths that are not byte (or word) aligned + // on either side, including the case where the longer mask starts mid + // byte and the common prefix ends mid byte. + let base: Vec = (0..600).map(|i| i % 7 == 0 || i % 3 == 1).collect(); + let other: Vec = (0..600).map(|i| i % 5 == 2 || i % 11 == 4).collect(); + let base = BooleanBuffer::from(base); + let other = BooleanBuffer::from(other); + + for l_offset in [0, 1, 5, 8, 13, 64, 67] { + for r_offset in [0, 1, 3, 8, 60, 64, 70] { + for (l_len, r_len) in [(0, 9), (9, 0), (1, 200), (200, 1), (63, 130), (321, 65)] { + let l = base.slice(l_offset, l_len); + let r = other.slice(r_offset, r_len); + let l_bits: Vec = l.iter().collect(); + let r_bits: Vec = r.iter().collect(); + let context = + format!("l_offset={l_offset} r_offset={r_offset} lens=({l_len},{r_len})"); + + assert_mask_eq( + &intersect_masks(&l, &r), + &expected_uneven(&l_bits, &r_bits, |a, b| a && b), + &format!("intersect {context}"), + ); + assert_mask_eq( + &union_masks(&l, &r), + &expected_uneven(&l_bits, &r_bits, |a, b| a || b), + &format!("union {context}"), + ); + } + } + } + } + + #[test] + fn test_uneven_masks_fuzz() { + let mut rng = rng(); + for _ in 0..200 { + let l_offset = rng.random_range(0..70); + let r_offset = rng.random_range(0..70); + let l_len = rng.random_range(0..300); + let r_len = rng.random_range(0..300); + + let l_bits: Vec = (0..l_offset + l_len) + .map(|_| rng.random_bool(0.5)) + .collect(); + let r_bits: Vec = (0..r_offset + r_len) + .map(|_| rng.random_bool(0.5)) + .collect(); + let l = BooleanBuffer::from(l_bits).slice(l_offset, l_len); + let r = BooleanBuffer::from(r_bits).slice(r_offset, r_len); + let l_bits: Vec = l.iter().collect(); + let r_bits: Vec = r.iter().collect(); + + assert_mask_eq( + &intersect_masks(&l, &r), + &expected_uneven(&l_bits, &r_bits, |a, b| a && b), + "intersect", + ); + assert_mask_eq( + &union_masks(&l, &r), + &expected_uneven(&l_bits, &r_bits, |a, b| a || b), + "union", + ); + } + } }