Avoid zero-filling IPC reads with typed buffer handling

Author: pchintar

arrow-buffer/src/buffer/immutable.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::alloc::Layout;
+use std::alloc::{Layout, handle_alloc_error};
 use std::fmt::Debug;
 use std::ptr::NonNull;
 use std::sync::Arc;
@@ -84,6 +84,81 @@ pub struct Buffer {
     length: usize,
 }
 
+/// An aligned byte buffer that can be filled through `Read::read_exact` and
+/// converted into [`Buffer`] without copying.
+///
+/// This is used for IPC reads that need Arrow buffer alignment but should not
+/// zero-fill the allocation before reading into it.
+pub struct AlignedVec {
+    ptr: NonNull<u8>,
+    len: usize,
+    layout: Layout,
+}
+
+impl AlignedVec {
+    /// Allocates `len` bytes with the requested alignment.
+    pub fn new(len: usize, align: usize) -> Self {
+        let layout =
+            Layout::from_size_align(len, align).expect("failed to create layout for AlignedVec");
+
+        let ptr = match layout.size() {
+            0 => crate::buffer::dangling_ptr(),
+            _ => {
+                // Safety: `layout` has non-zero size and was constructed above.
+                let raw_ptr = unsafe { std::alloc::alloc(layout) };
+                NonNull::new(raw_ptr).unwrap_or_else(|| handle_alloc_error(layout))
+            }
+        };
+
+        Self { ptr, len, layout }
+    }
+}
+
+impl std::ops::Deref for AlignedVec {
+    type Target = [u8];
+
+    fn deref(&self) -> &[u8] {
+        // Safety: `ptr` points to `len` bytes owned by this AlignedVec.
+        unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
+    }
+}
+
+impl std::ops::DerefMut for AlignedVec {
+    fn deref_mut(&mut self) -> &mut [u8] {
+        // Safety: `ptr` points to `len` bytes owned by this AlignedVec.
+        unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len) }
+    }
+}
+
+impl From<AlignedVec> for Buffer {
+    fn from(value: AlignedVec) -> Self {
+        // Safety: `value.ptr` was allocated with `value.layout`, and the
+        // resulting Bytes will deallocate it with the same layout.
+        let bytes =
+            unsafe { Bytes::new(value.ptr, value.len, Deallocation::Standard(value.layout)) };
+        std::mem::forget(value);
+        Buffer::from(bytes)
+    }
+}
+
+impl From<AlignedVec> for MutableBuffer {
+    fn from(value: AlignedVec) -> Self {
+        let buffer = Buffer::from(value);
+        buffer
+            .into_mutable()
+            .expect("AlignedVec should be uniquely owned")
+    }
+}
+
+impl Drop for AlignedVec {
+    fn drop(&mut self) {
+        if self.layout.size() != 0 {
+            // Safety: `ptr` was allocated with this exact layout in `new`.
+            unsafe { std::alloc::dealloc(self.ptr.as_ptr(), self.layout) }
+        }
+    }
+}
+
 impl Default for Buffer {
     #[inline]
     fn default() -> Self {