From aff06983763c7df04f420f811483c946afbc8ee7 Mon Sep 17 00:00:00 2001 From: "Jorge C. Leitao" Date: Fri, 30 Jul 2021 05:29:45 +0000 Subject: [PATCH 1/2] Systematized `new` signature of growable. Also removes an extra allocation. --- src/array/growable/binary.rs | 19 ++--- src/array/growable/boolean.rs | 13 +-- src/array/growable/fixed_binary.rs | 28 +++---- src/array/growable/list.rs | 21 ++--- src/array/growable/mod.rs | 129 +++++++++++++++++++---------- src/array/growable/primitive.rs | 18 ++-- src/array/growable/structure.rs | 10 +-- src/array/growable/utf8.rs | 12 +-- src/compute/filter.rs | 10 +-- src/compute/take/list.rs | 6 +- 10 files changed, 138 insertions(+), 128 deletions(-) diff --git a/src/array/growable/binary.rs b/src/array/growable/binary.rs index e012b22bc1d..5ae12c257f2 100644 --- a/src/array/growable/binary.rs +++ b/src/array/growable/binary.rs @@ -26,10 +26,10 @@ pub struct GrowableBinary<'a, O: Offset> { impl<'a, O: Offset> GrowableBinary<'a, O> { /// # Panics /// This function panics if any of the `arrays` is not downcastable to `PrimitiveArray`. - pub fn new(arrays: &[&'a dyn Array], mut use_validity: bool, capacity: usize) -> Self { + pub fn new(arrays: Vec<&'a BinaryArray>, mut use_validity: bool, capacity: usize) -> Self { // if any of the arrays has nulls, insertions from any array requires setting bits // as there is at least one array with nulls. - if arrays.iter().any(|array| array.null_count() > 0) { + if !use_validity & arrays.iter().any(|array| array.null_count() > 0) { use_validity = true; }; @@ -38,11 +38,6 @@ impl<'a, O: Offset> GrowableBinary<'a, O> { .map(|array| build_extend_null_bits(*array, use_validity)) .collect(); - let arrays = arrays - .iter() - .map(|array| array.as_any().downcast_ref::>().unwrap()) - .collect::>(); - let mut offsets = MutableBuffer::with_capacity(capacity + 1); let length = O::default(); unsafe { offsets.push_unchecked(length) }; @@ -114,7 +109,7 @@ mod tests { fn test_variable_sized_validity() { let array = BinaryArray::::from_iter(vec![Some("a"), Some("bc"), None, Some("defh")]); - let mut a = GrowableBinary::new(&[&array], false, 0); + let mut a = GrowableBinary::new(vec![&array], false, 0); a.extend(0, 1, 2); @@ -131,7 +126,7 @@ mod tests { let array = BinaryArray::::from_iter(vec![Some("a"), Some("bc"), None, Some("defh")]); let array = array.slice(1, 3); - let mut a = GrowableBinary::new(&[&array], false, 0); + let mut a = GrowableBinary::new(vec![&array], false, 0); a.extend(0, 0, 3); @@ -146,7 +141,7 @@ mod tests { let array = BinaryArray::::from_iter(vec![Some("a"), Some("bc"), None, Some("defh")]); let array = array.slice(1, 3); - let mut a = GrowableBinary::new(&[&array], false, 0); + let mut a = GrowableBinary::new(vec![&array], false, 0); a.extend(0, 0, 3); @@ -161,7 +156,7 @@ mod tests { let array1 = BinaryArray::::from_slice(vec![b"hello", b"world"]); let array2 = BinaryArray::::from_iter(vec![Some("1"), None]); - let mut a = GrowableBinary::new(&[&array1, &array2], false, 5); + let mut a = GrowableBinary::new(vec![&array1, &array2], false, 5); a.extend(0, 0, 2); a.extend(1, 0, 2); @@ -178,7 +173,7 @@ mod tests { let array = BinaryArray::::from_iter(vec![Some("a"), Some("bc"), None, Some("defh")]); let array = array.slice(1, 3); - let mut a = GrowableBinary::new(&[&array], true, 0); + let mut a = GrowableBinary::new(vec![&array], true, 0); a.extend(0, 1, 2); a.extend_validity(1); diff --git a/src/array/growable/boolean.rs b/src/array/growable/boolean.rs index 666e145f467..24bbf8267af 100644 --- a/src/array/growable/boolean.rs +++ b/src/array/growable/boolean.rs @@ -21,12 +21,10 @@ pub struct GrowableBoolean<'a> { } impl<'a> GrowableBoolean<'a> { - /// # Panics - /// This function panics if any of the `arrays` is not downcastable to `PrimitiveArray`. - pub fn new(arrays: &[&'a dyn Array], mut use_validity: bool, capacity: usize) -> Self { + pub fn new(arrays: Vec<&'a BooleanArray>, mut use_validity: bool, capacity: usize) -> Self { // if any of the arrays has nulls, insertions from any array requires setting bits // as there is at least one array with nulls. - if arrays.iter().any(|array| array.null_count() > 0) { + if !use_validity & arrays.iter().any(|array| array.null_count() > 0) { use_validity = true; }; @@ -35,11 +33,6 @@ impl<'a> GrowableBoolean<'a> { .map(|array| build_extend_null_bits(*array, use_validity)) .collect(); - let arrays = arrays - .iter() - .map(|array| array.as_any().downcast_ref::().unwrap()) - .collect::>(); - Self { arrays, values: MutableBitmap::with_capacity(capacity), @@ -94,7 +87,7 @@ mod tests { fn test_bool() { let array = BooleanArray::from(vec![Some(false), Some(true), None, Some(false)]); - let mut a = GrowableBoolean::new(&[&array], false, 0); + let mut a = GrowableBoolean::new(vec![&array], false, 0); a.extend(0, 1, 2); diff --git a/src/array/growable/fixed_binary.rs b/src/array/growable/fixed_binary.rs index cc428f20e1b..1c15d37ddcd 100644 --- a/src/array/growable/fixed_binary.rs +++ b/src/array/growable/fixed_binary.rs @@ -23,9 +23,11 @@ pub struct GrowableFixedSizeBinary<'a> { } impl<'a> GrowableFixedSizeBinary<'a> { - /// # Panics - /// This function panics if any of the `arrays` is not downcastable to `FixedSizeBinaryArray`. - pub fn new(arrays: &[&'a dyn Array], mut use_validity: bool, capacity: usize) -> Self { + pub fn new( + arrays: Vec<&'a FixedSizeBinaryArray>, + mut use_validity: bool, + capacity: usize, + ) -> Self { // if any of the arrays has nulls, insertions from any array requires setting bits // as there is at least one array with nulls. if arrays.iter().any(|array| array.null_count() > 0) { @@ -37,16 +39,6 @@ impl<'a> GrowableFixedSizeBinary<'a> { .map(|array| build_extend_null_bits(*array, use_validity)) .collect(); - let arrays = arrays - .iter() - .map(|array| { - array - .as_any() - .downcast_ref::() - .unwrap() - }) - .collect::>(); - let size = *FixedSizeBinaryArray::get_size(arrays[0].data_type()) as usize; Self { arrays, @@ -115,7 +107,7 @@ mod tests { let array = FixedSizeBinaryArray::from_iter(vec![Some(b"ab"), Some(b"bc"), None, Some(b"de")], 2); - let mut a = GrowableFixedSizeBinary::new(&[&array], false, 0); + let mut a = GrowableFixedSizeBinary::new(vec![&array], false, 0); a.extend(0, 1, 2); @@ -133,7 +125,7 @@ mod tests { FixedSizeBinaryArray::from_iter(vec![Some(b"ab"), Some(b"bc"), None, Some(b"fh")], 2); let array = array.slice(1, 3); - let mut a = GrowableFixedSizeBinary::new(&[&array], false, 0); + let mut a = GrowableFixedSizeBinary::new(vec![&array], false, 0); a.extend(0, 0, 3); @@ -148,7 +140,7 @@ mod tests { let array1 = FixedSizeBinaryArray::from_iter(vec![Some("hello"), Some("world")], 5); let array2 = FixedSizeBinaryArray::from_iter(vec![Some("12345"), None], 5); - let mut a = GrowableFixedSizeBinary::new(&[&array1, &array2], false, 5); + let mut a = GrowableFixedSizeBinary::new(vec![&array1, &array2], false, 5); a.extend(0, 0, 2); a.extend(1, 0, 2); @@ -168,7 +160,7 @@ mod tests { FixedSizeBinaryArray::from_iter(vec![Some("aa"), Some("bc"), None, Some("fh")], 2); let array = array.slice(1, 3); - let mut a = GrowableFixedSizeBinary::new(&[&array], true, 0); + let mut a = GrowableFixedSizeBinary::new(vec![&array], true, 0); a.extend(0, 1, 2); a.extend_validity(1); @@ -186,7 +178,7 @@ mod tests { let array = array.slice(1, 2); // = [[0, 1], [0, 2]] due to the offset = 1 - let mut a = GrowableFixedSizeBinary::new(&[&array], false, 0); + let mut a = GrowableFixedSizeBinary::new(vec![&array], false, 0); a.extend(0, 1, 1); a.extend(0, 0, 1); diff --git a/src/array/growable/list.rs b/src/array/growable/list.rs index 311b79aacdc..3c360c5cc9b 100644 --- a/src/array/growable/list.rs +++ b/src/array/growable/list.rs @@ -67,12 +67,10 @@ pub struct GrowableList<'a, O: Offset> { } impl<'a, O: Offset> GrowableList<'a, O> { - /// # Panics - /// This function panics if any of the `arrays` is not downcastable to `PrimitiveArray`. - pub fn new(arrays: &[&'a dyn Array], mut use_validity: bool, capacity: usize) -> Self { + pub fn new(arrays: Vec<&'a ListArray>, mut use_validity: bool, capacity: usize) -> Self { // if any of the arrays has nulls, insertions from any array requires setting bits // as there is at least one array with nulls. - if use_validity || arrays.iter().any(|array| array.null_count() > 0) { + if !use_validity & arrays.iter().any(|array| array.null_count() > 0) { use_validity = true; }; @@ -81,11 +79,6 @@ impl<'a, O: Offset> GrowableList<'a, O> { .map(|array| build_extend_null_bits(*array, use_validity)) .collect(); - let arrays = arrays - .iter() - .map(|array| array.as_any().downcast_ref::>().unwrap()) - .collect::>(); - let inner = arrays .iter() .map(|array| array.values().as_ref()) @@ -176,7 +169,7 @@ mod tests { let array = create_list_array(data); - let mut a = GrowableList::new(&[&array], false, 0); + let mut a = GrowableList::new(vec![&array], false, 0); a.extend(0, 0, 1); let result: ListArray = a.into(); @@ -197,7 +190,7 @@ mod tests { let array = create_list_array(data); let array = array.slice(1, 2); - let mut a = GrowableList::new(&[&array], false, 0); + let mut a = GrowableList::new(vec![&array], false, 0); a.extend(0, 1, 1); let result: ListArray = a.into(); @@ -218,7 +211,7 @@ mod tests { let array = create_list_array(data); let array = array.slice(1, 2); - let mut a = GrowableList::new(&[&array], false, 0); + let mut a = GrowableList::new(vec![&array], false, 0); a.extend(0, 1, 1); let result: ListArray = a.into(); @@ -245,9 +238,7 @@ mod tests { ]; let array_2 = create_list_array(data_2); - let arrays: Vec<&dyn Array> = vec![&array_1, &array_2]; - - let mut a = GrowableList::new(&arrays, false, 6); + let mut a = GrowableList::new(vec![&array_1, &array_2], false, 6); a.extend(0, 0, 2); a.extend(1, 1, 1); diff --git a/src/array/growable/mod.rs b/src/array/growable/mod.rs index 4903b3aee54..db54a6e12e4 100644 --- a/src/array/growable/mod.rs +++ b/src/array/growable/mod.rs @@ -53,15 +53,10 @@ macro_rules! dyn_growable { ($ty:ty, $arrays:expr, $use_validity:expr, $capacity:expr) => {{ let arrays = $arrays .iter() - .map(|array| { - array - .as_any() - .downcast_ref::>() - .unwrap() - }) + .map(|array| array.as_any().downcast_ref().unwrap()) .collect::>(); Box::new(primitive::GrowablePrimitive::<$ty>::new( - &arrays, + arrays, $use_validity, $capacity, )) @@ -103,11 +98,17 @@ pub fn make_growable<'a>( match data_type { DataType::Null => Box::new(null::GrowableNull::new()), - DataType::Boolean => Box::new(boolean::GrowableBoolean::new( - arrays, - use_validity, - capacity, - )), + DataType::Boolean => { + let arrays = arrays + .iter() + .map(|array| array.as_any().downcast_ref().unwrap()) + .collect::>(); + Box::new(boolean::GrowableBoolean::new( + arrays, + use_validity, + capacity, + )) + } DataType::Int8 => dyn_growable!(i8, arrays, use_validity, capacity), DataType::Int16 => dyn_growable!(i16, arrays, use_validity, capacity), DataType::Int32 @@ -137,10 +138,10 @@ pub fn make_growable<'a>( DataType::Utf8 => { let arrays = arrays .iter() - .map(|array| array.as_any().downcast_ref::>().unwrap()) + .map(|array| array.as_any().downcast_ref().unwrap()) .collect::>(); Box::new(utf8::GrowableUtf8::::new( - &arrays, + arrays, use_validity, capacity, )) @@ -148,45 +149,81 @@ pub fn make_growable<'a>( DataType::LargeUtf8 => { let arrays = arrays .iter() - .map(|array| array.as_any().downcast_ref::>().unwrap()) + .map(|array| array.as_any().downcast_ref().unwrap()) .collect::>(); Box::new(utf8::GrowableUtf8::::new( - &arrays, + arrays, + use_validity, + capacity, + )) + } + DataType::Binary => { + let arrays = arrays + .iter() + .map(|array| array.as_any().downcast_ref().unwrap()) + .collect::>(); + Box::new(binary::GrowableBinary::::new( + arrays, + use_validity, + capacity, + )) + } + DataType::LargeBinary => { + let arrays = arrays + .iter() + .map(|array| array.as_any().downcast_ref().unwrap()) + .collect::>(); + Box::new(binary::GrowableBinary::::new( + arrays, + use_validity, + capacity, + )) + } + DataType::FixedSizeBinary(_) => { + let arrays = arrays + .iter() + .map(|array| array.as_any().downcast_ref().unwrap()) + .collect::>(); + Box::new(fixed_binary::GrowableFixedSizeBinary::new( + arrays, use_validity, capacity, )) } - DataType::Binary => Box::new(binary::GrowableBinary::::new( - arrays, - use_validity, - capacity, - )), - DataType::LargeBinary => Box::new(binary::GrowableBinary::::new( - arrays, - use_validity, - capacity, - )), - DataType::FixedSizeBinary(_) => Box::new(fixed_binary::GrowableFixedSizeBinary::new( - arrays, - use_validity, - capacity, - )), - DataType::List(_) => Box::new(list::GrowableList::::new( - arrays, - use_validity, - capacity, - )), - DataType::LargeList(_) => Box::new(list::GrowableList::::new( - arrays, - use_validity, - capacity, - )), - DataType::Struct(_) => Box::new(structure::GrowableStruct::new( - arrays, - use_validity, - capacity, - )), + DataType::List(_) => { + let arrays = arrays + .iter() + .map(|array| array.as_any().downcast_ref().unwrap()) + .collect::>(); + Box::new(list::GrowableList::::new( + arrays, + use_validity, + capacity, + )) + } + DataType::LargeList(_) => { + let arrays = arrays + .iter() + .map(|array| array.as_any().downcast_ref().unwrap()) + .collect::>(); + Box::new(list::GrowableList::::new( + arrays, + use_validity, + capacity, + )) + } + DataType::Struct(_) => { + let arrays = arrays + .iter() + .map(|array| array.as_any().downcast_ref().unwrap()) + .collect::>(); + Box::new(structure::GrowableStruct::new( + arrays, + use_validity, + capacity, + )) + } DataType::FixedSizeList(_, _) => todo!(), DataType::Union(_) => todo!(), DataType::Dictionary(key, _) => match key.as_ref() { diff --git a/src/array/growable/primitive.rs b/src/array/growable/primitive.rs index 40f5931a791..6e4c8a26ae6 100644 --- a/src/array/growable/primitive.rs +++ b/src/array/growable/primitive.rs @@ -21,10 +21,14 @@ pub struct GrowablePrimitive<'a, T: NativeType> { } impl<'a, T: NativeType> GrowablePrimitive<'a, T> { - pub fn new(arrays: &[&'a PrimitiveArray], mut use_validity: bool, capacity: usize) -> Self { + pub fn new( + arrays: Vec<&'a PrimitiveArray>, + mut use_validity: bool, + capacity: usize, + ) -> Self { // if any of the arrays has nulls, insertions from any array requires setting bits // as there is at least one array with nulls. - if arrays.iter().any(|array| array.null_count() > 0) { + if !use_validity & arrays.iter().any(|array| array.null_count() > 0) { use_validity = true; }; @@ -103,7 +107,7 @@ mod tests { #[test] fn test_primitive() { let b = PrimitiveArray::::from(vec![Some(1), Some(2), Some(3)]).to(DataType::UInt8); - let mut a = GrowablePrimitive::new(&[&b], false, 3); + let mut a = GrowablePrimitive::new(vec![&b], false, 3); a.extend(0, 0, 2); let result: PrimitiveArray = a.into(); let expected = PrimitiveArray::::from(vec![Some(1), Some(2)]).to(DataType::UInt8); @@ -115,7 +119,7 @@ mod tests { fn test_primitive_offset() { let b = PrimitiveArray::::from(vec![Some(1), Some(2), Some(3)]).to(DataType::UInt8); let b = b.slice(1, 2); - let mut a = GrowablePrimitive::new(&[&b], false, 2); + let mut a = GrowablePrimitive::new(vec![&b], false, 2); a.extend(0, 0, 2); let result: PrimitiveArray = a.into(); let expected = PrimitiveArray::::from(vec![Some(2), Some(3)]).to(DataType::UInt8); @@ -127,7 +131,7 @@ mod tests { fn test_primitive_null_offset() { let b = PrimitiveArray::::from(vec![Some(1), None, Some(3)]).to(DataType::UInt8); let b = b.slice(1, 2); - let mut a = GrowablePrimitive::new(&[&b], false, 2); + let mut a = GrowablePrimitive::new(vec![&b], false, 2); a.extend(0, 0, 2); let result: PrimitiveArray = a.into(); let expected = PrimitiveArray::::from(vec![None, Some(3)]).to(DataType::UInt8); @@ -138,7 +142,7 @@ mod tests { fn test_primitive_null_offset_validity() { let b = PrimitiveArray::::from(vec![Some(1), Some(2), Some(3)]).to(DataType::UInt8); let b = b.slice(1, 2); - let mut a = GrowablePrimitive::new(&[&b], true, 2); + let mut a = GrowablePrimitive::new(vec![&b], true, 2); a.extend(0, 0, 2); a.extend_validity(3); a.extend(0, 1, 1); @@ -153,7 +157,7 @@ mod tests { fn test_primitive_joining_arrays() { let b = PrimitiveArray::::from(vec![Some(1), Some(2), Some(3)]).to(DataType::UInt8); let c = PrimitiveArray::::from(vec![Some(4), Some(5), Some(6)]).to(DataType::UInt8); - let mut a = GrowablePrimitive::new(&[&b, &c], false, 4); + let mut a = GrowablePrimitive::new(vec![&b, &c], false, 4); a.extend(0, 0, 2); a.extend(1, 1, 2); let result: PrimitiveArray = a.into(); diff --git a/src/array/growable/structure.rs b/src/array/growable/structure.rs index 7c4bb1a6081..0248fe8f328 100644 --- a/src/array/growable/structure.rs +++ b/src/array/growable/structure.rs @@ -24,7 +24,7 @@ pub struct GrowableStruct<'a> { impl<'a> GrowableStruct<'a> { /// # Panics /// This function panics if any of the `arrays` is not downcastable to `PrimitiveArray`. - pub fn new(arrays: &[&'a dyn Array], mut use_validity: bool, capacity: usize) -> Self { + pub fn new(arrays: Vec<&'a StructArray>, mut use_validity: bool, capacity: usize) -> Self { // if any of the arrays has nulls, insertions from any array requires setting bits // as there is at least one array with nulls. if arrays.iter().any(|array| array.null_count() > 0) { @@ -155,7 +155,7 @@ mod tests { let array = StructArray::from_data(fields.clone(), values.clone(), None); - let mut a = GrowableStruct::new(&[&array], false, 0); + let mut a = GrowableStruct::new(vec![&array], false, 0); a.extend(0, 1, 2); let result: StructArray = a.into(); @@ -174,7 +174,7 @@ mod tests { let array = StructArray::from_data(fields.clone(), values.clone(), None).slice(1, 3); - let mut a = GrowableStruct::new(&[&array], false, 0); + let mut a = GrowableStruct::new(vec![&array], false, 0); a.extend(0, 1, 2); let result: StructArray = a.into(); @@ -198,7 +198,7 @@ mod tests { Some(Bitmap::from_u8_slice(&[0b00000010], 5)), ); - let mut a = GrowableStruct::new(&[&array], false, 0); + let mut a = GrowableStruct::new(vec![&array], false, 0); a.extend(0, 1, 2); let result: StructArray = a.into(); @@ -218,7 +218,7 @@ mod tests { let array = StructArray::from_data(fields.clone(), values.clone(), None); - let mut mutable = GrowableStruct::new(&[&array, &array], false, 0); + let mut mutable = GrowableStruct::new(vec![&array, &array], false, 0); mutable.extend(0, 1, 2); mutable.extend(1, 0, 2); diff --git a/src/array/growable/utf8.rs b/src/array/growable/utf8.rs index 069e067504f..26005966519 100644 --- a/src/array/growable/utf8.rs +++ b/src/array/growable/utf8.rs @@ -24,7 +24,7 @@ pub struct GrowableUtf8<'a, O: Offset> { } impl<'a, O: Offset> GrowableUtf8<'a, O> { - pub fn new(arrays: &[&'a Utf8Array], mut use_validity: bool, capacity: usize) -> Self { + pub fn new(arrays: Vec<&'a Utf8Array>, mut use_validity: bool, capacity: usize) -> Self { // if any of the arrays has nulls, insertions from any array requires setting bits // as there is at least one array with nulls. if arrays.iter().any(|array| array.null_count() > 0) { @@ -117,7 +117,7 @@ mod tests { fn test_variable_sized_validity() { let array = Utf8Array::::from_iter(vec![Some("a"), Some("bc"), None, Some("defh")]); - let mut a = GrowableUtf8::new(&[&array], false, 0); + let mut a = GrowableUtf8::new(vec![&array], false, 0); a.extend(0, 1, 2); @@ -134,7 +134,7 @@ mod tests { let array = Utf8Array::::from_iter(vec![Some("a"), Some("bc"), None, Some("defh")]); let array = array.slice(1, 3); - let mut a = GrowableUtf8::new(&[&array], false, 0); + let mut a = GrowableUtf8::new(vec![&array], false, 0); a.extend(0, 0, 3); @@ -149,7 +149,7 @@ mod tests { let array = Utf8Array::::from_iter(vec![Some("a"), Some("bc"), None, Some("defh")]); let array = array.slice(1, 3); - let mut a = GrowableUtf8::new(&[&array], false, 0); + let mut a = GrowableUtf8::new(vec![&array], false, 0); a.extend(0, 0, 3); @@ -164,7 +164,7 @@ mod tests { let array1 = Utf8Array::::from_slice(&["hello", "world"]); let array2 = Utf8Array::::from(&[Some("1"), None]); - let mut a = GrowableUtf8::new(&[&array1, &array2], false, 5); + let mut a = GrowableUtf8::new(vec![&array1, &array2], false, 5); a.extend(0, 0, 2); a.extend(1, 0, 2); @@ -181,7 +181,7 @@ mod tests { let array = Utf8Array::::from_iter(vec![Some("a"), Some("bc"), None, Some("defh")]); let array = array.slice(1, 3); - let mut a = GrowableUtf8::new(&[&array], true, 0); + let mut a = GrowableUtf8::new(vec![&array], true, 0); a.extend(0, 1, 2); a.extend_validity(1); diff --git a/src/compute/filter.rs b/src/compute/filter.rs index 9b975d8529c..62421afb6a2 100644 --- a/src/compute/filter.rs +++ b/src/compute/filter.rs @@ -90,11 +90,9 @@ fn filter_growable<'a>(growable: &mut impl Growable<'a>, chunks: &[(usize, usize macro_rules! dyn_build_filter { ($ty:ty, $array:expr, $filter_count:expr, $chunks:expr) => {{ - let array = $array - .as_any() - .downcast_ref::>() - .unwrap(); - let mut growable = growable::GrowablePrimitive::<$ty>::new(&[array], false, $filter_count); + let array = $array.as_any().downcast_ref().unwrap(); + let mut growable = + growable::GrowablePrimitive::<$ty>::new(vec![array], false, $filter_count); filter_growable(&mut growable, &$chunks); let array: PrimitiveArray<$ty> = growable.into(); Box::new(array) @@ -154,7 +152,7 @@ pub fn build_filter(filter: &BooleanArray) -> Result { } DataType::Utf8 => { let array = array.as_any().downcast_ref::>().unwrap(); - let mut growable = growable::GrowableUtf8::::new(&[array], false, filter_count); + let mut growable = growable::GrowableUtf8::::new(vec![array], false, filter_count); filter_growable(&mut growable, &chunks); let array: Utf8Array = growable.into(); Box::new(array) diff --git a/src/compute/take/list.rs b/src/compute/take/list.rs index 0d5251d3d5f..5f2fbf4c283 100644 --- a/src/compute/take/list.rs +++ b/src/compute/take/list.rs @@ -39,10 +39,10 @@ pub fn take( }) .collect::>>(); - let array_ref: Vec<&dyn Array> = arrays.iter().map(|v| v as &dyn Array).collect(); + let arrays = arrays.iter().collect(); if let Some(validity) = indices.validity() { - let mut growable: GrowableList = GrowableList::new(&array_ref, true, capacity); + let mut growable: GrowableList = GrowableList::new(arrays, true, capacity); for index in 0..indices.len() { if validity.get_bit(index) { @@ -54,7 +54,7 @@ pub fn take( growable.into() } else { - let mut growable: GrowableList = GrowableList::new(&array_ref, false, capacity); + let mut growable: GrowableList = GrowableList::new(arrays, false, capacity); for index in 0..indices.len() { growable.extend(index, 0, 1); } From f2712586b3ac3574d2d9f6d78107e38a7db1961c Mon Sep 17 00:00:00 2001 From: "Jorge C. Leitao" Date: Fri, 30 Jul 2021 07:11:51 +0000 Subject: [PATCH 2/2] Added example. --- examples/growable.rs | 38 +++++++++++++++++++++++++++++++++ src/array/growable/primitive.rs | 7 +++--- 2 files changed, 41 insertions(+), 4 deletions(-) create mode 100644 examples/growable.rs diff --git a/examples/growable.rs b/examples/growable.rs new file mode 100644 index 00000000000..cc8b7573e49 --- /dev/null +++ b/examples/growable.rs @@ -0,0 +1,38 @@ +use arrow2::array::growable::{Growable, GrowablePrimitive}; +use arrow2::array::{Array, PrimitiveArray}; + +fn main() { + // say we have two sorted arrays + let array0 = PrimitiveArray::::from_slice(&[1, 2, 5]); + let array1 = PrimitiveArray::::from_slice(&[3, 4, 6]); + + // and we found a way to compute the slices that sort them: + // (array_index, start of the slice, length of the slice) + let slices = &[ + // [1, 2] from array0 + (0, 0, 2), + // [3, 4] from array1 + (1, 0, 2), + // [5] from array0 + (0, 2, 1), + // [6] from array1 + (1, 2, 1), + ]; + + // we can build a new array out of these slices as follows: + // first, declare the growable out of the arrays. Since we are not extending with nulls, + // we use `false`. We also pre-allocate, as we know that we will need all entries. + let capacity = array0.len() + array1.len(); + let use_validity = false; + let mut growable = GrowablePrimitive::new(vec![&array0, &array1], use_validity, capacity); + + // next, extend the growable: + for slice in slices { + growable.extend(slice.0, slice.1, slice.2); + } + // finally, convert it to the array (this is `O(1)`) + let result: PrimitiveArray = growable.into(); + + let expected = PrimitiveArray::::from_slice(&[1, 2, 3, 4, 5, 6]); + assert_eq!(result, expected); +} diff --git a/src/array/growable/primitive.rs b/src/array/growable/primitive.rs index 6e4c8a26ae6..f1e11118208 100644 --- a/src/array/growable/primitive.rs +++ b/src/array/growable/primitive.rs @@ -155,15 +155,14 @@ mod tests { #[test] fn test_primitive_joining_arrays() { - let b = PrimitiveArray::::from(vec![Some(1), Some(2), Some(3)]).to(DataType::UInt8); - let c = PrimitiveArray::::from(vec![Some(4), Some(5), Some(6)]).to(DataType::UInt8); + let b = PrimitiveArray::::from(vec![Some(1), Some(2), Some(3)]); + let c = PrimitiveArray::::from(vec![Some(4), Some(5), Some(6)]); let mut a = GrowablePrimitive::new(vec![&b, &c], false, 4); a.extend(0, 0, 2); a.extend(1, 1, 2); let result: PrimitiveArray = a.into(); - let expected = PrimitiveArray::::from(vec![Some(1), Some(2), Some(5), Some(6)]) - .to(DataType::UInt8); + let expected = PrimitiveArray::::from(vec![Some(1), Some(2), Some(5), Some(6)]); assert_eq!(result, expected); } }