diff --git a/vortex-array/src/arrays/primitive/compute/slice.rs b/vortex-array/src/arrays/primitive/compute/slice.rs index 9ac45e3baae..22dc5e1f770 100644 --- a/vortex-array/src/arrays/primitive/compute/slice.rs +++ b/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) -> VortexResult> { 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 = (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::(&mut ctx)?; + + assert_eq!(sliced.buffer_handle().alignment(), Alignment::of::()); + assert_eq!(sliced.as_slice::(), &[3127.0, 3128.0, 3129.0]); + + Ok(()) + } +} diff --git a/vortex-array/src/buffer.rs b/vortex-array/src/buffer.rs index 43039eb52ab..78bb1f22575 100644 --- a/vortex-array/src/buffer.rs +++ b/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, + alignment: Alignment, + ) -> VortexResult { + 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(&self, range: Range) -> Self { + self.try_slice_typed::(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(&self, range: Range) -> VortexResult { let start = range.start * size_of::(); let end = range.end * size_of::(); - self.slice(start..end) + self.slice_with_alignment(start..end, Alignment::of::()) } #[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::::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::())?; + assert_eq!(aligned.alignment(), Alignment::DEFAULT_ALIGNMENT); + + let offset = handle.slice_with_alignment(4..12, Alignment::of::())?; + assert_eq!(offset.alignment(), Alignment::of::()); + + Ok(()) + } +} + impl ArrayHash for BufferHandle { // TODO(aduffy): implement for array hash fn array_hash(&self, state: &mut H, accuracy: EqMode) { diff --git a/vortex-array/src/serde.rs b/vortex-array/src/serde.rs index f84ec182269..9b0dffc1782 100644 --- a/vortex-array/src/serde.rs +++ b/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::())? + .ensure_aligned(alignment)? }; offset = end;