Add COW conversion for Buffer and PrimitiveArray and unary_mut

arrow-array/src/array/mod.rs

@@ -372,6 +372,28 @@ impl<'a, T: Array> Array for &'a T {
     }
 }
 
+/// A kind of `Array` which can be downcasted by `downcast_array`.
+pub(crate) trait SizedArray: AsDynAny + Array {}
+
+/// A trait used to help conversion from `Arc<Self>` to `Arc<Any>`.
+pub(crate) trait AsDynAny: Any {
+    fn as_dyn_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static>;
+}
+
+impl<T: Sized + Send + Sync + 'static> AsDynAny for T {
+    fn as_dyn_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync + 'static> {
+        self
+    }
+}
+
+/// Downcasts an Arc-ed sized `Array` to Arc-ed underlying arrow type.
+#[allow(dead_code)]
+pub(crate) fn downcast_array<T: Send + Sync + 'static>(
+    array: Arc<dyn SizedArray>,
+) -> Option<Arc<T>> {
+    AsDynAny::as_dyn_any(array).downcast().ok()
+}
+
 /// A generic trait for accessing the values of an [`Array`]
 ///
 /// # Validity
@@ -1113,4 +1135,15 @@ mod tests {
         assert!(compute_my_thing(&arr));
         assert!(compute_my_thing(arr.as_ref()));
     }
+
+    #[test]
+    fn test_downcast_array() {
+        let array: Int32Array = vec![1, 2, 3].into_iter().map(Some).collect();
+
+        let boxed: Arc<dyn SizedArray> = Arc::new(array);
+        let array: Arc<Int32Array> = downcast_array(boxed).unwrap();
+
+        let expected: Int32Array = vec![1, 2, 3].into_iter().map(Some).collect();
+        assert_eq!(&array, &Arc::new(expected));
+    }
 }

arrow-array/src/array/primitive_array.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::array::SizedArray;
 use crate::builder::{BooleanBufferBuilder, BufferBuilder, PrimitiveBuilder};
 use crate::iterator::PrimitiveIter;
 use crate::raw_pointer::RawPtrBox;
@@ -489,6 +490,41 @@ impl<T: ArrowPrimitiveType> PrimitiveArray<T> {
             )
         }
     }
+
+    /// Returns `PrimitiveBuilder` of this primitive array for mutating its values if the underlying
+    /// data buffer is not shared by others.
+    pub fn into_builder(self) -> Result<PrimitiveBuilder<T>, Self> {
+        let null_buffer = self
+            .data
+            .null_buffer()
+            .cloned()
+            .and_then(|b| b.into_mutable().ok());
+
+        let len = self.len();
+        let null_bit_buffer = self.data.null_buffer().cloned();
+
+        let mut buffers = self.data.get_buffers();
+        let buffer = buffers.remove(0);
+
+        let builder = buffer
+            .into_mutable()
+            .map(|buffer| PrimitiveBuilder::<T>::new_from_buffer(buffer, null_buffer));
+
+        match builder {
+            Ok(builder) => Ok(builder),
+            Err(buffer) => {
+                let builder = ArrayData::builder(T::DATA_TYPE)
+                    .len(len)
+                    .add_buffer(buffer)
+                    .null_bit_buffer(null_bit_buffer);
+
+                let array_data = unsafe { builder.build_unchecked() };
+                let array = PrimitiveArray::<T>::from(array_data);
+
+                Err(array)
+            }
+        }
+    }
 }
 
 #[inline]
@@ -529,6 +565,9 @@ impl<T: ArrowPrimitiveType> Array for PrimitiveArray<T> {
     }
 }
 
+/// Makes `PrimitiveArray<T>` can be "downcast_array".
+impl<T: ArrowPrimitiveType> SizedArray for PrimitiveArray<T> {}
+
 impl<'a, T: ArrowPrimitiveType> ArrayAccessor for &'a PrimitiveArray<T> {
     type Item = T::Native;
 
@@ -1036,7 +1075,8 @@ impl<T: DecimalType + ArrowPrimitiveType> PrimitiveArray<T> {
 mod tests {
     use super::*;
     use crate::builder::{Decimal128Builder, Decimal256Builder};
-    use crate::BooleanArray;
+    use crate::{downcast_array, BooleanArray};
+    use std::sync::Arc;
 
     #[test]
     fn test_primitive_array_from_vec() {
@@ -1939,4 +1979,42 @@ mod tests {
 
         array.value(4);
     }
+
+    #[test]
+    fn test_into_builder() {
+        let array: Int32Array = vec![1, 2, 3].into_iter().map(Some).collect();
+
+        let boxed: Arc<dyn SizedArray> = Arc::new(array);
+        let array: Arc<Int32Array> = downcast_array(boxed).unwrap();
+
+        let col: Int32Array = Arc::try_unwrap(array).unwrap();
+        let mut builder = col.into_builder().unwrap();
+
+        builder.append_value(4);
+        builder.append_null();
+        builder.append_value(2);
+
+        let expected: Int32Array = vec![Some(4), None, Some(2)].into_iter().collect();
+
+        let new_array = builder.finish();
+        assert_eq!(expected, new_array);
+    }
+
+    #[test]
+    fn test_into_builder_cloned_array() {
+        let array: Int32Array = vec![1, 2, 3].into_iter().map(Some).collect();
+
+        let boxed: Arc<dyn SizedArray> = Arc::new(array);
+
+        let col: Int32Array = PrimitiveArray::<Int32Type>::from(boxed.data().clone());
+        let err = col.into_builder();
+
+        match err {
+            Ok(_) => panic!("Should not get builder from cloned array"),
+            Err(returned) => {
+                let expected: Int32Array = vec![1, 2, 3].into_iter().map(Some).collect();
+                assert_eq!(expected, returned)
+            }
+        }
+    }
 }

arrow-array/src/builder/boolean_buffer_builder.rs

@@ -33,6 +33,10 @@ impl BooleanBufferBuilder {
         Self { buffer, len: 0 }
     }
 
+    pub fn new_from_buffer(buffer: MutableBuffer) -> Self {
+        Self { buffer, len: 0 }
+    }
+
     #[inline]
     pub fn len(&self) -> usize {
         self.len

arrow-array/src/builder/buffer_builder.rs

@@ -124,6 +124,14 @@ impl<T: ArrowNativeType> BufferBuilder<T> {
         }
     }
 
+    pub fn new_from_buffer(buffer: MutableBuffer) -> Self {
+        Self {
+            buffer,
+            len: 0,
+            _marker: PhantomData,
+        }
+    }
+
     /// Returns the current number of array elements in the internal buffer.
     ///
     /// # Example:

arrow-array/src/builder/null_buffer_builder.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 use crate::builder::BooleanBufferBuilder;
-use arrow_buffer::Buffer;
+use arrow_buffer::{Buffer, MutableBuffer};
 
 /// Builder for creating the null bit buffer.
 /// This builder only materializes the buffer when we append `false`.
@@ -42,6 +42,15 @@ impl NullBufferBuilder {
         }
     }
 
+    pub fn new_from_buffer(buffer: Option<MutableBuffer>, capacity: usize) -> Self {
+        let bitmap_builder = buffer.map(BooleanBufferBuilder::new_from_buffer);
+        Self {
+            bitmap_builder,
+            len: 0,
+            capacity,
+        }
+    }
+
     /// Appends `n` `true`s into the builder
     /// to indicate that these `n` items are not nulls.
     #[inline]

arrow-array/src/builder/primitive_builder.rs

@@ -19,6 +19,7 @@ use crate::builder::null_buffer_builder::NullBufferBuilder;
 use crate::builder::{ArrayBuilder, BufferBuilder};
 use crate::types::*;
 use crate::{ArrayRef, ArrowPrimitiveType, PrimitiveArray};
+use arrow_buffer::MutableBuffer;
 use arrow_data::ArrayData;
 use std::any::Any;
 use std::sync::Arc;
@@ -114,6 +115,21 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
         }
     }
 
+    pub fn new_from_buffer(
+        values_buffer: MutableBuffer,
+        null_buffer: Option<MutableBuffer>,
+    ) -> Self {
+        let capacity = values_buffer.capacity();
+
+        Self {
+            values_builder: BufferBuilder::<T::Native>::new_from_buffer(values_buffer),
+            null_buffer_builder: NullBufferBuilder::new_from_buffer(
+                null_buffer,
+                capacity,
+            ),
+        }
+    }
+
     /// Returns the capacity of this builder measured in slots of type `T`
     pub fn capacity(&self) -> usize {
         self.values_builder.capacity()

arrow-buffer/src/buffer/immutable.rs

@@ -19,7 +19,7 @@ use std::fmt::Debug;
 use std::iter::FromIterator;
 use std::ptr::NonNull;
 use std::sync::Arc;
-use std::{convert::AsRef, usize};
+use std::{convert::AsRef, mem, usize};
 
 use crate::alloc::{Allocation, Deallocation};
 use crate::util::bit_chunk_iterator::{BitChunks, UnalignedBitChunk};
@@ -227,6 +227,28 @@ impl Buffer {
     pub fn count_set_bits_offset(&self, offset: usize, len: usize) -> usize {
         UnalignedBitChunk::new(self.as_slice(), offset, len).count_ones()
     }
+
+    /// Returns `MutableBuffer` for mutating the buffer if this buffer is not shared.
+    pub fn into_mutable(self) -> Result<MutableBuffer, Self> {
+        let offset_ptr = self.as_ptr();
+        let offset = self.offset;
+        let length = self.length;
+        Arc::try_unwrap(self.data)
+            .map(|bytes| {
+                // The pointer of underlying buffer should not be offset.
+                assert_eq!(offset_ptr, bytes.ptr().as_ptr());
+
+                let mutable_buffer =
+                    MutableBuffer::from_ptr(bytes.ptr(), bytes.capacity());
+                mem::forget(bytes);
+                mutable_buffer
+            })
+            .map_err(|bytes| Buffer {
+                data: bytes,
+                offset,
+                length,
+            })
+    }
 }
 
 /// Creating a `Buffer` instance by copying the memory from a `AsRef<[u8]>` into a newly

arrow-buffer/src/buffer/mutable.rs

@@ -92,6 +92,15 @@ impl MutableBuffer {
         }
     }
 
+    /// Allocates a new [MutableBuffer] from given pointer `ptr`, `capacity`.
+    pub(crate) fn from_ptr(ptr: NonNull<u8>, capacity: usize) -> Self {
+        Self {
+            data: ptr,
+            len: 0,
+            capacity,
+        }
+    }
+
     /// creates a new [MutableBuffer] with capacity and length capable of holding `len` bits.
     /// This is useful to create a buffer for packed bitmaps.
     pub fn new_null(len: usize) -> Self {

arrow-data/src/data.rs

@@ -387,6 +387,10 @@ impl ArrayData {
         &self.buffers[..]
     }
 
+    pub fn get_buffers(self) -> Vec<Buffer> {
+        self.buffers
+    }
+
     /// Returns a slice of children data arrays
     pub fn child_data(&self) -> &[ArrayData] {
         &self.child_data[..]