bool array slice to support device buffers

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::<u64>());
+        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();

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<usize>) -> VortexResult<Option<ArrayRef>> {
-        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()))
     }
 }

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<usize>) -> 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(),
         }
     }
 

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>) -> (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 {

vortex-cuda/src/stream.rs

@@ -259,6 +259,9 @@ fn register_stream_callback(stream: &CudaStream) -> VortexResult<kanal::AsyncRec
 mod tests {
     use std::mem::size_of;
 
+    use vortex::array::IntoArray;
+    use vortex::array::arrays::BoolArray;
+    use vortex::array::validity::Validity;
     use vortex::error::VortexResult;
 
     use super::padded_device_allocation_len;
@@ -327,4 +330,20 @@ mod tests {
 
         Ok(())
     }
+
+    #[crate::test]
+    async fn test_slice_device_bool_preserves_device_buffer() -> 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(())
+    }
 }