Simplified new signature of growable API (#238)

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::<i64>::from_slice(&[1, 2, 5]);
+    let array1 = PrimitiveArray::<i64>::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<i64> = growable.into();
+
+    let expected = PrimitiveArray::<i64>::from_slice(&[1, 2, 3, 4, 5, 6]);
+    assert_eq!(result, expected);
+}

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<T>`.
-    pub fn new(arrays: &[&'a dyn Array], mut use_validity: bool, capacity: usize) -> Self {
+    pub fn new(arrays: Vec<&'a BinaryArray<O>>, 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::<BinaryArray<O>>().unwrap())
-            .collect::<Vec<_>>();
-
         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::<i32>::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::<i32>::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::<i32>::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::<i32>::from_slice(vec![b"hello", b"world"]);
         let array2 = BinaryArray::<i32>::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::<i32>::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);

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<T>`.
-    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::<BooleanArray>().unwrap())
-            .collect::<Vec<_>>();
-
         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);
 

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::<FixedSizeBinaryArray>()
-                    .unwrap()
-            })
-            .collect::<Vec<_>>();
-
         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);

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<T>`.
-    pub fn new(arrays: &[&'a dyn Array], mut use_validity: bool, capacity: usize) -> Self {
+    pub fn new(arrays: Vec<&'a ListArray<O>>, 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::<ListArray<O>>().unwrap())
-            .collect::<Vec<_>>();
-
         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<i32> = 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<i32> = 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<i32> = 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);