Add MutableArray::as_box (#450)

jorgecarleitao · Sep 26, 2021 · f761d41 · f761d41
1 parent db92910
commit f761d41
Show file tree

Hide file tree

Showing 9 changed files with 76 additions and 2 deletions.
diff --git a/src/array/binary/mutable.rs b/src/array/binary/mutable.rs
@@ -146,6 +146,15 @@ impl<O: Offset> MutableArray for MutableBinaryArray<O> {
         self.validity.as_ref()
     }
 
+    fn as_box(&mut self) -> Box<dyn Array> {
+        Box::new(BinaryArray::from_data(
+            self.data_type.clone(),
+            std::mem::take(&mut self.offsets).into(),
+            std::mem::take(&mut self.values).into(),
+            std::mem::take(&mut self.validity).map(|x| x.into()),
+        ))
+    }
+
     fn as_arc(&mut self) -> Arc<dyn Array> {
         Arc::new(BinaryArray::from_data(
             self.data_type.clone(),

diff --git a/src/array/boolean/mutable.rs b/src/array/boolean/mutable.rs
@@ -338,6 +338,14 @@ impl MutableArray for MutableBooleanArray {
         self.validity.as_ref()
     }
 
+    fn as_box(&mut self) -> Box<dyn Array> {
+        Box::new(BooleanArray::from_data(
+            self.data_type.clone(),
+            std::mem::take(&mut self.values).into(),
+            std::mem::take(&mut self.validity).map(|x| x.into()),
+        ))
+    }
+
     fn as_arc(&mut self) -> Arc<dyn Array> {
         Arc::new(BooleanArray::from_data(
             self.data_type.clone(),

diff --git a/src/array/dictionary/mutable.rs b/src/array/dictionary/mutable.rs
@@ -115,6 +115,13 @@ impl<K: DictionaryKey, M: 'static + MutableArray> MutableArray for MutableDictio
         self.keys.validity()
     }
 
+    fn as_box(&mut self) -> Box<dyn Array> {
+        Box::new(DictionaryArray::<K>::from_data(
+            std::mem::take(&mut self.keys).into(),
+            self.values.as_arc(),
+        ))
+    }
+
     fn as_arc(&mut self) -> Arc<dyn Array> {
         Arc::new(DictionaryArray::<K>::from_data(
             std::mem::take(&mut self.keys).into(),

diff --git a/src/array/fixed_size_binary/mutable.rs b/src/array/fixed_size_binary/mutable.rs
@@ -173,6 +173,14 @@ impl MutableArray for MutableFixedSizeBinaryArray {
         self.validity.as_ref()
     }
 
+    fn as_box(&mut self) -> Box<dyn Array> {
+        Box::new(FixedSizeBinaryArray::from_data(
+            DataType::FixedSizeBinary(self.size as i32),
+            std::mem::take(&mut self.values).into(),
+            std::mem::take(&mut self.validity).map(|x| x.into()),
+        ))
+    }
+
     fn as_arc(&mut self) -> Arc<dyn Array> {
         Arc::new(FixedSizeBinaryArray::from_data(
             DataType::FixedSizeBinary(self.size as i32),

diff --git a/src/array/fixed_size_list/mutable.rs b/src/array/fixed_size_list/mutable.rs
@@ -69,6 +69,14 @@ impl<M: MutableArray + 'static> MutableArray for MutableFixedSizeListArray<M> {
         self.validity.as_ref()
     }
 
+    fn as_box(&mut self) -> Box<dyn Array> {
+        Box::new(FixedSizeListArray::from_data(
+            self.data_type.clone(),
+            self.values.as_arc(),
+            std::mem::take(&mut self.validity).map(|x| x.into()),
+        ))
+    }
+
     fn as_arc(&mut self) -> Arc<dyn Array> {
         Arc::new(FixedSizeListArray::from_data(
             self.data_type.clone(),

diff --git a/src/array/list/mutable.rs b/src/array/list/mutable.rs
@@ -177,6 +177,15 @@ impl<O: Offset, M: MutableArray + 'static> MutableArray for MutableListArray<O,
         self.validity.as_ref()
     }
 
+    fn as_box(&mut self) -> Box<dyn Array> {
+        Box::new(ListArray::from_data(
+            self.data_type.clone(),
+            std::mem::take(&mut self.offsets).into(),
+            self.values.as_arc(),
+            std::mem::take(&mut self.validity).map(|x| x.into()),
+        ))
+    }
+
     fn as_arc(&mut self) -> Arc<dyn Array> {
         Arc::new(ListArray::from_data(
             self.data_type.clone(),

diff --git a/src/array/mod.rs b/src/array/mod.rs
@@ -121,8 +121,16 @@ pub trait MutableArray: std::fmt::Debug {
     /// The optional validity of the array.
     fn validity(&self) -> Option<&MutableBitmap>;
 
-    /// Convert itself to an (immutable) [`Array`].
-    fn as_arc(&mut self) -> Arc<dyn Array>;
+    /// Convert itself to an (immutable) ['Array'].
+    fn as_box(&mut self) -> Box<dyn Array>;
+
+    /// Convert itself to an (immutable) atomically reference counted [`Array`].
+    // This provided implementation has an extra allocation as it first
+    // boxes `self`, then converts the box into an `Arc`. Implementors may wish
+    // to avoid an allocation by skipping the box completely.
+    fn as_arc(&mut self) -> Arc<dyn Array> {
+        self.as_box().into()
+    }
 
     /// Convert to `Any`, to enable dynamic casting.
     fn as_any(&self) -> &dyn Any;

diff --git a/src/array/primitive/mutable.rs b/src/array/primitive/mutable.rs
@@ -326,6 +326,14 @@ impl<T: NativeType> MutableArray for MutablePrimitiveArray<T> {
         self.validity.as_ref()
     }
 
+    fn as_box(&mut self) -> Box<dyn Array> {
+        Box::new(PrimitiveArray::from_data(
+            self.data_type.clone(),
+            std::mem::take(&mut self.values).into(),
+            std::mem::take(&mut self.validity).map(|x| x.into()),
+        ))
+    }
+
     fn as_arc(&mut self) -> Arc<dyn Array> {
         Arc::new(PrimitiveArray::from_data(
             self.data_type.clone(),

diff --git a/src/array/utf8/mutable.rs b/src/array/utf8/mutable.rs
@@ -187,6 +187,15 @@ impl<O: Offset> MutableArray for MutableUtf8Array<O> {
         self.validity.as_ref()
     }
 
+    fn as_box(&mut self) -> Box<dyn Array> {
+        Box::new(Utf8Array::from_data(
+            Self::default_data_type(),
+            std::mem::take(&mut self.offsets).into(),
+            std::mem::take(&mut self.values).into(),
+            std::mem::take(&mut self.validity).map(|x| x.into()),
+        ))
+    }
+
     fn as_arc(&mut self) -> Arc<dyn Array> {
         Arc::new(Utf8Array::from_data(
             Self::default_data_type(),