[GitHub] [arrow] Dandandan commented on a change in pull request #9076: ARROW-11108: [Rust] Fixed performance issue in mutableBuffer.

[GitBox <gi...@apache.org>]

Dandandan commented on a change in pull request #9076:
URL: https://github.com/apache/arrow/pull/9076#discussion_r550932922



##########
File path: rust/arrow/src/buffer.rs
##########
@@ -863,44 +866,49 @@ impl MutableBuffer {
 
     /// View buffer as typed slice.
     pub fn typed_data_mut<T: ArrowNativeType>(&mut self) -> &mut [T] {
-        assert_eq!(self.len() % mem::size_of::<T>(), 0);
-        assert!(memory::is_ptr_aligned::<T>(self.data.cast()));
-        // JUSTIFICATION
-        //  Benefit
-        //      Many of the buffers represent specific types, and consumers of `Buffer` often need to re-interpret them.
-        //  Soundness
-        //      * The pointer is non-null by construction
-        //      * alignment asserted above
         unsafe {
-            std::slice::from_raw_parts_mut(
-                self.as_ptr() as *mut T,
-                self.len() / mem::size_of::<T>(),
-            )
+            let (prefix, offsets, suffix) = self.as_slice_mut().align_to_mut::<T>();
+            assert!(prefix.is_empty() && suffix.is_empty());
+            offsets
         }
     }
 
-    /// Extends the buffer from a byte slice, incrementing its capacity if needed.
+    /// Extends the buffer from a slice, increasing its capacity if needed.
     #[inline]
-    pub fn extend_from_slice(&mut self, bytes: &[u8]) {
-        let new_len = self.len + bytes.len();
+    pub fn extend_from_slice<T: ToByteSlice>(&mut self, items: &[T]) {
+        let additional = items.len() * std::mem::size_of::<T>();
+        let new_len = self.len + additional;
         if new_len > self.capacity {
-            self.reserve(new_len);
+            self.reserve(additional);
         }
         unsafe {
-            let dst = NonNull::new_unchecked(self.data.as_ptr().add(self.len));
-            let src = NonNull::new_unchecked(bytes.as_ptr() as *mut u8);
-            memory::memcpy(dst, src, bytes.len());
+            let dst = self.data.as_ptr().add(self.len) as *mut T;
+            let src = items.as_ptr() as *const T;
+            std::ptr::copy_nonoverlapping(src, dst, items.len())
         }
         self.len = new_len;
     }
 
-    /// Extends the buffer by `len` with all bytes equal to `0u8`, incrementing its capacity if needed.
-    pub fn extend(&mut self, len: usize) {
-        let remaining_capacity = self.capacity - self.len;
-        if len > remaining_capacity {
-            self.reserve(self.len + len);
+    /// Extends the buffer with a new item, increasing its capacity if needed.
+    #[inline]
+    pub fn push<T: ToByteSlice>(&mut self, item: T) {
+        let additional = std::mem::size_of::<T>();
+        let new_len = self.len + additional;
+        if new_len > self.capacity {
+            self.reserve(additional);
+        }
+        unsafe {
+            let dst = self.data.as_ptr().add(self.len) as *mut T;
+            let src = (&item) as *const T;
+            std::ptr::copy_nonoverlapping(src, dst, 1)

Review comment:
       Is this the same as `ptr::write(dst, item)`?




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org