Preserve valid alignment when slicing buffer handles

vortex-array/src/arrays/primitive/compute/slice.rs

@@ -3,6 +3,7 @@
 
 use std::ops::Range;
 
+use vortex_buffer::Alignment;
 use vortex_error::VortexResult;
 
 use crate::ArrayRef;
@@ -16,7 +17,9 @@ impl SliceReduce for Primitive {
     fn slice(array: ArrayView<'_, Self>, range: Range<usize>) -> VortexResult<Option<ArrayRef>> {
         let byte_width = array.ptype().byte_width();
         let byte_range = range.start * byte_width..range.end * byte_width;
-        let values = array.buffer_handle().slice(byte_range);
+        let values = array
+            .buffer_handle()
+            .slice_with_alignment(byte_range, Alignment::new(byte_width))?;
         let validity = array.validity()?.slice(range)?;
 
         // SAFETY:
@@ -29,3 +32,32 @@ impl SliceReduce for Primitive {
         Ok(Some(array))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use vortex_buffer::Alignment;
+    use vortex_buffer::Buffer;
+
+    use crate::IntoArray;
+    use crate::VortexSessionExecute;
+    use crate::array_session;
+    use crate::arrays::PrimitiveArray;
+    use crate::arrays::primitive::PrimitiveArrayExt;
+    use crate::validity::Validity;
+
+    #[test]
+    fn slice_over_aligned_f32_buffer_at_f32_aligned_offset() -> vortex_error::VortexResult<()> {
+        let values: Vec<f32> = (0..4096).map(|value| value as f32).collect();
+        let buffer = Buffer::copy_from_aligned(values, Alignment::DEFAULT_ALIGNMENT);
+        let array = PrimitiveArray::new(buffer, Validity::NonNullable).into_array();
+
+        let sliced = array.slice(3127..3130)?;
+        let mut ctx = array_session().create_execution_ctx();
+        let sliced = sliced.execute::<PrimitiveArray>(&mut ctx)?;
+
+        assert_eq!(sliced.buffer_handle().alignment(), Alignment::of::<f32>());
+        assert_eq!(sliced.as_slice::<f32>(), &[3127.0, 3128.0, 3129.0]);
+
+        Ok(())
+    }
+}

vortex-array/src/buffer.rs

@@ -13,6 +13,7 @@ use vortex_buffer::Alignment;
 use vortex_buffer::ByteBuffer;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
 use vortex_utils::dyn_traits::DynEq;
 use vortex_utils::dyn_traits::DynHash;
 
@@ -201,6 +202,42 @@ impl BufferHandle {
         }
     }
 
+    /// Creates a new handle to a subrange whose reported alignment is at least `alignment`.
+    ///
+    /// The returned handle reports the strongest alignment still guaranteed by the sliced view.
+    /// Returns an error if the sliced view cannot preserve the requested alignment.
+    pub fn slice_with_alignment(
+        &self,
+        range: Range<usize>,
+        alignment: Alignment,
+    ) -> VortexResult<Self> {
+        match &self.0 {
+            Inner::Host(host) => {
+                let sliced_alignment = offset_alignment(host.alignment(), range.start);
+                if !sliced_alignment.is_aligned_to(alignment) {
+                    vortex_bail!(
+                        "sliced host buffer alignment {sliced_alignment} is weaker than requested \
+                         alignment {alignment}"
+                    );
+                }
+                Ok(BufferHandle::new_host(
+                    host.slice_with_alignment(range, sliced_alignment),
+                ))
+            }
+            Inner::Device(device) => {
+                let sliced = BufferHandle::new_device(device.slice(range));
+                if !sliced.is_aligned_to(alignment) {
+                    vortex_bail!(
+                        "sliced device buffer alignment {} is weaker than requested alignment {}",
+                        sliced.alignment(),
+                        alignment,
+                    );
+                }
+                Ok(sliced)
+            }
+        }
+    }
+
     /// Reinterpret the pointee as a buffer of `T` and slice the provided element range.
     ///
     /// # Example
@@ -215,10 +252,16 @@ impl BufferHandle {
     /// assert_eq!(result, buffer![2, 3, 4]);
     /// ```
     pub fn slice_typed<T: Sized>(&self, range: Range<usize>) -> Self {
+        self.try_slice_typed::<T>(range)
+            .vortex_expect("typed buffer slice should preserve type alignment")
+    }
+
+    /// Reinterpret the pointee as a buffer of `T` and slice the provided element range.
+    pub fn try_slice_typed<T: Sized>(&self, range: Range<usize>) -> VortexResult<Self> {
         let start = range.start * size_of::<T>();
         let end = range.end * size_of::<T>();
 
-        self.slice(start..end)
+        self.slice_with_alignment(start..end, Alignment::of::<T>())
     }
 
     #[expect(clippy::panic)]
@@ -406,6 +449,38 @@ impl BufferHandle {
     }
 }
 
+fn offset_alignment(base: Alignment, offset: usize) -> Alignment {
+    if offset == 0 {
+        return base;
+    }
+    let exponent = base
+        .exponent()
+        .min(u8::try_from(offset.trailing_zeros()).unwrap_or(u8::MAX));
+    Alignment::from_exponent(exponent)
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_buffer::Buffer;
+
+    use super::*;
+
+    #[test]
+    fn host_slice_reports_strongest_supported_view_alignment() -> VortexResult<()> {
+        let values =
+            Buffer::<f32>::copy_from_aligned([0.0, 1.0, 2.0, 3.0], Alignment::DEFAULT_ALIGNMENT);
+        let handle = BufferHandle::new_host(values.into_byte_buffer());
+
+        let aligned = handle.slice_with_alignment(0..8, Alignment::of::<f32>())?;
+        assert_eq!(aligned.alignment(), Alignment::DEFAULT_ALIGNMENT);
+
+        let offset = handle.slice_with_alignment(4..12, Alignment::of::<f32>())?;
+        assert_eq!(offset.alignment(), Alignment::of::<f32>());
+
+        Ok(())
+    }
+}
+
 impl ArrayHash for BufferHandle {
     // TODO(aduffy): implement for array hash
     fn array_hash<H: Hasher>(&self, state: &mut H, accuracy: EqMode) {

vortex-array/src/serde.rs

@@ -647,7 +647,9 @@ impl SerializedArray {
                             segment.len(),
                         );
                     }
-                    segment.slice(start..end).ensure_aligned(alignment)?
+                    segment
+                        .slice_with_alignment(start..end, Alignment::of::<u8>())?
+                        .ensure_aligned(alignment)?
                 };
 
                 offset = end;