From e78aecdd031a092f3a5438dcb79bce7b48abd77f Mon Sep 17 00:00:00 2001
From: Liang-Chi Hsieh <viirya@gmail.com>
Date: Tue, 15 Nov 2022 23:53:34 -0800
Subject: [PATCH] Fix test and more for review

---
 arrow-array/src/array/primitive_array.rs      | 43 ++++++++++++++-----
 .../src/builder/boolean_buffer_builder.rs     |  4 +-
 .../src/builder/null_buffer_builder.rs        | 17 +++++++-
 arrow-array/src/builder/primitive_builder.rs  |  6 +--
 4 files changed, 51 insertions(+), 19 deletions(-)

diff --git a/arrow-array/src/array/primitive_array.rs b/arrow-array/src/array/primitive_array.rs
index 646464146f25..16a4b0949535 100644
--- a/arrow-array/src/array/primitive_array.rs
+++ b/arrow-array/src/array/primitive_array.rs
@@ -523,12 +523,6 @@ 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();
 
@@ -537,13 +531,36 @@ impl<T: ArrowPrimitiveType> PrimitiveArray<T> {
 
         drop(self.data);
 
-        let builder = buffer
-            .into_mutable()
-            .map(|buffer| PrimitiveBuilder::<T>::new_from_buffer(buffer, null_buffer));
+        let try_mutable_null_buffer = match null_bit_buffer {
+            None => Ok(None),
+            Some(null_buffer) => {
+                // Null buffer exists, tries to make it mutable
+                null_buffer.into_mutable().map(|b| Some(b))
+            }
+        };
+
+        let try_mutable_buffers = if try_mutable_null_buffer.is_ok() {
+            // Got mutable null buffer, tries to get mutable value buffer
+            let mutable_null_buffer = try_mutable_null_buffer.unwrap();
+            let try_mutable_buffer = buffer.into_mutable();
+
+            // try_mutable_buffer.map(...).map_err(...) doesn't work as the compiler complains
+            // mutable_null_buffer is moved into map closure.
+            match try_mutable_buffer {
+                Ok(mutable_buffer) => Ok(PrimitiveBuilder::<T>::new_from_buffer(
+                    mutable_buffer,
+                    mutable_null_buffer,
+                )),
+                Err(buffer) => Err((buffer, mutable_null_buffer.map(|b| b.into()))),
+            }
+        } else {
+            // Unable to get mutable null buffer
+            Err((buffer, Some(try_mutable_null_buffer.unwrap_err())))
+        };
 
-        match builder {
+        match try_mutable_buffers {
             Ok(builder) => Ok(builder),
-            Err(buffer) => {
+            Err((buffer, null_bit_buffer)) => {
                 let builder = ArrayData::builder(T::DATA_TYPE)
                     .len(len)
                     .add_buffer(buffer)
@@ -2058,5 +2075,9 @@ mod tests {
         let expected: Int32Array = vec![3, 5, 7].into_iter().map(Some).collect();
 
         assert_eq!(expected, c);
+
+        let array: Int32Array = Int32Array::from(vec![Some(5), Some(7), None]);
+        let c = array.unary_mut(|x| x * 2 + 1).unwrap();
+        assert_eq!(c, Int32Array::from(vec![Some(11), Some(15), None]));
     }
 }
diff --git a/arrow-array/src/builder/boolean_buffer_builder.rs b/arrow-array/src/builder/boolean_buffer_builder.rs
index f2a463ad87de..4b8af38cbc23 100644
--- a/arrow-array/src/builder/boolean_buffer_builder.rs
+++ b/arrow-array/src/builder/boolean_buffer_builder.rs
@@ -33,8 +33,8 @@ impl BooleanBufferBuilder {
         Self { buffer, len: 0 }
     }
 
-    pub fn new_from_buffer(buffer: MutableBuffer) -> Self {
-        Self { buffer, len: 0 }
+    pub fn new_from_buffer(buffer: MutableBuffer, len: usize) -> Self {
+        Self { buffer, len }
     }
 
     #[inline]
diff --git a/arrow-array/src/builder/null_buffer_builder.rs b/arrow-array/src/builder/null_buffer_builder.rs
index 623445542d2f..fef7214d5aa7 100644
--- a/arrow-array/src/builder/null_buffer_builder.rs
+++ b/arrow-array/src/builder/null_buffer_builder.rs
@@ -42,9 +42,22 @@ impl NullBufferBuilder {
         }
     }
 
+    /// Creates a new builder with given length.
+    pub fn new_with_len(len: usize) -> Self {
+        Self {
+            bitmap_builder: None,
+            len,
+            capacity: len,
+        }
+    }
+
     /// Creates a new builder from a `MutableBuffer`.
-    pub fn new_from_buffer(buffer: MutableBuffer, len: usize, capacity: usize) -> Self {
-        let bitmap_builder = Some(BooleanBufferBuilder::new_from_buffer(buffer));
+    pub fn new_from_buffer(buffer: MutableBuffer, len: usize) -> Self {
+        let capacity = buffer.len() * 8;
+
+        assert!(len < capacity);
+
+        let bitmap_builder = Some(BooleanBufferBuilder::new_from_buffer(buffer, len));
         Self {
             bitmap_builder,
             len,
diff --git a/arrow-array/src/builder/primitive_builder.rs b/arrow-array/src/builder/primitive_builder.rs
index 064cf081774e..55d8bac0189f 100644
--- a/arrow-array/src/builder/primitive_builder.rs
+++ b/arrow-array/src/builder/primitive_builder.rs
@@ -119,15 +119,13 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
         values_buffer: MutableBuffer,
         null_buffer: Option<MutableBuffer>,
     ) -> Self {
-        let capacity = values_buffer.capacity();
-
         let values_builder = BufferBuilder::<T::Native>::new_from_buffer(values_buffer);
 
         let null_buffer_builder = null_buffer
             .map(|buffer| {
-                NullBufferBuilder::new_from_buffer(buffer, values_builder.len(), capacity)
+                NullBufferBuilder::new_from_buffer(buffer, values_builder.len())
             })
-            .unwrap_or_else(|| NullBufferBuilder::new(capacity));
+            .unwrap_or_else(|| NullBufferBuilder::new_with_len(values_builder.len()));
 
         Self {
             values_builder,