Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

More to into_mut implementations #801

Merged
merged 2 commits into from
Feb 3, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 37 additions & 0 deletions src/array/boolean/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use crate::{
bitmap::Bitmap,
datatypes::{DataType, PhysicalType},
};
use either::Either;

use super::{display_fmt, Array};

Expand Down Expand Up @@ -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
Expand Down
75 changes: 75 additions & 0 deletions src/array/utf8/mod.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use crate::{bitmap::Bitmap, buffer::Buffer, datatypes::DataType};
use either::Either;

use super::{
display_fmt,
Expand Down Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions tests/it/array/utf8/mod.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use arrow2::{array::*, bitmap::Bitmap, buffer::Buffer, datatypes::DataType, error::Result};

mod mutable;
mod to_mutable;

#[test]
fn basics() {
Expand Down
14 changes: 14 additions & 0 deletions tests/it/array/utf8/to_mutable.rs
Original file line number Diff line number Diff line change
@@ -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())
}