diff --git a/src/array/binary/mod.rs b/src/array/binary/mod.rs index 72341bdfeeb..e7a4e9fe8c0 100644 --- a/src/array/binary/mod.rs +++ b/src/array/binary/mod.rs @@ -114,7 +114,7 @@ impl BinaryArray { /// Creates a new [`BinaryArray`] from a slice of optional `&[u8]`. // Note: this can't be `impl From` because Rust does not allow double `AsRef` on it. pub fn from, P: AsRef<[Option]>>(slice: P) -> Self { - Self::from_trusted_len_iter(slice.as_ref().iter().map(|x| x.as_ref())) + MutableBinaryArray::::from(slice).into() } /// Returns an iterator of `Option<&[u8]>` over every element of this array. diff --git a/src/array/binary/mutable.rs b/src/array/binary/mutable.rs index 1d52f1cbafd..15035aed0ed 100644 --- a/src/array/binary/mutable.rs +++ b/src/array/binary/mutable.rs @@ -1,21 +1,24 @@ use std::{iter::FromIterator, sync::Arc}; use crate::{ - array::{Array, MutableArray, Offset, TryExtend, TryPush}, - bitmap::{Bitmap, MutableBitmap}, + array::{Array, MutableArray, Offset, TryExtend, TryExtendFromSelf, TryPush}, + bitmap::{ + utils::{BitmapIter, ZipValidity}, + Bitmap, MutableBitmap, + }, datatypes::DataType, error::{Error, Result}, trusted_len::TrustedLen, }; -use super::{BinaryArray, MutableBinaryValuesArray}; +use super::{BinaryArray, MutableBinaryValuesArray, MutableBinaryValuesIter}; use crate::array::physical_binary::*; /// The Arrow's equivalent to `Vec>>`. /// Converting a [`MutableBinaryArray`] into a [`BinaryArray`] is `O(1)`. /// # Implementation /// This struct does not allocate a validity until one is required (i.e. push a null to it). -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct MutableBinaryArray { values: MutableBinaryValuesArray, validity: Option, @@ -96,6 +99,12 @@ impl MutableBinaryArray { Self { values, validity } } + /// Creates a new [`MutableBinaryArray`] from a slice of optional `&[u8]`. + // Note: this can't be `impl From` because Rust does not allow double `AsRef` on it. + pub fn from, P: AsRef<[Option]>>(slice: P) -> Self { + Self::from_trusted_len_iter(slice.as_ref().iter().map(|x| x.as_ref())) + } + fn default_data_type() -> DataType { BinaryArray::::default_data_type() } @@ -193,6 +202,16 @@ impl MutableBinaryArray { pub fn offsets(&self) -> &Vec { self.values.offsets() } + + /// Returns an iterator of `Option<&[u8]>` + pub fn iter(&self) -> ZipValidity<&[u8], MutableBinaryValuesIter, BitmapIter> { + ZipValidity::new(self.values_iter(), self.validity.as_ref().map(|x| x.iter())) + } + + /// Returns an iterator over the values of this array + pub fn values_iter(&self) -> MutableBinaryValuesIter { + self.values.iter() + } } impl MutableArray for MutableBinaryArray { @@ -478,3 +497,17 @@ impl> TryPush> for MutableBinaryArray { Ok(()) } } + +impl PartialEq for MutableBinaryArray { + fn eq(&self, other: &Self) -> bool { + self.iter().eq(other.iter()) + } +} + +impl TryExtendFromSelf for MutableBinaryArray { + fn try_extend_from_self(&mut self, other: &Self) -> Result<()> { + extend_validity(self.len(), &mut self.validity, &other.validity); + + self.values.try_extend_from_self(&other.values) + } +} diff --git a/src/array/binary/mutable_values.rs b/src/array/binary/mutable_values.rs index 18518889c6e..53a43d69c7a 100644 --- a/src/array/binary/mutable_values.rs +++ b/src/array/binary/mutable_values.rs @@ -3,7 +3,8 @@ use std::{iter::FromIterator, sync::Arc}; use crate::{ array::{ specification::{check_offsets_minimal, try_check_offsets}, - Array, ArrayAccessor, ArrayValuesIter, MutableArray, Offset, TryExtend, TryPush, + Array, ArrayAccessor, ArrayValuesIter, MutableArray, Offset, TryExtend, TryExtendFromSelf, + TryPush, }, bitmap::MutableBitmap, datatypes::DataType, @@ -408,3 +409,10 @@ unsafe impl<'a, O: Offset> ArrayAccessor<'a> for MutableBinaryValuesArray { self.len() } } + +impl TryExtendFromSelf for MutableBinaryValuesArray { + fn try_extend_from_self(&mut self, other: &Self) -> Result<()> { + self.values.extend_from_slice(&other.values); + try_extend_offsets(&mut self.offsets, &other.offsets) + } +} diff --git a/src/array/boolean/mutable.rs b/src/array/boolean/mutable.rs index 83fac4b7ef5..f6336cda3a4 100644 --- a/src/array/boolean/mutable.rs +++ b/src/array/boolean/mutable.rs @@ -2,7 +2,10 @@ use std::iter::FromIterator; use std::sync::Arc; use crate::{ - array::{Array, MutableArray, TryExtend, TryPush}, + array::{ + physical_binary::extend_validity, Array, MutableArray, TryExtend, TryExtendFromSelf, + TryPush, + }, bitmap::MutableBitmap, datatypes::{DataType, PhysicalType}, error::Result, @@ -15,7 +18,7 @@ use super::BooleanArray; /// Converting a [`MutableBooleanArray`] into a [`BooleanArray`] is `O(1)`. /// # Implementation /// This struct does not allocate a validity until one is required (i.e. push a null to it). -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct MutableBooleanArray { data_type: DataType, values: MutableBitmap, @@ -533,3 +536,13 @@ impl PartialEq for MutableBooleanArray { self.iter().eq(other.iter()) } } + +impl TryExtendFromSelf for MutableBooleanArray { + fn try_extend_from_self(&mut self, other: &Self) -> Result<()> { + extend_validity(self.len(), &mut self.validity, &other.validity); + + let slice = other.values.as_slice(); + self.values.extend_from_slice(slice, 0, other.values.len()); + Ok(()) + } +} diff --git a/src/array/fixed_size_binary/mod.rs b/src/array/fixed_size_binary/mod.rs index d467d2c3c00..441b58d2fcc 100644 --- a/src/array/fixed_size_binary/mod.rs +++ b/src/array/fixed_size_binary/mod.rs @@ -1,9 +1,4 @@ -use crate::{ - bitmap::{Bitmap, MutableBitmap}, - buffer::Buffer, - datatypes::DataType, - error::Error, -}; +use crate::{bitmap::Bitmap, buffer::Buffer, datatypes::DataType, error::Error}; use super::Array; @@ -322,18 +317,7 @@ impl FixedSizeBinaryArray { /// Creates a new [`FixedSizeBinaryArray`] from a slice of optional `[u8]`. // Note: this can't be `impl From` because Rust does not allow double `AsRef` on it. pub fn from]>>(slice: P) -> Self { - let values = slice - .as_ref() - .iter() - .copied() - .flat_map(|x| x.unwrap_or([0; N])) - .collect::>(); - let validity = slice - .as_ref() - .iter() - .map(|x| x.is_some()) - .collect::(); - Self::new(DataType::FixedSizeBinary(N), values.into(), validity.into()) + MutableFixedSizeBinaryArray::from(slice).into() } } diff --git a/src/array/fixed_size_binary/mutable.rs b/src/array/fixed_size_binary/mutable.rs index 3a02158cc88..9e3cce628e3 100644 --- a/src/array/fixed_size_binary/mutable.rs +++ b/src/array/fixed_size_binary/mutable.rs @@ -1,7 +1,7 @@ use std::sync::Arc; use crate::{ - array::{Array, MutableArray}, + array::{physical_binary::extend_validity, Array, MutableArray, TryExtendFromSelf}, bitmap::MutableBitmap, datatypes::DataType, error::{Error, Result}, @@ -13,7 +13,7 @@ use super::{FixedSizeBinaryArray, FixedSizeBinaryValues}; /// Converting a [`MutableFixedSizeBinaryArray`] into a [`FixedSizeBinaryArray`] is `O(1)`. /// # Implementation /// This struct does not allocate a validity until one is required (i.e. push a null to it). -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct MutableFixedSizeBinaryArray { data_type: DataType, size: usize, @@ -73,6 +73,23 @@ impl MutableFixedSizeBinaryArray { ) } + /// Creates a new [`MutableFixedSizeBinaryArray`] from a slice of optional `[u8]`. + // Note: this can't be `impl From` because Rust does not allow double `AsRef` on it. + pub fn from]>>(slice: P) -> Self { + let values = slice + .as_ref() + .iter() + .copied() + .flat_map(|x| x.unwrap_or([0; N])) + .collect::>(); + let validity = slice + .as_ref() + .iter() + .map(|x| x.is_some()) + .collect::(); + Self::from_data(DataType::FixedSizeBinary(N), values, validity.into()) + } + /// tries to push a new entry to [`MutableFixedSizeBinaryArray`]. /// # Error /// Errors iff the size of `value` is not equal to its own size. @@ -278,3 +295,13 @@ impl PartialEq for MutableFixedSizeBinaryArray { self.iter().eq(other.iter()) } } + +impl TryExtendFromSelf for MutableFixedSizeBinaryArray { + fn try_extend_from_self(&mut self, other: &Self) -> Result<()> { + extend_validity(self.len(), &mut self.validity, &other.validity); + + let slice = other.values.as_slice(); + self.values.extend_from_slice(slice); + Ok(()) + } +} diff --git a/src/array/fixed_size_list/mutable.rs b/src/array/fixed_size_list/mutable.rs index 0cba41a1e90..17d401b7bab 100644 --- a/src/array/fixed_size_list/mutable.rs +++ b/src/array/fixed_size_list/mutable.rs @@ -1,7 +1,10 @@ use std::sync::Arc; use crate::{ - array::{Array, MutableArray, TryExtend, TryPush}, + array::{ + physical_binary::extend_validity, Array, MutableArray, TryExtend, TryExtendFromSelf, + TryPush, + }, bitmap::MutableBitmap, datatypes::{DataType, Field}, error::{Error, Result}, @@ -10,7 +13,7 @@ use crate::{ use super::FixedSizeListArray; /// The mutable version of [`FixedSizeListArray`]. -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct MutableFixedSizeListArray { data_type: DataType, size: usize, @@ -64,6 +67,11 @@ impl MutableFixedSizeListArray { self.size } + /// The length of this array + pub fn len(&self) -> usize { + self.values.len() / self.size + } + /// The inner values pub fn values(&self) -> &M { &self.values @@ -210,3 +218,14 @@ where Ok(()) } } + +impl TryExtendFromSelf for MutableFixedSizeListArray +where + M: MutableArray + TryExtendFromSelf, +{ + fn try_extend_from_self(&mut self, other: &Self) -> Result<()> { + extend_validity(self.len(), &mut self.validity, &other.validity); + + self.values.try_extend_from_self(&other.values) + } +} diff --git a/src/array/list/iterator.rs b/src/array/list/iterator.rs index d0b7d22455b..82b5c7dca5f 100644 --- a/src/array/list/iterator.rs +++ b/src/array/list/iterator.rs @@ -43,3 +43,27 @@ impl<'a, O: Offset> ListArray { ListValuesIter::new(self) } } + +struct Iter>> { + current: i32, + offsets: std::vec::IntoIter, + values: I, +} + +impl> + Clone> Iterator for Iter { + type Item = Option>>; + + fn next(&mut self) -> Option { + let next = self.offsets.next(); + next.map(|next| { + let length = next - self.current; + let iter = self + .values + .clone() + .skip(self.current as usize) + .take(length as usize); + self.current = next; + Some(iter) + }) + } +} diff --git a/src/array/list/mutable.rs b/src/array/list/mutable.rs index b4f3c9d1d3b..966aac04651 100644 --- a/src/array/list/mutable.rs +++ b/src/array/list/mutable.rs @@ -1,7 +1,11 @@ use std::sync::Arc; use crate::{ - array::{specification::try_check_offsets, Array, MutableArray, Offset, TryExtend, TryPush}, + array::{ + physical_binary::{extend_validity, try_extend_offsets}, + specification::try_check_offsets, + Array, MutableArray, Offset, TryExtend, TryExtendFromSelf, TryPush, + }, bitmap::MutableBitmap, datatypes::{DataType, Field}, error::{Error, Result}, @@ -11,7 +15,7 @@ use crate::{ use super::ListArray; /// The mutable version of [`ListArray`]. -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct MutableListArray { data_type: DataType, offsets: Vec, @@ -71,6 +75,8 @@ where I: IntoIterator>, { fn try_extend>>(&mut self, iter: II) -> Result<()> { + let iter = iter.into_iter(); + self.reserve(iter.size_hint().0); for items in iter { self.try_push(items)?; } @@ -97,6 +103,20 @@ where } } +impl TryExtendFromSelf for MutableListArray +where + O: Offset, + M: MutableArray + TryExtendFromSelf, +{ + fn try_extend_from_self(&mut self, other: &Self) -> Result<()> { + extend_validity(self.len(), &mut self.validity, &other.validity); + + self.values.try_extend_from_self(&other.values)?; + + try_extend_offsets(&mut self.offsets, &other.offsets) + } +} + impl MutableListArray { /// Creates a new [`MutableListArray`] from a [`MutableArray`] and capacity. pub fn new_from(values: M, data_type: DataType, capacity: usize) -> Self { @@ -222,6 +242,12 @@ impl MutableListArray { } } + /// Returns the length of this array + #[inline] + pub fn len(&self) -> usize { + self.offsets.len() - 1 + } + /// The values pub fn mut_values(&mut self) -> &mut M { &mut self.values @@ -279,10 +305,6 @@ impl MutableListArray { validity.shrink_to_fit() } } - - fn len(&self) -> usize { - self.offsets.len() - 1 - } } impl MutableArray for MutableListArray { diff --git a/src/array/mod.rs b/src/array/mod.rs index c303f2d6120..45e83a80803 100644 --- a/src/array/mod.rs +++ b/src/array/mod.rs @@ -470,6 +470,13 @@ pub trait TryPush { fn try_push(&mut self, item: A) -> Result<()>; } +/// A trait describing the ability of a struct to extend from a reference of itself. +/// Specialization of [`TryExtend`]. +pub trait TryExtendFromSelf { + /// Tries to extend itself with elements from `other`, failing only on overflow. + fn try_extend_from_self(&mut self, other: &Self) -> Result<()>; +} + /// Trait that [`BinaryArray`] and [`Utf8Array`] implement for the purposes of DRY. /// # Safety /// The implementer must ensure that diff --git a/src/array/physical_binary.rs b/src/array/physical_binary.rs index 3e3a9c5e901..bf883f4e6e8 100644 --- a/src/array/physical_binary.rs +++ b/src/array/physical_binary.rs @@ -1,5 +1,6 @@ use crate::array::Offset; use crate::bitmap::MutableBitmap; +use crate::error::Error; /// # Safety /// The caller must ensure that `iterator` is `TrustedLen`. @@ -280,3 +281,39 @@ where } (offsets, values) } + +/// Extends `offsets` with all offsets from `other` +#[inline] +pub(crate) fn try_extend_offsets(offsets: &mut Vec, other: &[O]) -> Result<(), Error> +where + O: Offset, +{ + let lengths = other.windows(2).map(|w| w[1] - w[0]); + let mut last = *offsets.last().unwrap(); + + offsets.reserve(other.len() - 1); + for length in lengths { + let r = last.checked_add(&length).ok_or(Error::Overflow)?; + last += length; + offsets.push(r) + } + Ok(()) +} + +/// Extends `validity` with all items from `other` +pub(crate) fn extend_validity( + length: usize, + validity: &mut Option, + other: &Option, +) { + if let Some(other) = other { + if let Some(validity) = validity { + let slice = other.as_slice(); + validity.extend_from_slice(slice, 0, other.len()) + } else { + let mut new_validity = MutableBitmap::from_len_set(length); + new_validity.extend_from_slice(other.as_slice(), 0, other.len()); + *validity = Some(new_validity); + } + } +} diff --git a/src/array/primitive/mutable.rs b/src/array/primitive/mutable.rs index 5f2014db319..c71c41c164e 100644 --- a/src/array/primitive/mutable.rs +++ b/src/array/primitive/mutable.rs @@ -1,5 +1,7 @@ use std::{iter::FromIterator, sync::Arc}; +use crate::array::physical_binary::extend_validity; +use crate::array::TryExtendFromSelf; use crate::bitmap::Bitmap; use crate::{ array::{Array, MutableArray, TryExtend, TryPush}, @@ -655,3 +657,13 @@ impl PartialEq for MutablePrimitiveArray { self.iter().eq(other.iter()) } } + +impl TryExtendFromSelf for MutablePrimitiveArray { + fn try_extend_from_self(&mut self, other: &Self) -> Result<()> { + extend_validity(self.len(), &mut self.validity, &other.validity); + + let slice = other.values.as_slice(); + self.values.extend_from_slice(slice); + Ok(()) + } +} diff --git a/src/array/utf8/mod.rs b/src/array/utf8/mod.rs index 965eb6552f1..f91e1466451 100644 --- a/src/array/utf8/mod.rs +++ b/src/array/utf8/mod.rs @@ -128,7 +128,7 @@ impl Utf8Array { /// A convenience method that uses [`Self::from_trusted_len_iter`]. // Note: this can't be `impl From` because Rust does not allow double `AsRef` on it. pub fn from, P: AsRef<[Option]>>(slice: P) -> Self { - Self::from_trusted_len_iter(slice.as_ref().iter().map(|x| x.as_ref())) + MutableUtf8Array::::from(slice).into() } /// Returns an iterator of `Option<&str>` diff --git a/src/array/utf8/mutable.rs b/src/array/utf8/mutable.rs index c3564144281..27d28dc64b2 100644 --- a/src/array/utf8/mutable.rs +++ b/src/array/utf8/mutable.rs @@ -1,6 +1,6 @@ use std::{iter::FromIterator, sync::Arc}; -use crate::array::physical_binary::*; +use crate::array::{physical_binary::*, TryExtendFromSelf}; use crate::{ array::{Array, MutableArray, Offset, TryExtend, TryPush}, bitmap::{ @@ -16,7 +16,7 @@ use super::{MutableUtf8ValuesArray, MutableUtf8ValuesIter, StrAsBytes, Utf8Array /// A [`MutableArray`] that builds a [`Utf8Array`]. It differs /// from [`MutableUtf8ValuesArray`] in that it can build nullable [`Utf8Array`]s. -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct MutableUtf8Array { values: MutableUtf8ValuesArray, validity: Option, @@ -126,6 +126,12 @@ impl MutableUtf8Array { Self::try_new(data_type, offsets, values, validity).unwrap() } + /// Creates a new [`MutableUtf8Array`] from a slice of optional `&[u8]`. + // Note: this can't be `impl From` because Rust does not allow double `AsRef` on it. + pub fn from, P: AsRef<[Option]>>(slice: P) -> Self { + Self::from_trusted_len_iter(slice.as_ref().iter().map(|x| x.as_ref())) + } + fn default_data_type() -> DataType { Utf8Array::::default_data_type() } @@ -544,3 +550,17 @@ impl> TryPush> for MutableUtf8Array { Ok(()) } } + +impl PartialEq for MutableUtf8Array { + fn eq(&self, other: &Self) -> bool { + self.iter().eq(other.iter()) + } +} + +impl TryExtendFromSelf for MutableUtf8Array { + fn try_extend_from_self(&mut self, other: &Self) -> Result<()> { + extend_validity(self.len(), &mut self.validity, &other.validity); + + self.values.try_extend_from_self(&other.values) + } +} diff --git a/src/array/utf8/mutable_values.rs b/src/array/utf8/mutable_values.rs index 1f393a931d1..5e33144452f 100644 --- a/src/array/utf8/mutable_values.rs +++ b/src/array/utf8/mutable_values.rs @@ -3,7 +3,7 @@ use std::{iter::FromIterator, sync::Arc}; use crate::{ array::{ specification::{check_offsets_minimal, try_check_offsets_and_utf8}, - Array, ArrayValuesIter, MutableArray, Offset, TryExtend, TryPush, + Array, ArrayValuesIter, MutableArray, Offset, TryExtend, TryExtendFromSelf, TryPush, }, bitmap::MutableBitmap, datatypes::DataType, @@ -407,3 +407,10 @@ impl> TryPush for MutableUtf8ValuesArray { Ok(()) } } + +impl TryExtendFromSelf for MutableUtf8ValuesArray { + fn try_extend_from_self(&mut self, other: &Self) -> Result<()> { + self.values.extend_from_slice(&other.values); + try_extend_offsets(&mut self.offsets, &other.offsets) + } +} diff --git a/tests/it/array/binary/mutable.rs b/tests/it/array/binary/mutable.rs index e363b47bcf8..74addaac939 100644 --- a/tests/it/array/binary/mutable.rs +++ b/tests/it/array/binary/mutable.rs @@ -1,6 +1,6 @@ use std::ops::Deref; -use arrow2::array::{BinaryArray, MutableArray, MutableBinaryArray}; +use arrow2::array::{BinaryArray, MutableArray, MutableBinaryArray, TryExtendFromSelf}; use arrow2::bitmap::Bitmap; use arrow2::error::Error; @@ -151,3 +151,15 @@ fn extend_trusted_len() { Some(&Bitmap::from_u8_slice(&[0b00001011], 4)) ); } + +#[test] +fn extend_from_self() { + let mut a = MutableBinaryArray::::from(&[Some(b"aa"), None]); + + a.try_extend_from_self(&a.clone()).unwrap(); + + assert_eq!( + a, + MutableBinaryArray::::from([Some(b"aa"), None, Some(b"aa"), None]) + ); +} diff --git a/tests/it/array/boolean/mutable.rs b/tests/it/array/boolean/mutable.rs index 173472106a0..de960b7a77f 100644 --- a/tests/it/array/boolean/mutable.rs +++ b/tests/it/array/boolean/mutable.rs @@ -1,4 +1,4 @@ -use arrow2::array::{MutableArray, MutableBooleanArray}; +use arrow2::array::{MutableArray, MutableBooleanArray, TryExtendFromSelf}; use arrow2::bitmap::MutableBitmap; use arrow2::datatypes::DataType; use arrow2::error::Result; @@ -162,3 +162,15 @@ fn shrink_to_fit() { a.shrink_to_fit(); assert_eq!(a.capacity(), 8); } + +#[test] +fn extend_from_self() { + let mut a = MutableBooleanArray::from([Some(true), None]); + + a.try_extend_from_self(&a.clone()).unwrap(); + + assert_eq!( + a, + MutableBooleanArray::from([Some(true), None, Some(true), None]) + ); +} diff --git a/tests/it/array/fixed_size_binary/mutable.rs b/tests/it/array/fixed_size_binary/mutable.rs index 068ed0681c4..3cffcaa899a 100644 --- a/tests/it/array/fixed_size_binary/mutable.rs +++ b/tests/it/array/fixed_size_binary/mutable.rs @@ -154,3 +154,15 @@ fn shrink_to_fit_and_capacity() { array.shrink_to_fit(); assert_eq!(array.capacity(), 1); } + +#[test] +fn extend_from_self() { + let mut a = MutableFixedSizeBinaryArray::from([Some([1u8, 2u8]), None]); + + a.try_extend_from_self(&a.clone()).unwrap(); + + assert_eq!( + a, + MutableFixedSizeBinaryArray::from([Some([1u8, 2u8]), None, Some([1u8, 2u8]), None]) + ); +} diff --git a/tests/it/array/fixed_size_list/mutable.rs b/tests/it/array/fixed_size_list/mutable.rs index 39826645ba7..f7a8784dfce 100644 --- a/tests/it/array/fixed_size_list/mutable.rs +++ b/tests/it/array/fixed_size_list/mutable.rs @@ -63,3 +63,26 @@ fn new_with_field() { let expected = Int32Array::from(vec![None, None, None]); assert_eq!(a, &expected) } + +#[test] +fn extend_from_self() { + let data = vec![ + Some(vec![Some(1i32), Some(2), Some(3)]), + None, + Some(vec![Some(4), None, Some(6)]), + ]; + let mut a = MutableFixedSizeListArray::new(MutablePrimitiveArray::::new(), 3); + a.try_extend(data.clone()).unwrap(); + + a.try_extend_from_self(&a.clone()).unwrap(); + let a: FixedSizeListArray = a.into(); + + let mut expected = data.clone(); + expected.extend(data); + + let mut b = MutableFixedSizeListArray::new(MutablePrimitiveArray::::new(), 3); + b.try_extend(expected).unwrap(); + let b: FixedSizeListArray = b.into(); + + assert_eq!(a, b); +} diff --git a/tests/it/array/list/mutable.rs b/tests/it/array/list/mutable.rs index 64b593bc2e0..63cf3f8c1e5 100644 --- a/tests/it/array/list/mutable.rs +++ b/tests/it/array/list/mutable.rs @@ -48,3 +48,26 @@ fn push() { assert_eq!(array.offsets().as_ref(), [0, 3]); assert_eq!(array.validity(), None); } + +#[test] +fn extend_from_self() { + let data = vec![ + Some(vec![Some(1i32), Some(2), Some(3)]), + None, + Some(vec![Some(4), None, Some(6)]), + ]; + let mut a = MutableListArray::>::new(); + a.try_extend(data.clone()).unwrap(); + + a.try_extend_from_self(&a.clone()).unwrap(); + let a: ListArray = a.into(); + + let mut expected = data.clone(); + expected.extend(data); + + let mut b = MutableListArray::>::new(); + b.try_extend(expected).unwrap(); + let b: ListArray = b.into(); + + assert_eq!(a, b); +} diff --git a/tests/it/array/primitive/mutable.rs b/tests/it/array/primitive/mutable.rs index 762c1a24045..69c7e47dafd 100644 --- a/tests/it/array/primitive/mutable.rs +++ b/tests/it/array/primitive/mutable.rs @@ -10,7 +10,7 @@ use std::iter::FromIterator; fn from_and_into_data() { let a = MutablePrimitiveArray::from_data( DataType::Int32, - Vec::from([1i32, 0]), + vec![1i32, 0], Some(MutableBitmap::from([true, false])), ); assert_eq!(a.len(), 2); @@ -30,7 +30,7 @@ fn from_vec() { fn to() { let a = MutablePrimitiveArray::from_data( DataType::Int32, - Vec::from([1i32, 0]), + vec![1i32, 0], Some(MutableBitmap::from([true, false])), ); let a = a.to(DataType::Date32); @@ -41,7 +41,7 @@ fn to() { fn values_mut_slice() { let mut a = MutablePrimitiveArray::from_data( DataType::Int32, - Vec::from([1i32, 0]), + vec![1i32, 0], Some(MutableBitmap::from([true, false])), ); let values = a.values_mut_slice(); @@ -316,3 +316,15 @@ fn wrong_data_type() { let values = vec![1u8]; MutablePrimitiveArray::from_data(DataType::Utf8, values, None); } + +#[test] +fn extend_from_self() { + let mut a = MutablePrimitiveArray::from([Some(1), None]); + + a.try_extend_from_self(&a.clone()).unwrap(); + + assert_eq!( + a, + MutablePrimitiveArray::from([Some(1), None, Some(1), None]) + ); +} diff --git a/tests/it/array/utf8/mutable.rs b/tests/it/array/utf8/mutable.rs index 4a12a459abe..58932d1f6b9 100644 --- a/tests/it/array/utf8/mutable.rs +++ b/tests/it/array/utf8/mutable.rs @@ -1,4 +1,4 @@ -use arrow2::array::{MutableArray, MutableUtf8Array, Utf8Array}; +use arrow2::array::{MutableArray, MutableUtf8Array, TryExtendFromSelf, Utf8Array}; use arrow2::bitmap::Bitmap; use arrow2::datatypes::DataType; @@ -198,3 +198,15 @@ fn as_box_twice() { let _ = a.as_arc(); let _ = a.as_arc(); } + +#[test] +fn extend_from_self() { + let mut a = MutableUtf8Array::::from(&[Some("aa"), None]); + + a.try_extend_from_self(&a.clone()).unwrap(); + + assert_eq!( + a, + MutableUtf8Array::::from([Some("aa"), None, Some("aa"), None]) + ); +}