diff --git a/src/array/boolean/mutable.rs b/src/array/boolean/mutable.rs index a776b7a12bc..b27e7b2bb2d 100644 --- a/src/array/boolean/mutable.rs +++ b/src/array/boolean/mutable.rs @@ -104,6 +104,71 @@ impl MutableBooleanArray { } } + /// Extends the [`MutableBooleanArray`] from an iterator of values of trusted len. + /// This differs from `extend_trusted_len` which accepts in iterator of optional values. + #[inline] + pub fn extend_trusted_len_values(&mut self, iterator: I) + where + I: TrustedLen, + { + // Safety: `I` is `TrustedLen` + unsafe { self.extend_trusted_len_values_unchecked(iterator) } + } + + /// Extends the [`MutableBooleanArray`] from an iterator of values of trusted len. + /// This differs from `extend_trusted_len_unchecked`, which accepts in iterator of optional values. + /// # Safety + /// The iterator must be trusted len. + #[inline] + pub unsafe fn extend_trusted_len_values_unchecked(&mut self, iterator: I) + where + I: Iterator, + { + let (_, upper) = iterator.size_hint(); + let additional = + upper.expect("extend_trusted_len_values_unchecked requires an upper limit"); + + if let Some(validity) = self.validity.as_mut() { + validity.extend_constant(additional, true); + } + + self.values.extend_from_trusted_len_iter_unchecked(iterator) + } + + /// Extends the [`MutableBooleanArray`] from an iterator of trusted len. + #[inline] + pub fn extend_trusted_len(&mut self, iterator: I) + where + P: std::borrow::Borrow, + I: TrustedLen>, + { + // Safety: `I` is `TrustedLen` + unsafe { self.extend_trusted_len_unchecked(iterator) } + } + + /// Extends the [`MutableBooleanArray`] from an iterator of trusted len. + /// # Safety + /// The iterator must be trusted len. + #[inline] + pub unsafe fn extend_trusted_len_unchecked(&mut self, iterator: I) + where + P: std::borrow::Borrow, + I: Iterator>, + { + if let Some(validity) = self.validity.as_mut() { + extend_trusted_len_unzip(iterator, validity, &mut self.values); + } else { + let mut validity = MutableBitmap::new(); + validity.extend_constant(self.len(), true); + + extend_trusted_len_unzip(iterator, &mut validity, &mut self.values); + + if validity.null_count() > 0 { + self.validity = Some(validity); + } + } + } + fn init_validity(&mut self) { let mut validity = MutableBitmap::new(); validity.extend_constant(self.len(), true); @@ -181,12 +246,6 @@ impl MutableBooleanArray { { let (validity, values) = trusted_len_unzip(iterator); - let validity = if validity.null_count() > 0 { - Some(validity) - } else { - None - }; - Self::from_data(DataType::Boolean, values, validity) } @@ -197,6 +256,7 @@ impl MutableBooleanArray { P: std::borrow::Borrow, I: TrustedLen>, { + // Safety: `I` is `TrustedLen` unsafe { Self::from_trusted_len_iter_unchecked(iterator) } } @@ -230,6 +290,7 @@ impl MutableBooleanArray { P: std::borrow::Borrow, I: TrustedLen, E>>, { + // Safety: `I` is `TrustedLen` unsafe { Self::try_from_trusted_len_iter_unchecked(iterator) } } } @@ -240,33 +301,63 @@ impl MutableBooleanArray { /// # Safety /// The caller must ensure that `iterator` is `TrustedLen`. #[inline] -pub(crate) unsafe fn trusted_len_unzip(iterator: I) -> (MutableBitmap, MutableBitmap) +pub(crate) unsafe fn trusted_len_unzip(iterator: I) -> (Option, MutableBitmap) where P: std::borrow::Borrow, I: Iterator>, +{ + let mut validity = MutableBitmap::new(); + let mut values = MutableBitmap::new(); + + extend_trusted_len_unzip(iterator, &mut validity, &mut values); + + let validity = if validity.null_count() > 0 { + Some(validity) + } else { + None + }; + + (validity, values) +} + +/// Extends validity [`MutableBitmap`] and values [`MutableBitmap`] from an iterator of `Option`. +/// # Safety +/// The caller must ensure that `iterator` is `TrustedLen`. +#[inline] +pub(crate) unsafe fn extend_trusted_len_unzip( + iterator: I, + validity: &mut MutableBitmap, + values: &mut MutableBitmap, +) where + P: std::borrow::Borrow, + I: Iterator>, { let (_, upper) = iterator.size_hint(); - let len = upper.expect("trusted_len_unzip requires an upper limit"); + let additional = upper.expect("extend_trusted_len_unzip requires an upper limit"); - let mut validity = MutableBitmap::with_capacity(len); - let mut values = MutableBitmap::with_capacity(len); + // Length of the array before new values are pushed, + // variable created for assertion post operation + let pre_length = values.len(); + + validity.reserve(additional); + values.reserve(additional); for item in iterator { let item = if let Some(item) = item { - validity.push(true); + validity.push_unchecked(true); *item.borrow() } else { - validity.push(false); - false + validity.push_unchecked(false); + bool::default() }; - values.push(item); + values.push_unchecked(item); } - assert_eq!( + + debug_assert_eq!( values.len(), - len, + pre_length + additional, "Trusted iterator length was not accurately reported" ); - (validity, values) } /// # Safety diff --git a/tests/it/array/boolean/mutable.rs b/tests/it/array/boolean/mutable.rs index 6b4522b99b9..e56ccadb72c 100644 --- a/tests/it/array/boolean/mutable.rs +++ b/tests/it/array/boolean/mutable.rs @@ -64,3 +64,36 @@ fn reserve() { assert!(a.validity().unwrap().capacity() > 0); assert!(a.values().capacity() > 0) } + +#[test] +fn extend_trusted_len() { + let mut a = MutableBooleanArray::new(); + + a.extend_trusted_len(vec![Some(true), Some(false)].into_iter()); + assert_eq!(a.validity(), None); + + a.extend_trusted_len(vec![None, Some(true)].into_iter()); + assert_eq!( + a.validity(), + Some(&MutableBitmap::from([true, true, false, true])) + ); + assert_eq!(a.values(), &MutableBitmap::from([true, false, false, true])); +} + +#[test] +fn extend_trusted_len_values() { + let mut a = MutableBooleanArray::new(); + + a.extend_trusted_len_values(vec![true, true, false].into_iter()); + assert_eq!(a.validity(), None); + assert_eq!(a.values(), &MutableBitmap::from([true, true, false])); + + let mut a = MutableBooleanArray::new(); + a.push(None); + a.extend_trusted_len_values(vec![true, false].into_iter()); + assert_eq!( + a.validity(), + Some(&MutableBitmap::from([false, true, true])) + ); + assert_eq!(a.values(), &MutableBitmap::from([false, true, false])); +}