save

Signed-off-by: Andrew Duffy <andrew@a10y.dev>

Author: a10y

encodings/bytebool/src/array.rs

@@ -29,8 +29,8 @@ use vortex_array::vtable::VTable;
 use vortex_array::vtable::ValidityVTable;
 use vortex_array::vtable::child_to_validity;
 use vortex_array::vtable::validity_to_child;
-use vortex_buffer::BitBuffer;
 use vortex_buffer::ByteBuffer;
+use vortex_buffer::{BitBuffer, BitBufferMut};
 use vortex_error::VortexExpect as _;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
@@ -151,7 +151,8 @@ impl VTable for ByteBool {
     }
 
     fn execute(array: Array<Self>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
-        let boolean_buffer = BitBuffer::from(array.as_slice());
+        // convert truthy values to set/unset bits
+        let boolean_buffer = BitBufferMut::from(array.truthy_bytes()).freeze();
         let validity = array.validity()?;
         Ok(ExecutionResult::done(
             BoolArray::new(boolean_buffer, validity).into_array(),
@@ -225,7 +226,7 @@ impl ByteBool {
     /// # Safety
     ///
     /// Every byte of `buffer` must be `0x00` or `0x01`. Any other byte value is
-    /// Undefined Behavior because [`ByteBoolData::as_slice`] reinterprets the buffer
+    /// Undefined Behavior because [`ByteBoolData::truthy_bytes`] reinterprets the buffer
     /// as `&[bool]`, and a `bool` with any bit pattern other than 0 or 1 is UB.
     /// If `validity` is [`Validity::Array`], its length must equal `buffer.len()`.
     pub unsafe fn new_unchecked(buffer: BufferHandle, validity: Validity) -> ByteBoolArray {
@@ -295,7 +296,7 @@ impl ByteBoolData {
 
     /// Validate that every byte of `buffer` is `0x00` or `0x01`.
     ///
-    /// [`ByteBoolData::as_slice`] transmutes the buffer's bytes to `&[bool]`; any byte
+    /// [`ByteBoolData::truthy_bytes`] transmutes the buffer's bytes to `&[bool]`; any byte
     /// other than `0x00` or `0x01` would produce a `bool` with an invalid bit pattern,
     /// which is Undefined Behavior per the Rust reference.
     ///
@@ -345,7 +346,7 @@ impl ByteBoolData {
     /// # Safety
     ///
     /// Every byte of `buffer` must be `0x00` or `0x01`. Any other byte value is
-    /// Undefined Behavior because [`ByteBoolData::as_slice`] reinterprets the buffer
+    /// Undefined Behavior because [`ByteBoolData::truthy_bytes`] reinterprets the buffer
     /// as `&[bool]`, and a `bool` with any bit pattern other than 0 or 1 is UB.
     /// If `validity` is [`Validity::Array`], its length must equal `buffer.len()`.
     pub unsafe fn new_unchecked(buffer: BufferHandle, validity: Validity) -> Self {
@@ -394,9 +395,11 @@ impl ByteBoolData {
         &self.buffer
     }
 
-    pub fn as_slice(&self) -> &[bool] {
-        // Safety: The internal buffer contains byte-sized bools
-        unsafe { std::mem::transmute(self.buffer().as_host().as_slice()) }
+    /// Get access to the underlying 8-bit truthy values.
+    ///
+    /// The zero byte indicates `false`, and any non-zero byte is a `true`.
+    pub fn truthy_bytes(&self) -> &[u8] {
+        self.buffer().as_host().as_slice()
     }
 }
 
@@ -514,7 +517,7 @@ mod tests {
         // SAFETY: all bytes are 0 or 1.
         let arr = unsafe { ByteBool::new_unchecked(handle, Validity::NonNullable) };
         assert_eq!(arr.len(), 4);
-        assert_eq!(arr.as_slice(), &[false, true, true, false]);
+        assert_eq!(arr.truthy_bytes(), &[false, true, true, false]);
     }
 
     #[test]

encodings/bytebool/src/compute.rs

@@ -8,11 +8,13 @@ use vortex_array::ExecutionCtx;
 use vortex_array::IntoArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::dict::TakeExecute;
+use vortex_array::buffer::BufferHandle;
 use vortex_array::dtype::DType;
 use vortex_array::match_each_integer_ptype;
 use vortex_array::scalar_fn::fns::cast::CastReduce;
 use vortex_array::scalar_fn::fns::mask::MaskReduce;
 use vortex_array::validity::Validity;
+use vortex_buffer::ByteBuffer;
 use vortex_error::VortexResult;
 
 use super::ByteBool;
@@ -58,23 +60,28 @@ impl TakeExecute for ByteBool {
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
         let indices = indices.clone().execute::<PrimitiveArray>(ctx)?;
-        let bools = array.as_slice();
+        let values = array.truthy_bytes();
 
         // This handles combining validity from both source array and nullable indices
         let validity = array.validity()?.take(&indices.clone().into_array())?;
 
-        let taken_bools = match_each_integer_ptype!(indices.ptype(), |I| {
+        let taken = match_each_integer_ptype!(indices.ptype(), |I| {
             indices
                 .as_slice::<I>()
                 .iter()
                 .map(|&idx| {
                     let idx: usize = idx.as_();
-                    bools[idx]
+                    values[idx]
                 })
-                .collect::<Vec<bool>>()
+                .collect::<ByteBuffer>()
         });
 
-        Ok(Some(ByteBool::from_vec(taken_bools, validity).into_array()))
+        // SAFETY:
+        unsafe {
+            Ok(Some(
+                ByteBool::new_unchecked(BufferHandle::new_host(taken), validity).into_array(),
+            ))
+        }
     }
 }
 

vortex-buffer/src/bit/buf_mut.rs

@@ -573,6 +573,13 @@ impl From<&[bool]> for BitBufferMut {
     }
 }
 
+// allow building a buffer from a set of truthy byte values.
+impl From<&[u8]> for BitBufferMut {
+    fn from(value: &[u8]) -> Self {
+        BitBufferMut::collect_bool(value.len(), |i| value[i] > 0)
+    }
+}
+
 impl From<Vec<bool>> for BitBufferMut {
     fn from(value: Vec<bool>) -> Self {
         value.as_slice().into()