Skip to content

Commit

Permalink
Add set_bit to BooleanBufferBuilder to allow mutating bit in index (#383
Browse files Browse the repository at this point in the history
)

* Add set_bit to BooleanBufferBuilder to allow mutating bits in the builder

* Fix tests

* Update builder.rs

* Update builder.rs

* Fix clippy failures

Co-authored-by: Boaz Berman <boaz@codota.com>
  • Loading branch information
boazberman and boaz-codota authored Jun 8, 2021
1 parent ddc8a36 commit 18c804a
Show file tree
Hide file tree
Showing 2 changed files with 88 additions and 2 deletions.
88 changes: 87 additions & 1 deletion arrow/src/array/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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);
Expand All @@ -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 =
Expand Down
2 changes: 1 addition & 1 deletion arrow/src/util/bit_util.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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];
Expand Down

0 comments on commit 18c804a

Please sign in to comment.