From e3438921ce65d8e0af4d78e71d2cd4cb6d79441f Mon Sep 17 00:00:00 2001
From: jackylee-ch <qcsd2011@gmail.com>
Date: Sun, 4 Oct 2026 09:13:36 +0800
Subject: [PATCH] fix(buffer): keep map_each_in_place output aligned to the
 mapped type

`map_each_in_place` transmuted `BufferMut<T>` to `BufferMut<R>` while keeping
`T`'s alignment, so mapping to a type with stricter alignment (e.g. `[u8; 4]`
to `u32`) produced an under-aligned buffer, and reading it back as `R` was UB.

Raise the alignment to `Alignment::of::<R>()` after mapping, reusing `aligned`
(metadata-only when the data is already aligned). This requires `R: Copy`,
which every caller already satisfies.

The regression test is taken from the issue.

Fixes #9746

Signed-off-by: jackylee-ch <qcsd2011@gmail.com>
---
 vortex-buffer/benches/vortex_buffer.rs |  4 ++--
 vortex-buffer/src/buffer.rs            |  1 +
 vortex-buffer/src/buffer_mut.rs        | 16 +++++++++++++++-
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/vortex-buffer/benches/vortex_buffer.rs b/vortex-buffer/benches/vortex_buffer.rs
index 3ba571751f6..81bae392b14 100644
--- a/vortex-buffer/benches/vortex_buffer.rs
+++ b/vortex-buffer/benches/vortex_buffer.rs
@@ -75,7 +75,7 @@ impl<T: ArrowNativeType, R: ArrowNativeType> MapEach<T, R> for Arrow<ScalarBuffe
     }
 }
 
-impl<T: Copy, R> MapEach<T, R> for Buffer<T> {
+impl<T: Copy, R: Copy> MapEach<T, R> for Buffer<T> {
     type Output = BufferMut<R>;
 
     fn map_each<F>(self, f: F) -> Self::Output
@@ -86,7 +86,7 @@ impl<T: Copy, R> MapEach<T, R> for Buffer<T> {
     }
 }
 
-impl<T: Copy, R> MapEach<T, R> for BufferMut<T> {
+impl<T: Copy, R: Copy> MapEach<T, R> for BufferMut<T> {
     type Output = BufferMut<R>;
 
     fn map_each<F>(self, f: F) -> Self::Output
diff --git a/vortex-buffer/src/buffer.rs b/vortex-buffer/src/buffer.rs
index 9e954bdcde4..ab0ca530f00 100644
--- a/vortex-buffer/src/buffer.rs
+++ b/vortex-buffer/src/buffer.rs
@@ -345,6 +345,7 @@ impl<T> Buffer<T> {
     pub fn map_each_in_place<R, F>(self, mut f: F) -> BufferMut<R>
     where
         T: Copy,
+        R: Copy,
         F: FnMut(T) -> R,
     {
         match self.try_into_mut() {
diff --git a/vortex-buffer/src/buffer_mut.rs b/vortex-buffer/src/buffer_mut.rs
index 136267b11b8..0b45fe23772 100644
--- a/vortex-buffer/src/buffer_mut.rs
+++ b/vortex-buffer/src/buffer_mut.rs
@@ -785,6 +785,7 @@ impl<T> BufferMut<T> {
     pub fn map_each_in_place<R, F>(self, mut f: F) -> BufferMut<R>
     where
         T: Copy,
+        R: Copy,
         F: FnMut(T) -> R,
     {
         assert_eq!(
@@ -796,7 +797,9 @@ impl<T> BufferMut<T> {
         let mut buf: BufferMut<R> = unsafe { std::mem::transmute(self) };
         buf.iter_mut()
             .for_each(|item| *item = f(unsafe { std::mem::transmute_copy(item) }));
-        buf
+        // `transmute` preserves `T`'s alignment, which can be weaker than `R`'s.
+        let alignment = buf.alignment().max(Alignment::of::<R>());
+        buf.aligned(alignment)
     }
 
     /// Return a `BufferMut<T>` with the same data as this one with the given alignment.
@@ -1074,6 +1077,7 @@ fn misaligned_scalar_type(alignment: Alignment, scalar_align: Alignment) -> ! {
 
 #[cfg(test)]
 mod tests {
+    use std::alloc::Layout;
     use std::cell::Cell;
 
     use allocator_api2::alloc::Global;
@@ -1298,6 +1302,16 @@ mod tests {
         assert_eq!(buf.as_slice(), &[1u32, 2, 3]);
     }
 
+    #[test]
+    fn map_each_in_place_keeps_output_alignment_valid() {
+        let bytes = BufferMut::<[u8; 4]>::copy_from_aligned([[1u8, 0, 0, 0]], Alignment::new(1));
+        assert_eq!(1, Layout::new::<[u8; 4]>().align());
+        let words = bytes.map_each_in_place(u32::from_ne_bytes);
+
+        assert_eq!(words.as_slice(), [1]);
+        assert!(words.alignment().is_aligned_to(Alignment::of::<u32>()));
+    }
+
     #[test]
     fn buffer_mut_zeroed() {
         const LEN: usize = 17;