Avoid zero-filling IPC reads with typed buffer handling

Author: pchintar

arrow-buffer/src/buffer/immutable.rs

@@ -15,13 +15,15 @@
 // 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::ops::{Deref, DerefMut};
 use std::ptr::NonNull;
 use std::sync::Arc;
 
 use crate::BufferBuilder;
 use crate::alloc::{Allocation, Deallocation};
+use crate::buffer::dangling_ptr;
 use crate::util::bit_chunk_iterator::{BitChunks, UnalignedBitChunk};
 use crate::{bit_util, bytes::Bytes, native::ArrowNativeType};
 
@@ -84,6 +86,89 @@ 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 useful for readers that need Arrow buffer alignment without
+/// first zero-initializing the allocation.
+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 => 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 }
+    }
+}
+
+// Allows callers such as `Read::read_exact` to view the allocated region as
+// bytes after it has been filled.
+impl 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) }
+    }
+}
+
+// Allows callers such as `Read::read_exact` to write directly into the aligned
+// allocation before it is converted into an Arrow buffer.
+impl 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) }
+    }
+}
+
+// Transfers ownership of the aligned allocation into Buffer without copying.
+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)
+    }
+}
+
+// Converts through Buffer so the aligned allocation is still owned through the
+// normal Arrow buffer representation.
+impl From<AlignedVec> for MutableBuffer {
+    fn from(value: AlignedVec) -> Self {
+        let buffer = Buffer::from(value);
+        buffer
+            .into_mutable()
+            .expect("AlignedVec should be uniquely owned")
+    }
+}
+// Frees the allocation if AlignedVec is dropped before ownership is transferred
+// into Buffer.
+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 {