Repository: jorgecarleitao/arrow2
Pull request #314 (2021-08-22 18:37:16)

Added trait `TryPush`
And implemented it for most of the `Mutable` API.

codecov[bot]: # [Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/314?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) Report
> Merging [#314](https://codecov.io/gh/jorgecarleitao/arrow2/pull/314?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (a7f23ca) into [main](https://codecov.io/gh/jorgecarleitao/arrow2/commit/83256a8040f1524f4f85d3e326fbcb3fe7262fb5?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (83256a8) will **decrease** coverage by `0.01%`.
> The diff coverage is `76.47%`.

[![Impacted file tree graph](https://codecov.io/gh/jorgecarleitao/arrow2/pull/314/graphs/tree.svg?width=650&height=150&src=pr&token=AgyTF60R3D&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao)](https://codecov.io/gh/jorgecarleitao/arrow2/pull/314?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao)

```diff
@@            Coverage Diff             @@
##             main     #314      +/-   ##
==========================================
- Coverage   80.60%   80.59%   -0.02%     
==========================================
  Files         324      324              
  Lines       21255    21268      +13     
==========================================
+ Hits        17132    17140       +8     
- Misses       4123     4128       +5     
```


| [Impacted Files](https://codecov.io/gh/jorgecarleitao/arrow2/pull/314?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) | Coverage Δ | |
|---|---|---|
| [src/array/boolean/mutable.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/314/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L2Jvb2xlYW4vbXV0YWJsZS5ycw==) | `86.86% <0.00%> (-1.95%)` | :arrow_down: |
| [src/array/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/314/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L21vZC5ycw==) | `48.24% <ø> (ø)` | |
| [src/array/primitive/mutable.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/314/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L3ByaW1pdGl2ZS9tdXRhYmxlLnJz) | `85.71% <0.00%> (-1.39%)` | :arrow_down: |
| [src/array/utf8/mutable.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/314/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L3V0ZjgvbXV0YWJsZS5ycw==) | `86.29% <81.25%> (ø)` | |
| [src/array/binary/mutable.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/314/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L2JpbmFyeS9tdXRhYmxlLnJz) | `67.94% <82.35%> (ø)` | |
| [src/array/list/mutable.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/314/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L2xpc3QvbXV0YWJsZS5ycw==) | `75.86% <100.00%> (+0.86%)` | :arrow_up: |
| [tests/it/array/list/mutable.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/314/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-dGVzdHMvaXQvYXJyYXkvbGlzdC9tdXRhYmxlLnJz) | `100.00% <100.00%> (ø)` | |
| [src/io/json\_integration/write.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/314/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL2pzb25faW50ZWdyYXRpb24vd3JpdGUucnM=) | `6.25% <0.00%> (+6.25%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/314?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/314?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao). Last update [83256a8...a7f23ca](https://codecov.io/gh/jorgecarleitao/arrow2/pull/314?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao).


Patch:
diff --git a/src/array/binary/mutable.rs b/src/array/binary/mutable.rs
index 0326497278..fe4b3e050f 100644
--- a/src/array/binary/mutable.rs
+++ b/src/array/binary/mutable.rs
@@ -1,7 +1,7 @@
 use std::{iter::FromIterator, sync::Arc};
 
 use crate::{
-    array::{Array, MutableArray, Offset, TryExtend},
+    array::{Array, MutableArray, Offset, TryExtend, TryPush},
     bitmap::MutableBitmap,
     buffer::MutableBuffer,
     datatypes::DataType,
@@ -61,34 +61,9 @@ impl<O: Offset> MutableBinaryArray<O> {
         *self.offsets.last().unwrap()
     }
 
-    pub fn try_push<T: AsRef<[u8]>>(&mut self, value: Option<T>) -> Result<()> {
-        match value {
-            Some(value) => {
-                let bytes = value.as_ref();
-
-                let size = O::from_usize(self.values.len() + bytes.len())
-                    .ok_or(ArrowError::KeyOverflowError)?;
-
-                self.values.extend_from_slice(bytes);
-
-                self.offsets.push(size);
-
-                match &mut self.validity {
-                    Some(validity) => validity.push(true),
-                    None => {}
-                }
-            }
-            None => {
-                self.offsets.push(self.last_offset());
-                match &mut self.validity {
-                    Some(validity) => validity.push(false),
-                    None => self.init_validity(),
-                }
-            }
-        }
-        Ok(())
-    }
-
+    /// Pushes a new element to the array.
+    /// # Panic
+    /// This operation panics iff the length of all values (in bytes) exceeds `O` maximum value.
     pub fn push<T: AsRef<[u8]>>(&mut self, value: Option<T>) {
         self.try_push(value).unwrap()
     }
@@ -176,3 +151,33 @@ impl<O: Offset, T: AsRef<[u8]>> TryExtend<Option<T>> for MutableBinaryArray<O> {
         iter.try_for_each(|x| self.try_push(x))
     }
 }
+
+impl<O: Offset, T: AsRef<[u8]>> TryPush<Option<T>> for MutableBinaryArray<O> {
+    fn try_push(&mut self, value: Option<T>) -> Result<()> {
+        match value {
+            Some(value) => {
+                let bytes = value.as_ref();
+
+                let size = O::from_usize(self.values.len() + bytes.len())
+                    .ok_or(ArrowError::KeyOverflowError)?;
+
+                self.values.extend_from_slice(bytes);
+
+                self.offsets.push(size);
+
+                match &mut self.validity {
+                    Some(validity) => validity.push(true),
+                    None => {}
+                }
+            }
+            None => {
+                self.offsets.push(self.last_offset());
+                match &mut self.validity {
+                    Some(validity) => validity.push(false),
+                    None => self.init_validity(),
+                }
+            }
+        }
+        Ok(())
+    }
+}
diff --git a/src/array/boolean/mutable.rs b/src/array/boolean/mutable.rs
index 603a30b83b..32f3ba6d89 100644
--- a/src/array/boolean/mutable.rs
+++ b/src/array/boolean/mutable.rs
@@ -1,9 +1,8 @@
 use std::iter::FromIterator;
 use std::sync::Arc;
 
-use crate::array::TryExtend;
 use crate::{
-    array::{Array, MutableArray},
+    array::{Array, MutableArray, TryExtend, TryPush},
     bitmap::MutableBitmap,
     datatypes::DataType,
     error::Result,
@@ -352,6 +351,14 @@ impl TryExtend<Option<bool>> for MutableBooleanArray {
     }
 }
 
+impl TryPush<Option<bool>> for MutableBooleanArray {
+    /// This is infalible and is implemented for consistency with all other types
+    fn try_push(&mut self, item: Option<bool>) -> Result<()> {
+        self.push(item);
+        Ok(())
+    }
+}
+
 impl PartialEq for MutableBooleanArray {
     fn eq(&self, other: &Self) -> bool {
         self.iter().eq(other.iter())
diff --git a/src/array/list/mutable.rs b/src/array/list/mutable.rs
index 3dde5bede0..80a846b8e8 100644
--- a/src/array/list/mutable.rs
+++ b/src/array/list/mutable.rs
@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use crate::{
-    array::{Array, MutableArray, Offset, TryExtend},
+    array::{Array, MutableArray, Offset, TryExtend, TryPush},
     bitmap::MutableBitmap,
     buffer::MutableBuffer,
     datatypes::{DataType, Field},
@@ -67,13 +67,25 @@ where
 {
     fn try_extend<II: IntoIterator<Item = Option<I>>>(&mut self, iter: II) -> Result<()> {
         for items in iter {
-            if let Some(items) = items {
-                let values = self.mut_values();
-                values.try_extend(items)?;
-                self.try_push_valid()?;
-            } else {
-                self.push_null();
-            }
+            self.try_push(items)?;
+        }
+        Ok(())
+    }
+}
+
+impl<O, M, I, T> TryPush<Option<I>> for MutableListArray<O, M>
+where
+    O: Offset,
+    M: MutableArray + TryExtend<Option<T>>,
+    I: IntoIterator<Item = Option<T>>,
+{
+    fn try_push(&mut self, item: Option<I>) -> Result<()> {
+        if let Some(items) = item {
+            let values = self.mut_values();
+            values.try_extend(items)?;
+            self.try_push_valid()?;
+        } else {
+            self.push_null();
         }
         Ok(())
     }
diff --git a/src/array/mod.rs b/src/array/mod.rs
index 2a9fcb3993..09b1d3ac0d 100644
--- a/src/array/mod.rs
+++ b/src/array/mod.rs
@@ -414,6 +414,12 @@ pub trait TryExtend<A> {
     fn try_extend<I: IntoIterator<Item = A>>(&mut self, iter: I) -> Result<()>;
 }
 
+/// A trait describing the ability of a struct to receive new items.
+pub trait TryPush<A> {
+    /// Tries to push a new element.
+    fn try_push(&mut self, item: A) -> Result<()>;
+}
+
 fn display_helper<T: std::fmt::Display, I: IntoIterator<Item = Option<T>>>(iter: I) -> Vec<String> {
     let iterator = iter.into_iter();
     let len = iterator.size_hint().0;
diff --git a/src/array/primitive/mutable.rs b/src/array/primitive/mutable.rs
index 19e352a34d..ee33018812 100644
--- a/src/array/primitive/mutable.rs
+++ b/src/array/primitive/mutable.rs
@@ -1,7 +1,7 @@
 use std::{iter::FromIterator, sync::Arc};
 
 use crate::{
-    array::{Array, MutableArray, TryExtend},
+    array::{Array, MutableArray, TryExtend, TryPush},
     bitmap::MutableBitmap,
     buffer::MutableBuffer,
     datatypes::DataType,
@@ -266,6 +266,14 @@ impl<T: NativeType> TryExtend<Option<T>> for MutablePrimitiveArray<T> {
     }
 }
 
+impl<T: NativeType> TryPush<Option<T>> for MutablePrimitiveArray<T> {
+    /// This is infalible and is implemented for consistency with all other types
+    fn try_push(&mut self, item: Option<T>) -> Result<()> {
+        self.push(item);
+        Ok(())
+    }
+}
+
 impl<T: NativeType> MutableArray for MutablePrimitiveArray<T> {
     fn len(&self) -> usize {
         self.values.len()
diff --git a/src/array/struct_.rs b/src/array/struct_.rs
index 86c2b09210..9a6191a5d4 100644
--- a/src/array/struct_.rs
+++ b/src/array/struct_.rs
@@ -52,6 +52,20 @@ impl StructArray {
         }
     }
 
+    pub fn into_data(self) -> (Vec<Field>, Vec<Arc<dyn Array>>, Option<Bitmap>) {
+        let Self {
+            data_type,
+            values,
+            validity,
+        } = self;
+        let fields = if let DataType::Struct(fields) = data_type {
+            fields
+        } else {
+            unreachable!()
+        };
+        (fields, values, validity)
+    }
+
     pub fn slice(&self, offset: usize, length: usize) -> Self {
         let validity = self.validity.clone().map(|x| x.slice(offset, length));
         Self {
diff --git a/src/array/utf8/mutable.rs b/src/array/utf8/mutable.rs
index 1e3b27020c..61aa037aac 100644
--- a/src/array/utf8/mutable.rs
+++ b/src/array/utf8/mutable.rs
@@ -3,7 +3,7 @@ use std::{iter::FromIterator, sync::Arc};
 use crate::{
     array::{
         specification::{check_offsets, check_offsets_and_utf8},
-        Array, MutableArray, Offset, TryExtend,
+        Array, MutableArray, Offset, TryExtend, TryPush,
     },
     bitmap::MutableBitmap,
     buffer::MutableBuffer,
@@ -126,35 +126,6 @@ impl<O: Offset> MutableUtf8Array<O> {
         *self.offsets.last().unwrap()
     }
 
-    /// Tries to push a new element to the array.
-    /// # Error
-    /// This operation errors iff the length of all values (in bytes) exceeds `O` maximum value.
-    pub fn try_push<T: AsRef<str>>(&mut self, value: Option<T>) -> Result<()> {
-        match value {
-            Some(value) => {
-                let bytes = value.as_ref().as_bytes();
-                self.values.extend_from_slice(bytes);
-
-                let size = O::from_usize(self.values.len()).ok_or(ArrowError::KeyOverflowError)?;
-
-                self.offsets.push(size);
-
-                match &mut self.validity {
-                    Some(validity) => validity.push(true),
-                    None => {}
-                }
-            }
-            None => {
-                self.offsets.push(self.last_offset());
-                match &mut self.validity {
-                    Some(validity) => validity.push(false),
-                    None => self.init_validity(),
-                }
-            }
-        }
-        Ok(())
-    }
-
     /// Pushes a new element to the array.
     /// # Panic
     /// This operation panics iff the length of all values (in bytes) exceeds `O` maximum value.
@@ -340,6 +311,34 @@ impl<O: Offset, T: AsRef<str>> TryExtend<Option<T>> for MutableUtf8Array<O> {
     }
 }
 
+impl<O: Offset, T: AsRef<str>> TryPush<Option<T>> for MutableUtf8Array<O> {
+    fn try_push(&mut self, value: Option<T>) -> Result<()> {
+        match value {
+            Some(value) => {
+                let bytes = value.as_ref().as_bytes();
+                self.values.extend_from_slice(bytes);
+
+                let size = O::from_usize(self.values.len()).ok_or(ArrowError::KeyOverflowError)?;
+
+                self.offsets.push(size);
+
+                match &mut self.validity {
+                    Some(validity) => validity.push(true),
+                    None => {}
+                }
+            }
+            None => {
+                self.offsets.push(self.last_offset());
+                match &mut self.validity {
+                    Some(validity) => validity.push(false),
+                    None => self.init_validity(),
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
 /// Creates [`MutableBitmap`] and two [`MutableBuffer`]s from an iterator of `Option`.
 /// The first buffer corresponds to a offset buffer, the second one
 /// corresponds to a values buffer.
diff --git a/src/record_batch.rs b/src/record_batch.rs
index 07e512cd06..ddd171af54 100644
--- a/src/record_batch.rs
+++ b/src/record_batch.rs
@@ -339,17 +339,14 @@ impl Default for RecordBatchOptions {
     }
 }
 
-impl From<&StructArray> for RecordBatch {
-    /// Create a record batch from struct array.
-    ///
-    /// This currently does not flatten and nested struct types
-    fn from(struct_array: &StructArray) -> Self {
-        if let DataType::Struct(fields) = struct_array.data_type() {
-            let schema = Arc::new(Schema::new(fields.clone()));
-            let columns = struct_array.values().to_vec();
-            RecordBatch { schema, columns }
-        } else {
-            unreachable!("unable to get datatype as struct")
+impl From<StructArray> for RecordBatch {
+    /// # Panics iff the null count of the array is not null.
+    fn from(array: StructArray) -> Self {
+        assert!(array.null_count() == 0);
+        let (fields, values, _) = array.into_data();
+        RecordBatch {
+            schema: Arc::new(Schema::new(fields)),
+            columns: values,
         }
     }
 }


Test patch:
diff --git a/tests/it/array/list/mutable.rs b/tests/it/array/list/mutable.rs
index d05ad6f34c..2102735a93 100644
--- a/tests/it/array/list/mutable.rs
+++ b/tests/it/array/list/mutable.rs
@@ -29,3 +29,11 @@ fn basics() {
     );
     assert_eq!(expected, array);
 }
+
+#[test]
+fn push() {
+    let mut array = MutableListArray::<i32, MutablePrimitiveArray<i32>>::new();
+    array
+        .try_push(Some(vec![Some(1i32), Some(2), Some(3)]))
+        .unwrap();
+}
diff --git a/tests/it/record_batch.rs b/tests/it/record_batch.rs
index b9ca197c06..8bda0484b8 100644
--- a/tests/it/record_batch.rs
+++ b/tests/it/record_batch.rs
@@ -91,22 +91,18 @@ fn number_of_fields_mismatch() {
 fn from_struct_array() {
     let boolean = Arc::new(BooleanArray::from_slice(&[false, false, true, true])) as ArrayRef;
     let int = Arc::new(Int32Array::from_slice(&[42, 28, 19, 31])) as ArrayRef;
-    let struct_array = StructArray::from_data(
-        vec![
-            Field::new("b", DataType::Boolean, false),
-            Field::new("c", DataType::Int32, false),
-        ],
-        vec![boolean.clone(), int.clone()],
-        None,
-    );
 
-    let batch = RecordBatch::from(&struct_array);
+    let fields = vec![
+        Field::new("b", DataType::Boolean, false),
+        Field::new("c", DataType::Int32, false),
+    ];
+
+    let array = StructArray::from_data(fields.clone(), vec![boolean.clone(), int.clone()], None);
+
+    let batch = RecordBatch::from(array);
     assert_eq!(2, batch.num_columns());
     assert_eq!(4, batch.num_rows());
-    assert_eq!(
-        struct_array.data_type(),
-        &DataType::Struct(batch.schema().fields().to_vec())
-    );
+    assert_eq!(&fields, batch.schema().fields());
     assert_eq!(boolean.as_ref(), batch.column(0).as_ref());
     assert_eq!(int.as_ref(), batch.column(1).as_ref());
 }