diff --git a/src/array/primitive/mutable.rs b/src/array/primitive/mutable.rs index ee330188126..64d8c245abb 100644 --- a/src/array/primitive/mutable.rs +++ b/src/array/primitive/mutable.rs @@ -164,14 +164,57 @@ impl MutablePrimitiveArray { if let Some(validity) = self.validity.as_mut() { extend_trusted_len_unzip(iterator, validity, &mut self.values) } else { - let mut validity = - MutableBitmap::from_trusted_len_iter(std::iter::repeat(true).take(self.len())); + 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); } } } + /// Extends the [`MutablePrimitiveArray`] 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, + { + unsafe { self.extend_trusted_len_values_unchecked(iterator) } + } + + /// Extends the [`MutablePrimitiveArray`] 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, + { + self.values.extend_from_trusted_len_iter_unchecked(iterator); + self.update_all_valid(); + } + + #[inline] + /// Extends the [`MutablePrimitiveArray`] from a slice + pub fn extend_from_slice(&mut self, items: &[T]) { + self.values.extend_from_slice(items); + self.update_all_valid(); + } + + fn update_all_valid(&mut self) { + // get len before mutable borrow + let len = self.len(); + if let Some(validity) = self.validity.as_mut() { + validity.extend_constant(len - validity.len(), true); + } else { + let mut validity = MutableBitmap::new(); + validity.extend_constant(len, true); + if validity.null_count() > 0 { + self.validity = Some(validity); + } + } + } fn init_validity(&mut self) { let mut validity = MutableBitmap::new(); @@ -449,10 +492,10 @@ pub(crate) unsafe fn extend_trusted_len_unzip( I: Iterator>, { let (_, upper) = iterator.size_hint(); - let len = upper.expect("trusted_len_unzip requires an upper limit"); + let additional = upper.expect("trusted_len_unzip requires an upper limit"); - validity.reserve(len); - buffer.reserve(len); + validity.reserve(additional); + buffer.reserve(additional); for item in iterator { let item = if let Some(item) = item { diff --git a/src/bitmap/mutable.rs b/src/bitmap/mutable.rs index d3eae6689e5..4826ef7872c 100644 --- a/src/bitmap/mutable.rs +++ b/src/bitmap/mutable.rs @@ -98,10 +98,8 @@ impl MutableBitmap { if self.length % 8 == 0 { self.buffer.push_unchecked(0); } - if value { - let byte = self.buffer.as_mut_slice().last_mut().unwrap(); - *byte = set(*byte, self.length % 8, true); - }; + let byte = self.buffer.as_mut_slice().last_mut().unwrap(); + *byte = set(*byte, self.length % 8, value); self.length += 1; } diff --git a/src/compute/like.rs b/src/compute/like.rs index bcd2ab71a96..1160721cef7 100644 --- a/src/compute/like.rs +++ b/src/compute/like.rs @@ -232,7 +232,10 @@ fn a_like_binary_scalar bool>( ) -> Result { let validity = lhs.validity(); let pattern = std::str::from_utf8(rhs).map_err(|e| { - ArrowError::InvalidArgumentError(format!("Unable to convert the LIKE pattern to string: {}", e)) + ArrowError::InvalidArgumentError(format!( + "Unable to convert the LIKE pattern to string: {}", + e + )) })?; let values = if !pattern.contains(is_like_pattern) { diff --git a/tests/it/array/primitive/mutable.rs b/tests/it/array/primitive/mutable.rs index f5f4cd381d7..ecc1fd8f994 100644 --- a/tests/it/array/primitive/mutable.rs +++ b/tests/it/array/primitive/mutable.rs @@ -126,6 +126,38 @@ fn extend_trusted_len() { assert_eq!(a.values(), &MutableBuffer::::from([1, 2, 0, 4])); } +#[test] +fn extend_trusted_len_values() { + let mut a = MutablePrimitiveArray::::new(); + a.extend_trusted_len_values(vec![1, 2, 3].into_iter()); + assert_eq!(a.validity(), &None); + assert_eq!(a.values(), &MutableBuffer::::from([1, 2, 3])); + + let mut a = MutablePrimitiveArray::::new(); + a.push(None); + a.extend_trusted_len_values(vec![1, 2].into_iter()); + assert_eq!( + a.validity(), + &Some(MutableBitmap::from([false, true, true])) + ); +} + +#[test] +fn extend_from_slice() { + let mut a = MutablePrimitiveArray::::new(); + a.extend_from_slice(&[1, 2, 3]); + assert_eq!(a.validity(), &None); + assert_eq!(a.values(), &MutableBuffer::::from([1, 2, 3])); + + let mut a = MutablePrimitiveArray::::new(); + a.push(None); + a.extend_from_slice(&[1, 2]); + assert_eq!( + a.validity(), + &Some(MutableBitmap::from([false, true, true])) + ); +} + #[test] fn set_validity() { let mut a = MutablePrimitiveArray::::new();