diff --git a/src/array/growable/utils.rs b/src/array/growable/utils.rs index 4afdc1a05ea..6a67988e2d8 100644 --- a/src/array/growable/utils.rs +++ b/src/array/growable/utils.rs @@ -46,10 +46,8 @@ pub(super) fn extend_validity( ) { if let Some(bitmap) = validity { assert!(start + len <= bitmap.len()); - unsafe { - let iter = (start..start + len).map(|i| bitmap.get_bit_unchecked(i)); - mutable_validity.extend_from_trusted_len_iter_unchecked(iter); - }; + let (slice, offset, _) = bitmap.as_slice(); + mutable_validity.extend_from_slice(slice, start + offset, len); } else if use_validity { mutable_validity.extend_constant(len, true); }; diff --git a/src/bitmap/mutable.rs b/src/bitmap/mutable.rs index fa5cba54fb2..bd6a183f0b0 100644 --- a/src/bitmap/mutable.rs +++ b/src/bitmap/mutable.rs @@ -131,16 +131,63 @@ impl MutableBitmap { self.length = len; } + fn extend_set(&mut self, mut additional: usize) { + let offset = self.length % 8; + let added = if offset != 0 { + // offset != 0 => at least one byte in the buffer + let last_index = self.buffer.len() - 1; + let last = &mut self.buffer[last_index]; + + let remaining = 0b11111111u8; + let remaining = remaining >> 8usize.saturating_sub(additional); + let remaining = remaining << offset; + *last |= remaining; + std::cmp::min(additional, 8 - offset) + } else { + 0 + }; + self.length += added; + additional = additional.saturating_sub(added); + if additional > 0 { + debug_assert_eq!(self.length % 8, 0); + let existing = self.buffer.len(); + let required = (self.length + additional).saturating_add(7) / 8; + // add remaining as full bytes + self.buffer.extend_from_trusted_len_iter( + std::iter::repeat(0b11111111u8).take(required - existing), + ); + self.length += additional; + } + } + + fn extend_unset(&mut self, mut additional: usize) { + let offset = self.length % 8; + let added = if offset != 0 { + // offset != 0 => at least one byte in the buffer + let last_index = self.buffer.len() - 1; + let last = &mut self.buffer[last_index]; + *last &= 0b11111111u8 >> (8 - offset); // unset them + std::cmp::min(additional, 8 - offset) + } else { + 0 + }; + self.length += added; + additional = additional.saturating_sub(added); + if additional > 0 { + debug_assert_eq!(self.length % 8, 0); + self.buffer + .resize((self.length + additional).saturating_add(7) / 8, 0); + self.length += additional; + } + } + /// Extends [`MutableBitmap`] by `additional` values of constant `value`. #[inline] pub fn extend_constant(&mut self, additional: usize, value: bool) { if value { - let iter = std::iter::repeat(true).take(additional); - self.extend_from_trusted_len_iter(iter); + self.extend_set(additional) } else { - self.buffer - .resize((self.length + additional).saturating_add(7) / 8, 0); - self.length += additional; + self.extend_unset(additional) } } diff --git a/src/compute/mod.rs b/src/compute/mod.rs index 332905b405d..232677e2679 100644 --- a/src/compute/mod.rs +++ b/src/compute/mod.rs @@ -7,13 +7,6 @@ //! Some dynamically-typed operators have an auxiliary function, `can_*`, that returns //! true if the operator can be applied to the particular `DataType`. -pub mod aggregate; -pub mod arithmetics; -pub mod arity; -pub mod boolean; -pub mod boolean_kleene; -pub mod cast; -pub mod comparison; pub mod concat; pub mod contains; pub mod filter; diff --git a/tests/it/bitmap/mutable.rs b/tests/it/bitmap/mutable.rs index 2704f4c1f07..7acf4781666 100644 --- a/tests/it/bitmap/mutable.rs +++ b/tests/it/bitmap/mutable.rs @@ -201,3 +201,49 @@ fn debug() { b.push(true); assert_eq!(format!("{:?}", b), "[0b11000001, 0b_______1]"); } + +#[test] +fn extend_set() { + let mut b = MutableBitmap::new(); + b.extend_constant(6, true); + assert_eq!(b.as_slice(), &[0b11111111]); + assert_eq!(b.len(), 6); + + let mut b = MutableBitmap::from(&[false]); + b.extend_constant(6, true); + assert_eq!(b.as_slice(), &[0b01111110]); + assert_eq!(b.len(), 1 + 6); + + let mut b = MutableBitmap::from(&[false]); + b.extend_constant(9, true); + assert_eq!(b.as_slice(), &[0b11111110, 0b11111111]); + assert_eq!(b.len(), 1 + 9); + + let mut b = MutableBitmap::from(&[false, false, false, false]); + b.extend_constant(2, true); + assert_eq!(b.as_slice(), &[0b00110000]); + assert_eq!(b.len(), 4 + 2); +} + +#[test] +fn extend_unset() { + let mut b = MutableBitmap::new(); + b.extend_constant(6, false); + assert_eq!(b.as_slice(), &[0b0000000]); + assert_eq!(b.len(), 6); + + let mut b = MutableBitmap::from(&[true]); + b.extend_constant(6, false); + assert_eq!(b.as_slice(), &[0b00000001]); + assert_eq!(b.len(), 1 + 6); + + let mut b = MutableBitmap::from(&[true]); + b.extend_constant(9, false); + assert_eq!(b.as_slice(), &[0b0000001, 0b00000000]); + assert_eq!(b.len(), 1 + 9); + + let mut b = MutableBitmap::from(&[true, true, true, true]); + b.extend_constant(2, false); + assert_eq!(b.as_slice(), &[0b00001111]); + assert_eq!(b.len(), 4 + 2); +}