diff --git a/arrow/src/array/builder.rs b/arrow/src/array/builder.rs index a9e7fe24a830..71fd9a90e7a0 100644 --- a/arrow/src/array/builder.rs +++ b/arrow/src/array/builder.rs @@ -302,6 +302,15 @@ impl BooleanBufferBuilder { self.len } + #[inline] + pub fn set_bit(&mut self, index: usize, v: bool) { + if v { + bit_util::set_bit(self.buffer.as_mut(), index); + } else { + bit_util::unset_bit(self.buffer.as_mut(), index); + } + } + #[inline] pub fn is_empty(&self) -> bool { self.len == 0 @@ -2266,7 +2275,7 @@ mod tests { } #[test] - fn test_write_bytes() { + fn test_boolean_buffer_builder_write_bytes() { let mut b = BooleanBufferBuilder::new(4); b.append(false); b.append(true); @@ -2285,6 +2294,83 @@ mod tests { assert_eq!(1, buffer.len()); } + #[test] + fn test_boolean_buffer_builder_unset_first_bit() { + let mut buffer = BooleanBufferBuilder::new(4); + buffer.append(true); + buffer.append(true); + buffer.append(false); + buffer.append(true); + buffer.set_bit(0, false); + assert_eq!(buffer.len(), 4); + assert_eq!(buffer.finish().as_slice(), &[0b1010_u8]); + } + + #[test] + fn test_boolean_buffer_builder_unset_last_bit() { + let mut buffer = BooleanBufferBuilder::new(4); + buffer.append(true); + buffer.append(true); + buffer.append(false); + buffer.append(true); + buffer.set_bit(3, false); + assert_eq!(buffer.len(), 4); + assert_eq!(buffer.finish().as_slice(), &[0b0011_u8]); + } + + #[test] + fn test_boolean_buffer_builder_unset_an_inner_bit() { + let mut buffer = BooleanBufferBuilder::new(5); + buffer.append(true); + buffer.append(true); + buffer.append(false); + buffer.append(true); + buffer.set_bit(1, false); + assert_eq!(buffer.len(), 4); + assert_eq!(buffer.finish().as_slice(), &[0b1001_u8]); + } + + #[test] + fn test_boolean_buffer_builder_unset_several_bits() { + let mut buffer = BooleanBufferBuilder::new(5); + buffer.append(true); + buffer.append(true); + buffer.append(true); + buffer.append(false); + buffer.append(true); + buffer.set_bit(1, false); + buffer.set_bit(2, false); + assert_eq!(buffer.len(), 5); + assert_eq!(buffer.finish().as_slice(), &[0b10001_u8]); + } + + #[test] + fn test_boolean_buffer_builder_unset_several_bits_bigger_than_one_byte() { + let mut buffer = BooleanBufferBuilder::new(16); + buffer.append_n(10, true); + buffer.set_bit(0, false); + buffer.set_bit(3, false); + buffer.set_bit(9, false); + assert_eq!(buffer.len(), 10); + assert_eq!(buffer.finish().as_slice(), &[0b11110110_u8, 0b01_u8]); + } + + #[test] + fn test_boolean_buffer_builder_flip_several_bits_bigger_than_one_byte() { + let mut buffer = BooleanBufferBuilder::new(16); + buffer.append_n(5, true); + buffer.append_n(5, false); + buffer.append_n(5, true); + buffer.set_bit(0, false); + buffer.set_bit(3, false); + buffer.set_bit(9, false); + buffer.set_bit(6, true); + buffer.set_bit(14, true); + buffer.set_bit(13, false); + assert_eq!(buffer.len(), 15); + assert_eq!(buffer.finish().as_slice(), &[0b01010110_u8, 0b1011100_u8]); + } + #[test] fn test_boolean_array_builder_append_slice() { let arr1 = diff --git a/arrow/src/util/bit_util.rs b/arrow/src/util/bit_util.rs index 4cfa8168e713..7298eb3c5715 100644 --- a/arrow/src/util/bit_util.rs +++ b/arrow/src/util/bit_util.rs @@ -62,7 +62,7 @@ pub unsafe fn get_bit_raw(data: *const u8, i: usize) -> bool { (*data.add(i >> 3) & BIT_MASK[i & 7]) != 0 } -/// Sets bit at position `i` for `data` +/// Sets bit at position `i` for `data` to 1 #[inline] pub fn set_bit(data: &mut [u8], i: usize) { data[i >> 3] |= BIT_MASK[i & 7];