Repository: jorgecarleitao/arrow2
Pull request #1196 (2022-07-31 07:05:36)

Added `MutableStructArray`
Very rough initial implementation of #703. Hopefully it's at least pointing in the right direction.

codecov[bot]: # [Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1196?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) Report
> Merging [#1196](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1196?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (5098426) into [main](https://codecov.io/gh/jorgecarleitao/arrow2/commit/9f6be525fd94623dfd3055d99af2840dbbbf7884?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (9f6be52) will **decrease** coverage by `0.15%`.
> The diff coverage is `41.60%`.

```diff
@@            Coverage Diff             @@
##             main    #1196      +/-   ##
==========================================
- Coverage   83.52%   83.36%   -0.16%     
==========================================
  Files         364      365       +1     
  Lines       36398    36535     +137     
==========================================
+ Hits        30403    30459      +56     
- Misses       5995     6076      +81     
```


| [Impacted Files](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1196?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) | Coverage Δ | |
|---|---|---|
| [src/array/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1196/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L21vZC5ycw==) | `68.54% <ø> (ø)` | |
| [src/array/struct\_/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1196/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L3N0cnVjdF8vbW9kLnJz) | `61.45% <ø> (ø)` | |
| [src/array/struct\_/mutable.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1196/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L3N0cnVjdF8vbXV0YWJsZS5ycw==) | `41.60% <41.60%> (ø)` | |
| [src/io/ipc/read/array/boolean.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1196/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL2lwYy9yZWFkL2FycmF5L2Jvb2xlYW4ucnM=) | `90.74% <0.00%> (-7.41%)` | :arrow_down: |
| [src/io/ipc/read/reader.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1196/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL2lwYy9yZWFkL3JlYWRlci5ycw==) | `96.32% <0.00%> (-0.34%)` | :arrow_down: |
| [src/bitmap/utils/slice\_iterator.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1196/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2JpdG1hcC91dGlscy9zbGljZV9pdGVyYXRvci5ycw==) | `98.78% <0.00%> (+1.21%)` | :arrow_up: |
| [src/array/binary/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1196/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L2JpbmFyeS9tb2QucnM=) | `90.12% <0.00%> (+1.23%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1196?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/1196?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao). Last update [9f6be52...5098426](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1196?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).

jorgecarleitao: Would you be able to rebase against main and use the `reserve`?
hohav: I rebased and updated `reserve`. I believe there's only one open question left, about performance of the validity check in `push`.

Patch:
diff --git a/src/array/mod.rs b/src/array/mod.rs
index 2a2f4d0e9b..951e14d71a 100644
--- a/src/array/mod.rs
+++ b/src/array/mod.rs
@@ -391,7 +391,7 @@ pub use list::{ListArray, ListValuesIter, MutableListArray};
 pub use map::MapArray;
 pub use null::NullArray;
 pub use primitive::*;
-pub use struct_::StructArray;
+pub use struct_::{MutableStructArray, StructArray};
 pub use union::UnionArray;
 pub use utf8::{MutableUtf8Array, Utf8Array, Utf8ValuesIter};
 
diff --git a/src/array/struct_/mod.rs b/src/array/struct_/mod.rs
index a7c88b6fb1..3d25d8c188 100644
--- a/src/array/struct_/mod.rs
+++ b/src/array/struct_/mod.rs
@@ -9,6 +9,8 @@ use super::{new_empty_array, new_null_array, Array};
 mod ffi;
 pub(super) mod fmt;
 mod iterator;
+mod mutable;
+pub use mutable::*;
 
 /// A [`StructArray`] is a nested [`Array`] with an optional validity representing
 /// multiple [`Array`] with the same number of rows.
diff --git a/src/array/struct_/mutable.rs b/src/array/struct_/mutable.rs
new file mode 100644
index 0000000000..cc02146a77
--- /dev/null
+++ b/src/array/struct_/mutable.rs
@@ -0,0 +1,216 @@
+use std::sync::Arc;
+
+use crate::{
+    array::{Array, MutableArray},
+    bitmap::MutableBitmap,
+    datatypes::DataType,
+};
+
+use super::StructArray;
+
+/// Converting a [`MutableStructArray`] into a [`StructArray`] is `O(1)`.
+#[derive(Debug)]
+pub struct MutableStructArray {
+    data_type: DataType,
+    values: Vec<Box<dyn MutableArray>>,
+    validity: Option<MutableBitmap>,
+}
+
+impl From<MutableStructArray> for StructArray {
+    fn from(other: MutableStructArray) -> Self {
+        let validity = if other.validity.as_ref().map(|x| x.unset_bits()).unwrap_or(0) > 0 {
+            other.validity.map(|x| x.into())
+        } else {
+            None
+        };
+
+        StructArray::from_data(
+            other.data_type,
+            other.values.into_iter().map(|mut v| v.as_box()).collect(),
+            validity,
+        )
+    }
+}
+
+impl MutableStructArray {
+    /// Creates a new [`MutableStructArray`].
+    pub fn new(data_type: DataType, values: Vec<Box<dyn MutableArray>>) -> Self {
+        Self::from_data(data_type, values, None)
+    }
+
+    /// Create a [`MutableStructArray`] out of low-end APIs.
+    /// # Panics
+    /// This function panics iff:
+    /// * `data_type` is not [`DataType::Struct`]
+    /// * The inner types of `data_type` are not equal to those of `values`
+    /// * `validity` is not `None` and its length is different from the `values`'s length
+    pub fn from_data(
+        data_type: DataType,
+        values: Vec<Box<dyn MutableArray>>,
+        validity: Option<MutableBitmap>,
+    ) -> Self {
+        match data_type.to_logical_type() {
+            DataType::Struct(ref fields) => assert!(fields
+                .iter()
+                .map(|f| f.data_type())
+                .eq(values.iter().map(|f| f.data_type()))),
+            _ => panic!("StructArray must be initialized with DataType::Struct"),
+        };
+        let self_ = Self {
+            data_type,
+            values,
+            validity,
+        };
+        self_.assert_lengths();
+        self_
+    }
+
+    fn assert_lengths(&self) {
+        let first_len = self.values.first().map(|v| v.len());
+        if let Some(len) = first_len {
+            if !self.values.iter().all(|x| x.len() == len) {
+                let lengths: Vec<_> = self.values.iter().map(|v| v.len()).collect();
+                panic!("StructArray child lengths differ: {:?}", lengths);
+            }
+        }
+        if let Some(validity) = &self.validity {
+            assert_eq!(first_len.unwrap_or(0), validity.len());
+        }
+    }
+
+    /// Extract the low-end APIs from the [`MutableStructArray`].
+    pub fn into_data(self) -> (DataType, Vec<Box<dyn MutableArray>>, Option<MutableBitmap>) {
+        (self.data_type, self.values, self.validity)
+    }
+
+    /// The mutable values
+    pub fn mut_values(&mut self) -> &mut Vec<Box<dyn MutableArray>> {
+        &mut self.values
+    }
+
+    /// The values
+    pub fn values(&self) -> &Vec<Box<dyn MutableArray>> {
+        &self.values
+    }
+
+    /// Return the `i`th child array.
+    pub fn value<A: MutableArray + 'static>(&mut self, i: usize) -> Option<&mut A> {
+        self.values[i].as_mut_any().downcast_mut::<A>()
+    }
+}
+
+impl MutableStructArray {
+    /// Reserves `additional` entries.
+    pub fn reserve(&mut self, additional: usize) {
+        for v in &mut self.values {
+            v.reserve(additional);
+        }
+        if let Some(x) = self.validity.as_mut() {
+            x.reserve(additional)
+        }
+    }
+
+    /// Call this after pushing into each child array.
+    /// # Panics
+    /// This function panics if any of the children does not have exactly one more
+    /// element than before the last call of `push`.
+    pub fn push(&mut self, valid: bool) {
+        match &mut self.validity {
+            Some(validity) => validity.push(valid),
+            None => match valid {
+                true => (),
+                false => self.init_validity(),
+            },
+        };
+        self.assert_lengths();
+    }
+
+    fn push_null(&mut self) {
+        for v in &mut self.values {
+            v.push_null();
+        }
+        self.push(false);
+    }
+
+    fn init_validity(&mut self) {
+        let mut validity = MutableBitmap::with_capacity(self.values.capacity());
+        let len = self.len();
+        if len > 0 {
+            validity.extend_constant(len, true);
+            validity.set(len - 1, false);
+        }
+        self.validity = Some(validity)
+    }
+
+    /// Converts itself into an [`Array`].
+    pub fn into_arc(self) -> Arc<dyn Array> {
+        let a: StructArray = self.into();
+        Arc::new(a)
+    }
+
+    /// Shrinks the capacity of the [`MutableStructArray`] to fit its current length.
+    pub fn shrink_to_fit(&mut self) {
+        for v in &mut self.values {
+            v.shrink_to_fit();
+        }
+        if let Some(validity) = self.validity.as_mut() {
+            validity.shrink_to_fit()
+        }
+    }
+}
+
+impl MutableArray for MutableStructArray {
+    fn len(&self) -> usize {
+        self.values.first().map(|v| v.len()).unwrap_or(0)
+    }
+
+    fn validity(&self) -> Option<&MutableBitmap> {
+        self.validity.as_ref()
+    }
+
+    fn as_box(&mut self) -> Box<dyn Array> {
+        Box::new(StructArray::from_data(
+            self.data_type.clone(),
+            std::mem::take(&mut self.values)
+                .into_iter()
+                .map(|mut v| v.as_box())
+                .collect(),
+            std::mem::take(&mut self.validity).map(|x| x.into()),
+        ))
+    }
+
+    fn as_arc(&mut self) -> Arc<dyn Array> {
+        Arc::new(StructArray::from_data(
+            self.data_type.clone(),
+            std::mem::take(&mut self.values)
+                .into_iter()
+                .map(|mut v| v.as_box())
+                .collect(),
+            std::mem::take(&mut self.validity).map(|x| x.into()),
+        ))
+    }
+
+    fn data_type(&self) -> &DataType {
+        &self.data_type
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
+        self
+    }
+
+    fn push_null(&mut self) {
+        self.push_null()
+    }
+
+    fn shrink_to_fit(&mut self) {
+        self.shrink_to_fit()
+    }
+
+    fn reserve(&mut self, additional: usize) {
+        self.reserve(additional)
+    }
+}


Test patch:
diff --git a/tests/it/array/struct_/mod.rs b/tests/it/array/struct_/mod.rs
index 9ec34eca0e..8869c215e5 100644
--- a/tests/it/array/struct_/mod.rs
+++ b/tests/it/array/struct_/mod.rs
@@ -1,4 +1,5 @@
 mod iterator;
+mod mutable;
 
 use arrow2::array::*;
 use arrow2::bitmap::Bitmap;
diff --git a/tests/it/array/struct_/mutable.rs b/tests/it/array/struct_/mutable.rs
new file mode 100644
index 0000000000..19f2f12f15
--- /dev/null
+++ b/tests/it/array/struct_/mutable.rs
@@ -0,0 +1,33 @@
+use arrow2::{
+    array::*,
+    datatypes::{DataType, Field},
+};
+
+#[test]
+fn push() {
+    let c1 = Box::new(MutablePrimitiveArray::<i32>::new()) as Box<dyn MutableArray>;
+    let values = vec![c1];
+    let data_type = DataType::Struct(vec![Field::new("f1", DataType::Int32, true)]);
+    let mut a = MutableStructArray::new(data_type, values);
+
+    a.value::<MutablePrimitiveArray<i32>>(0)
+        .unwrap()
+        .push(Some(1));
+    a.push(true);
+    a.value::<MutablePrimitiveArray<i32>>(0).unwrap().push(None);
+    a.push(false);
+    a.value::<MutablePrimitiveArray<i32>>(0)
+        .unwrap()
+        .push(Some(2));
+    a.push(true);
+
+    assert_eq!(a.len(), 3);
+    assert!(a.is_valid(0));
+    assert!(!a.is_valid(1));
+    assert!(a.is_valid(2));
+
+    assert_eq!(
+        a.value::<MutablePrimitiveArray<i32>>(0).unwrap().values(),
+        &Vec::from([1, 0, 2])
+    );
+}