Utf8Array::into_mut

jorgecarleitao · Feb 3, 2022 · ff8e359 · ff8e359
1 parent 5fdda18
commit ff8e359
Show file tree

Hide file tree

Showing 3 changed files with 89 additions and 0 deletions.
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,13 @@
+use arrow2::array::Utf8Array;
+
+#[test]
+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())
+}