diff --git a/vortex-array/src/arrays/bool/array.rs b/vortex-array/src/arrays/bool/array.rs index 2dd08c55832..9a0258fd3d7 100644 --- a/vortex-array/src/arrays/bool/array.rs +++ b/vortex-array/src/arrays/bool/array.rs @@ -367,9 +367,12 @@ mod tests { use std::iter::once; use std::iter::repeat_n; + use vortex_buffer::Alignment; use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; + use vortex_buffer::ByteBuffer; use vortex_buffer::buffer; + use vortex_error::VortexResult; use crate::IntoArray; use crate::VortexSessionExecute; @@ -378,6 +381,7 @@ mod tests { use crate::arrays::PrimitiveArray; use crate::arrays::bool::BoolArrayExt; use crate::assert_arrays_eq; + use crate::buffer::BufferHandle; use crate::patches::Patches; use crate::validity::Validity; @@ -471,6 +475,25 @@ mod tests { assert_arrays_eq!(sliced, BoolArray::from_iter([true; 8]), &mut ctx); } + #[test] + fn slice_aligned_host_handle_at_unaligned_byte() -> VortexResult<()> { + let bits: ByteBuffer = buffer![0b1010_1100_u8, 0b0110_1001, 0]; + let bits = bits.aligned(Alignment::of::()); + let array = + BoolArray::new_handle(BufferHandle::new_host(bits), 0, 16, Validity::NonNullable) + .into_array(); + + let sliced = array.slice(9..15)?; + + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!( + sliced, + BoolArray::from_iter([false, false, true, false, true, true]), + &mut ctx + ); + Ok(()) + } + #[test] fn patch_bools_owned() { let mut ctx = array_session().create_execution_ctx(); diff --git a/vortex-array/src/arrays/bool/compute/slice.rs b/vortex-array/src/arrays/bool/compute/slice.rs index 0b38d2a6117..f88eb418da5 100644 --- a/vortex-array/src/arrays/bool/compute/slice.rs +++ b/vortex-array/src/arrays/bool/compute/slice.rs @@ -10,18 +10,23 @@ use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::Bool; use crate::arrays::BoolArray; -use crate::arrays::bool::BoolArrayExt; use crate::arrays::slice::SliceReduce; +use crate::buffer::BufferHandle; impl SliceReduce for Bool { fn slice(array: ArrayView<'_, Bool>, range: Range) -> VortexResult> { - let bit_buffer = array.to_bit_buffer().slice(range.clone()); + let (byte_start, meta) = array.meta.slice(range.clone()); + let byte_end = byte_start + meta.byte_len(); + + let bits = if let Some(host) = array.bits.as_host_opt() { + BufferHandle::new_host(host.slice_unaligned(byte_start..byte_end)) + } else { + array.bits.slice(byte_start..byte_end) + }; let validity = array.validity()?.slice(range)?; - // Safety: - // range is verified in the callers and is the same for both bits and validity. - let array = unsafe { BoolArray::new_unchecked(bit_buffer, validity).into_array() }; + let array = BoolArray::try_new_from_handle(bits, meta.offset(), meta.len(), validity)?; - Ok(Some(array)) + Ok(Some(array.into_array())) } } diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs index 8dc5e38a42c..cf69ce35449 100644 --- a/vortex-buffer/src/bit/buf.rs +++ b/vortex-buffer/src/bit/buf.rs @@ -7,11 +7,11 @@ use std::fmt::Result as FmtResult; use std::ops::BitAnd; use std::ops::BitOr; use std::ops::BitXor; -use std::ops::Bound; use std::ops::Not; use std::ops::RangeBounds; use crate::Alignment; +use crate::BitBufferMeta; use crate::BitBufferMut; use crate::Buffer; use crate::BufferMut; @@ -346,25 +346,7 @@ impl BitBuffer { /// /// Panics if the slice would extend beyond the end of the buffer. pub fn slice(&self, range: impl RangeBounds) -> Self { - let start = match range.start_bound() { - Bound::Included(&s) => s, - Bound::Excluded(&s) => s + 1, - Bound::Unbounded => 0, - }; - let end = match range.end_bound() { - Bound::Included(&e) => e + 1, - Bound::Excluded(&e) => e, - Bound::Unbounded => self.len, - }; - - assert!(start <= end); - assert!(start <= self.len); - assert!(end <= self.len); - let len = end - start; - - let offset = self.offset + start; - let byte_offset = offset / 8; - let bit_offset = offset % 8; + let (byte_offset, meta) = BitBufferMeta::new(self.offset, self.len).slice(range); // Trim whole bytes off the front directly rather than going through `new_with_offset`, // which would slice (and re-clone) the clone we'd have to pass it. @@ -376,8 +358,8 @@ impl BitBuffer { Self { buffer, - offset: bit_offset, - len, + offset: meta.offset(), + len: meta.len(), } } diff --git a/vortex-buffer/src/bit/meta.rs b/vortex-buffer/src/bit/meta.rs index 866a308647d..75659f3397a 100644 --- a/vortex-buffer/src/bit/meta.rs +++ b/vortex-buffer/src/bit/meta.rs @@ -1,6 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Bound; +use std::ops::RangeBounds; + +use vortex_error::VortexExpect; + /// In-memory metadata describing a packed bitset: a normalized bit `offset` (always `< 8`) and a /// logical bit `len`. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] @@ -34,6 +39,34 @@ impl BitBufferMeta { ) } + /// Return the leading byte offset and normalized metadata for a logical slice. + /// + /// # Panics + /// + /// Panics if the range is out of bounds or its end precedes its start. + pub fn slice(&self, range: impl RangeBounds) -> (usize, Self) { + let start = match range.start_bound() { + Bound::Included(&start) => start, + Bound::Excluded(&start) => start + .checked_add(1) + .vortex_expect("excluded slice start must not overflow"), + Bound::Unbounded => 0, + }; + let end = match range.end_bound() { + Bound::Included(&end) => end + .checked_add(1) + .vortex_expect("included slice end must not overflow"), + Bound::Excluded(&end) => end, + Bound::Unbounded => self.len, + }; + + assert!(start <= end); + assert!(start <= self.len); + assert!(end <= self.len); + + Self::from_raw_offset(self.offset + start, end - start) + } + /// The sub-byte bit offset. Always `< 8`. #[inline(always)] pub fn offset(&self) -> usize { diff --git a/vortex-cuda/src/stream.rs b/vortex-cuda/src/stream.rs index 4d31a7937f9..d54f8d0940b 100644 --- a/vortex-cuda/src/stream.rs +++ b/vortex-cuda/src/stream.rs @@ -259,6 +259,9 @@ fn register_stream_callback(stream: &CudaStream) -> VortexResult VortexResult<()> { + let ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?; + let bits = ctx + .stream() + .copy_to_device(vec![0b1010_1100_u8, 0b0110_1001, 0b1100_0011])? + .await?; + let array = BoolArray::new_handle(bits, 3, 18, Validity::NonNullable).into_array(); + + let sliced = array.slice(7..16)?; + + assert_eq!(sliced.len(), 9); + assert!(sliced.buffer_handles()[0].is_on_device()); + Ok(()) + } }