From cb1fb2346508ca2eeee44504c33d64b1ac2ceba9 Mon Sep 17 00:00:00 2001 From: Onur Satici Date: Wed, 15 Jul 2026 21:22:41 +0100 Subject: [PATCH 1/3] Preserve device buffers when slicing booleans Signed-off-by: Onur Satici --- vortex-array/src/arrays/bool/array.rs | 23 ++++++++ vortex-array/src/arrays/bool/compute/slice.rs | 14 +++-- vortex-array/src/buffer.rs | 8 +++ vortex-buffer/src/bit/buf.rs | 26 ++------- vortex-buffer/src/bit/meta.rs | 58 +++++++++++++++++++ vortex-cuda/src/stream.rs | 19 ++++++ 6 files changed, 120 insertions(+), 28 deletions(-) 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..f51127f7c99 100644 --- a/vortex-array/src/arrays/bool/compute/slice.rs +++ b/vortex-array/src/arrays/bool/compute/slice.rs @@ -10,18 +10,20 @@ 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; 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(); + + // Slice the BufferHandle itself so device-resident Boolean buffers stay on the device. + // Converting through BitBuffer would require host access and panic for CUDA-flat reads. + let bits = array.bits.slice_unaligned(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-array/src/buffer.rs b/vortex-array/src/buffer.rs index 43039eb52ab..1361b75cc0d 100644 --- a/vortex-array/src/buffer.rs +++ b/vortex-array/src/buffer.rs @@ -201,6 +201,14 @@ impl BufferHandle { } } + /// Creates a new handle to a byte subrange without retaining an alignment guarantee. + pub fn slice_unaligned(&self, range: Range) -> Self { + match &self.0 { + Inner::Host(host) => BufferHandle::new_host(host.slice_unaligned(range)), + Inner::Device(device) => BufferHandle::new_device(device.slice(range)), + } + } + /// Reinterpret the pointee as a buffer of `T` and slice the provided element range. /// /// # Example 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..ac01ed7dd94 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 { @@ -58,3 +91,28 @@ impl BitBufferMeta { (self.offset + self.len).div_ceil(8) } } + +#[cfg(test)] +mod tests { + use super::BitBufferMeta; + + #[test] + fn slice_normalizes_bit_offset() { + let meta = BitBufferMeta::new(3, 18); + + let (byte_offset, sliced) = meta.slice(7..16); + + assert_eq!(byte_offset, 1); + assert_eq!(sliced, BitBufferMeta::new(2, 9)); + assert_eq!(sliced.byte_len(), 2); + } + + #[test] + fn slice_supports_inclusive_and_unbounded_ranges() { + let meta = BitBufferMeta::new(5, 20); + + assert_eq!(meta.slice(..), (0, meta)); + assert_eq!(meta.slice(3..=10), (1, BitBufferMeta::new(0, 8))); + assert_eq!(meta.slice(20..), (3, BitBufferMeta::new(1, 0))); + } +} 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(()) + } } From 0814d70b4c34a7ebe6dee96aaacc4d19ef78f4fe Mon Sep 17 00:00:00 2001 From: Onur Satici Date: Tue, 21 Jul 2026 11:34:59 +0100 Subject: [PATCH 2/3] less tests Signed-off-by: Onur Satici --- vortex-buffer/src/bit/meta.rs | 25 ------------------------- 1 file changed, 25 deletions(-) diff --git a/vortex-buffer/src/bit/meta.rs b/vortex-buffer/src/bit/meta.rs index ac01ed7dd94..75659f3397a 100644 --- a/vortex-buffer/src/bit/meta.rs +++ b/vortex-buffer/src/bit/meta.rs @@ -91,28 +91,3 @@ impl BitBufferMeta { (self.offset + self.len).div_ceil(8) } } - -#[cfg(test)] -mod tests { - use super::BitBufferMeta; - - #[test] - fn slice_normalizes_bit_offset() { - let meta = BitBufferMeta::new(3, 18); - - let (byte_offset, sliced) = meta.slice(7..16); - - assert_eq!(byte_offset, 1); - assert_eq!(sliced, BitBufferMeta::new(2, 9)); - assert_eq!(sliced.byte_len(), 2); - } - - #[test] - fn slice_supports_inclusive_and_unbounded_ranges() { - let meta = BitBufferMeta::new(5, 20); - - assert_eq!(meta.slice(..), (0, meta)); - assert_eq!(meta.slice(3..=10), (1, BitBufferMeta::new(0, 8))); - assert_eq!(meta.slice(20..), (3, BitBufferMeta::new(1, 0))); - } -} From 13205628522929b299f3c1c924141564bed037a5 Mon Sep 17 00:00:00 2001 From: Onur Satici Date: Tue, 21 Jul 2026 11:46:14 +0100 Subject: [PATCH 3/3] Keep unaligned boolean slicing host-specific Signed-off-by: Onur Satici --- vortex-array/src/arrays/bool/compute/slice.rs | 9 ++++++--- vortex-array/src/buffer.rs | 8 -------- 2 files changed, 6 insertions(+), 11 deletions(-) diff --git a/vortex-array/src/arrays/bool/compute/slice.rs b/vortex-array/src/arrays/bool/compute/slice.rs index f51127f7c99..f88eb418da5 100644 --- a/vortex-array/src/arrays/bool/compute/slice.rs +++ b/vortex-array/src/arrays/bool/compute/slice.rs @@ -11,15 +11,18 @@ use crate::array::ArrayView; use crate::arrays::Bool; use crate::arrays::BoolArray; use crate::arrays::slice::SliceReduce; +use crate::buffer::BufferHandle; impl SliceReduce for Bool { fn slice(array: ArrayView<'_, Bool>, range: Range) -> VortexResult> { let (byte_start, meta) = array.meta.slice(range.clone()); let byte_end = byte_start + meta.byte_len(); - // Slice the BufferHandle itself so device-resident Boolean buffers stay on the device. - // Converting through BitBuffer would require host access and panic for CUDA-flat reads. - let bits = array.bits.slice_unaligned(byte_start..byte_end); + 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)?; let array = BoolArray::try_new_from_handle(bits, meta.offset(), meta.len(), validity)?; diff --git a/vortex-array/src/buffer.rs b/vortex-array/src/buffer.rs index 1361b75cc0d..43039eb52ab 100644 --- a/vortex-array/src/buffer.rs +++ b/vortex-array/src/buffer.rs @@ -201,14 +201,6 @@ impl BufferHandle { } } - /// Creates a new handle to a byte subrange without retaining an alignment guarantee. - pub fn slice_unaligned(&self, range: Range) -> Self { - match &self.0 { - Inner::Host(host) => BufferHandle::new_host(host.slice_unaligned(range)), - Inner::Device(device) => BufferHandle::new_device(device.slice(range)), - } - } - /// Reinterpret the pointee as a buffer of `T` and slice the provided element range. /// /// # Example