Fix BoolArray::new to take BitBuffer (#5690)

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-array/src/arrays/bool/array.rs

@@ -4,7 +4,6 @@
 use arrow_array::BooleanArray;
 use vortex_buffer::BitBuffer;
 use vortex_buffer::BitBufferMut;
-use vortex_buffer::ByteBuffer;
 use vortex_dtype::DType;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
@@ -49,12 +48,21 @@ use crate::validity::Validity;
 #[derive(Clone, Debug)]
 pub struct BoolArray {
     pub(super) dtype: DType,
-    pub(super) buffer: BitBuffer,
+    pub(super) bits: BitBuffer,
     pub(super) validity: Validity,
     pub(super) stats_set: ArrayStats,
 }
 
 impl BoolArray {
+    /// Constructs a new `BoolArray`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the validity length is not equal to the bit buffer length.
+    pub fn new(bits: BitBuffer, validity: Validity) -> Self {
+        Self::try_new(bits, validity).vortex_expect("Failed to create BoolArray")
+    }
+
     /// Constructs a new `BoolArray`.
     ///
     /// See [`BoolArray::new_unchecked`] for more information.
@@ -63,80 +71,51 @@ impl BoolArray {
     ///
     /// Returns an error if the provided components do not satisfy the invariants documented in
     /// [`BoolArray::new_unchecked`].
-    pub fn try_new(
-        buffer: ByteBuffer,
-        offset: usize,
-        len: usize,
-        validity: Validity,
-    ) -> VortexResult<Self> {
-        Self::validate(&buffer, offset, len, &validity)?;
-
-        // SAFETY: validate ensures all invariants are met.
-        Ok(unsafe { Self::new_unchecked(buffer, offset, len, validity) })
+    pub fn try_new(bits: BitBuffer, validity: Validity) -> VortexResult<Self> {
+        let bits = bits.shrink_offset();
+        Self::validate(&bits, &validity)?;
+        Ok(Self {
+            dtype: DType::Bool(validity.nullability()),
+            bits,
+            validity,
+            stats_set: ArrayStats::default(),
+        })
     }
 
     /// Creates a new [`BoolArray`] without validation from these components:
     ///
-    /// * `buffer` is a raw [`ByteBuffer`] holding the packed bits.
-    /// * `offset` is the number of bits in the start of the buffer that should be skipped when
-    ///   looking up the i-th value.
-    /// * `len` is the length of the array, which should correspond to the number of bits.
-    /// * `validity` holds the null values.
-    ///
     /// # Safety
     ///
-    /// The caller must ensure all of the following invariants are satisfied:
-    ///
-    /// - `buffer` must contain at least `(offset + len).div_ceil(8)` bytes.
-    /// - `offset` must be less than 8 (it represents the bit offset within the first byte).
-    /// - If `validity` is `Validity::Array`, its length must exactly equal `len`.
-    pub unsafe fn new_unchecked(
-        buffer: ByteBuffer,
-        offset: usize,
-        len: usize,
-        validity: Validity,
-    ) -> Self {
-        #[cfg(debug_assertions)]
-        Self::validate(&buffer, offset, len, &validity)
-            .vortex_expect("[Debug Assertion]: Invalid `BoolArray` parameters");
-
-        let buffer = BitBuffer::new_with_offset(buffer, len, offset);
-        let buffer = buffer.shrink_offset();
-        Self {
-            dtype: DType::Bool(validity.nullability()),
-            buffer,
-            validity,
-            stats_set: ArrayStats::default(),
+    /// The caller must ensure that the validity length is equal to the bit buffer length.
+    pub unsafe fn new_unchecked(bits: BitBuffer, validity: Validity) -> Self {
+        if cfg!(debug_assertions) {
+            Self::new(bits, validity)
+        } else {
+            Self {
+                dtype: DType::Bool(validity.nullability()),
+                bits,
+                validity,
+                stats_set: ArrayStats::default(),
+            }
         }
     }
 
     /// Validates the components that would be used to create a [`BoolArray`].
     ///
     /// This function checks all the invariants required by [`BoolArray::new_unchecked`].
-    pub fn validate(
-        buffer: &ByteBuffer,
-        offset: usize,
-        len: usize,
-        validity: &Validity,
-    ) -> VortexResult<()> {
-        vortex_ensure!(
-            offset < 8,
-            "offset must be less than whole byte, was {offset} bits"
-        );
-
-        // Validate the buffer is large enough to hold all the bits
-        let required_bytes = offset.saturating_add(len).div_ceil(8);
+    pub fn validate(bits: &BitBuffer, validity: &Validity) -> VortexResult<()> {
         vortex_ensure!(
-            buffer.len() >= required_bytes,
-            "BoolArray with offset={offset} len={len} cannot be built from buffer of size {}",
-            buffer.len()
+            bits.offset() < 8,
+            "BitBuffer offset must be <8, got {}",
+            bits.offset()
         );
 
         // Validate validity
         if let Some(validity_len) = validity.maybe_len() {
             vortex_ensure!(
-                validity_len == len,
-                "BoolArray of size {len} cannot be built with validity of size {validity_len}"
+                validity_len == bits.len(),
+                "BoolArray of size {} cannot be built with validity of size {validity_len}",
+                bits.len()
             );
         }
 
@@ -157,7 +136,7 @@ impl BoolArray {
         let buffer = buffer.shrink_offset();
         Self {
             dtype: DType::Bool(validity.nullability()),
-            buffer,
+            bits: buffer,
             validity,
             stats_set: ArrayStats::default(),
         }
@@ -179,16 +158,16 @@ impl BoolArray {
     /// Returns the underlying [`BitBuffer`] of the array.
     pub fn bit_buffer(&self) -> &BitBuffer {
         assert!(
-            self.buffer.offset() < 8,
+            self.bits.offset() < 8,
             "Offset must be <8, did we forget to call shrink_offset? Found {}",
-            self.buffer.offset()
+            self.bits.offset()
         );
-        &self.buffer
+        &self.bits
     }
 
     /// Returns the underlying [`BitBuffer`] ofthe array
     pub fn into_bit_buffer(self) -> BitBuffer {
-        self.buffer
+        self.bits
     }
 
     pub fn to_mask(&self) -> Mask {
@@ -247,10 +226,7 @@ impl FromIterator<Option<bool>> for BoolArray {
 
 impl IntoArray for BitBuffer {
     fn into_array(self) -> ArrayRef {
-        let (offset, len, buffer) = self.into_inner();
-        BoolArray::try_new(buffer, offset, len, Validity::NonNullable)
-            .vortex_expect("known correct")
-            .into_array()
+        BoolArray::new(self, Validity::NonNullable).into_array()
     }
 }
 

vortex-array/src/arrays/bool/vtable/array.rs

@@ -15,7 +15,7 @@ use crate::vtable::BaseArrayVTable;
 
 impl BaseArrayVTable<BoolVTable> for BoolVTable {
     fn len(array: &BoolArray) -> usize {
-        array.buffer.len()
+        array.bits.len()
     }
 
     fn dtype(array: &BoolArray) -> &DType {

vortex-array/src/arrays/bool/vtable/mod.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use vortex_buffer::BitBuffer;
 use vortex_buffer::BufferHandle;
 use vortex_dtype::DType;
 use vortex_error::VortexExpect;
@@ -104,7 +105,9 @@ impl VTable for BoolVTable {
         };
 
         let buffer = buffers[0].clone().try_to_bytes()?;
-        BoolArray::try_new(buffer, metadata.offset as usize, len, validity)
+        let bits = BitBuffer::new_with_offset(buffer, len, metadata.offset as usize);
+
+        BoolArray::try_new(bits, validity)
     }
 
     fn bind_kernel(array: &Self::Array, _ctx: &mut BindCtx) -> VortexResult<KernelRef> {

vortex-array/src/arrays/validation_tests.rs

@@ -22,22 +22,6 @@ mod tests {
     use crate::arrays::*;
     use crate::validity::Validity;
 
-    #[test]
-    fn test_bool_array_validation_success() {
-        // Valid case: buffer has enough bytes for the array.
-        let buffer = ByteBuffer::from(vec![0xff, 0xff]);
-        let result = BoolArray::try_new(buffer, 0, 10, Validity::NonNullable);
-        assert!(result.is_ok());
-    }
-
-    #[test]
-    fn test_bool_array_validation_failure_buffer_too_small() {
-        // Invalid case: buffer too small.
-        let buffer = ByteBuffer::from(vec![0xff]); // Only 1 byte.
-        let result = BoolArray::try_new(buffer, 0, 16, Validity::NonNullable); // Needs 2 bytes.
-        assert!(result.is_err());
-    }
-
     #[test]
     fn test_chunked_array_validation_success() {
         // Valid case: all chunks have the same dtype.