More to into_mut implementations (#801)

jorgecarleitao · Feb 3, 2022 · 89d8e09 · 89d8e09
1 parent 374b8d7
commit 89d8e09
Show file tree

Hide file tree

Showing 4 changed files with 127 additions and 0 deletions.
diff --git a/src/array/boolean/mod.rs b/src/array/boolean/mod.rs
@@ -2,6 +2,7 @@ use crate::{
     bitmap::Bitmap,
     datatypes::{DataType, PhysicalType},
 };
+use either::Either;
 
 use super::{display_fmt, Array};
 
@@ -95,6 +96,42 @@ impl BooleanArray {
         arr.validity = validity;
         arr
     }
+
+    /// Try to convert this `BooleanArray` to a `MutableBooleanArray`
+    pub fn into_mut(self) -> Either<Self, MutableBooleanArray> {
+        use Either::*;
+
+        if let Some(bitmap) = self.validity {
+            match bitmap.into_mut() {
+                Left(bitmap) => Left(BooleanArray::from_data(
+                    self.data_type,
+                    self.values,
+                    Some(bitmap),
+                )),
+                Right(mutable_bitmap) => match self.values.into_mut() {
+                    Left(immutable) => Left(BooleanArray::from_data(
+                        self.data_type,
+                        immutable,
+                        Some(mutable_bitmap.into()),
+                    )),
+                    Right(mutable) => Right(MutableBooleanArray::from_data(
+                        self.data_type,
+                        mutable,
+                        Some(mutable_bitmap),
+                    )),
+                },
+            }
+        } else {
+            match self.values.into_mut() {
+                Left(immutable) => Left(BooleanArray::from_data(self.data_type, immutable, None)),
+                Right(mutable) => Right(MutableBooleanArray::from_data(
+                    self.data_type,
+                    mutable,
+                    None,
+                )),
+            }
+        }
+    }
 }
 
 // accessors

diff --git a/src/array/utf8/mod.rs b/src/array/utf8/mod.rs
@@ -1,4 +1,5 @@
 use crate::{bitmap::Bitmap, buffer::Buffer, datatypes::DataType};
+use either::Either;
 
 use super::{
     display_fmt,
@@ -177,6 +178,80 @@ impl<O: Offset> Utf8Array<O> {
         arr.validity = validity;
         arr
     }
+
+    /// Try to convert this `Utf8Array` to a `MutableUtf8Array`
+    pub fn into_mut(self) -> Either<Self, MutableUtf8Array<O>> {
+        use Either::*;
+        if let Some(bitmap) = self.validity {
+            match bitmap.into_mut() {
+                Left(bitmap) => Left(Utf8Array::from_data(
+                    self.data_type,
+                    self.offsets,
+                    self.values,
+                    Some(bitmap),
+                )),
+                Right(mutable_bitmap) => match (self.values.get_vec(), self.offsets.get_vec()) {
+                    (Left(immutable_values), Left(immutable_offsets)) => {
+                        Left(Utf8Array::from_data(
+                            self.data_type,
+                            immutable_offsets,
+                            immutable_values,
+                            Some(mutable_bitmap.into()),
+                        ))
+                    }
+                    (Left(immutable_values), Right(mutable_offsets)) => Left(Utf8Array::from_data(
+                        self.data_type,
+                        mutable_offsets.into(),
+                        immutable_values,
+                        Some(mutable_bitmap.into()),
+                    )),
+                    (Right(mutable_values), Left(immutable_offsets)) => Left(Utf8Array::from_data(
+                        self.data_type,
+                        immutable_offsets,
+                        mutable_values.into(),
+                        Some(mutable_bitmap.into()),
+                    )),
+                    (Right(mutable_values), Right(mutable_offsets)) => {
+                        Right(MutableUtf8Array::from_data(
+                            self.data_type,
+                            mutable_offsets,
+                            mutable_values,
+                            Some(mutable_bitmap),
+                        ))
+                    }
+                },
+            }
+        } else {
+            match (self.values.get_vec(), self.offsets.get_vec()) {
+                (Left(immutable_values), Left(immutable_offsets)) => Left(Utf8Array::from_data(
+                    self.data_type,
+                    immutable_offsets,
+                    immutable_values,
+                    None,
+                )),
+                (Left(immutable_values), Right(mutable_offsets)) => Left(Utf8Array::from_data(
+                    self.data_type,
+                    mutable_offsets.into(),
+                    immutable_values,
+                    None,
+                )),
+                (Right(mutable_values), Left(immutable_offsets)) => Left(Utf8Array::from_data(
+                    self.data_type,
+                    immutable_offsets,
+                    mutable_values.into(),
+                    None,
+                )),
+                (Right(mutable_values), Right(mutable_offsets)) => {
+                    Right(MutableUtf8Array::from_data(
+                        self.data_type,
+                        mutable_offsets,
+                        mutable_values,
+                        None,
+                    ))
+                }
+            }
+        }
+    }
 }
 
 // Accessors

diff --git a/tests/it/array/utf8/mod.rs b/tests/it/array/utf8/mod.rs
@@ -1,6 +1,7 @@
 use arrow2::{array::*, bitmap::Bitmap, buffer::Buffer, datatypes::DataType, error::Result};
 
 mod mutable;
+mod to_mutable;
 
 #[test]
 fn basics() {

diff --git a/tests/it/array/utf8/to_mutable.rs b/tests/it/array/utf8/to_mutable.rs
@@ -0,0 +1,14 @@
+use arrow2::array::Utf8Array;
+
+#[test]
+#[allow(clippy::redundant_clone)]
+fn array_to_mutable() {
+    let array = Utf8Array::<i32>::from(&[Some("hello"), Some(" "), None]);
+    let mutable = array.into_mut().unwrap_right();
+
+    let array: Utf8Array<i32> = mutable.into();
+    let array2 = array.clone();
+    let maybe_mut = array2.into_mut();
+    // the ref count is 2 we should not get a mutable.
+    assert!(maybe_mut.is_left())
+}