Remove BitBuffer::into_mut and require callers to handle failure to acquire unique ownership (#7136)

BitBuffer::into_mut is an escape hatch that will make the callers copy
the whole
buffer while they can implement better logic by fusing their buffer
operation
with the copy

Signed-off-by: Robert Kruszewski <github@robertk.io>

---------

Signed-off-by: Robert Kruszewski <github@robertk.io>

Author: robert3005

fuzz/src/array/fill_null.rs

@@ -34,11 +34,11 @@ pub fn fill_null_canonical_array(
 
     Ok(match canonical {
         Canonical::Null(array) => ConstantArray::new(fill_value.clone(), array.len()).into_array(),
-        Canonical::Bool(array) => fill_bool_array(&array, fill_value, result_nullability),
-        Canonical::Primitive(array) => fill_primitive_array(&array, fill_value, result_nullability),
-        Canonical::Decimal(array) => fill_decimal_array(&array, fill_value, result_nullability),
+        Canonical::Bool(array) => fill_bool_array(array, fill_value, result_nullability),
+        Canonical::Primitive(array) => fill_primitive_array(array, fill_value, result_nullability),
+        Canonical::Decimal(array) => fill_decimal_array(array, fill_value, result_nullability),
         Canonical::VarBinView(array) => {
-            fill_varbinview_array(&array, fill_value, result_nullability)
+            fill_varbinview_array(array, fill_value, result_nullability)
         }
         Canonical::Struct(_)
         | Canonical::List(_)
@@ -48,7 +48,7 @@ pub fn fill_null_canonical_array(
 }
 
 fn fill_bool_array(
-    array: &BoolArray,
+    array: BoolArray,
     fill_value: &Scalar,
     result_nullability: Nullability,
 ) -> ArrayRef {
@@ -59,28 +59,37 @@ fn fill_bool_array(
 
     match array.validity() {
         Validity::NonNullable | Validity::AllValid => {
-            BoolArray::new(array.to_bit_buffer(), result_nullability.into()).into_array()
+            BoolArray::new(array.into_bit_buffer(), result_nullability.into()).into_array()
         }
         Validity::AllInvalid => ConstantArray::new(fill_value.clone(), array.len()).into_array(),
         Validity::Array(validity_array) => {
-            let validity_bool_array = validity_array.to_bool();
-            let validity_bits = validity_bool_array.to_bit_buffer();
-            let data_bits = array.to_bit_buffer();
-
-            let mut new_bits = data_bits.into_mut();
-
-            (!validity_bits)
-                .set_indices()
-                .for_each(|i| new_bits.set_to(i, fill_bool));
+            let validity_bits = validity_array.to_bool().into_bit_buffer();
+            let data_bits = array.into_bit_buffer();
+
+            let new_bits = match data_bits.try_into_mut() {
+                Ok(mut buf) => {
+                    (!validity_bits)
+                        .set_indices()
+                        .for_each(|i| buf.set_to(i, fill_bool));
+                    buf.freeze()
+                }
+                Err(data_bits) => {
+                    if fill_bool {
+                        data_bits | !validity_bits
+                    } else {
+                        data_bits & validity_bits
+                    }
+                }
+            };
 
-            BoolArray::new(new_bits.freeze(), result_nullability.into()).into_array()
+            BoolArray::new(new_bits, result_nullability.into()).into_array()
         }
     }
 }
 
 #[expect(clippy::cognitive_complexity)]
 fn fill_primitive_array(
-    array: &PrimitiveArray,
+    array: PrimitiveArray,
     fill_value: &Scalar,
     result_nullability: Nullability,
 ) -> ArrayRef {
@@ -117,7 +126,7 @@ fn fill_primitive_array(
 }
 
 fn fill_decimal_array(
-    array: &DecimalArray,
+    array: DecimalArray,
     fill_value: &Scalar,
     result_nullability: Nullability,
 ) -> ArrayRef {
@@ -161,7 +170,7 @@ fn fill_decimal_array(
 }
 
 fn fill_varbinview_array(
-    array: &VarBinViewArray,
+    array: VarBinViewArray,
     fill_value: &Scalar,
     result_nullability: Nullability,
 ) -> ArrayRef {

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

@@ -224,7 +224,9 @@ impl BoolArray {
 
     /// Returns the underlying [`BitBuffer`] of the array
     pub fn into_bit_buffer(self) -> BitBuffer {
-        self.to_bit_buffer()
+        let buffer = self.bits.unwrap_host();
+
+        BitBuffer::new_with_offset(buffer, self.len, self.offset)
     }
 
     pub fn to_mask(&self) -> Mask {

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

@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use itertools::Itertools;
+use vortex_buffer::BitBufferMut;
 use vortex_error::VortexResult;
 
 use crate::ExecutionCtx;
@@ -26,7 +27,10 @@ impl BoolArray {
             ctx,
         )?;
 
-        let mut own_values = self.into_bit_buffer().into_mut();
+        let bit_buffer = self.into_bit_buffer();
+        let mut own_values = bit_buffer
+            .try_into_mut()
+            .unwrap_or_else(|bb| BitBufferMut::copy_from(&bb));
         match_each_unsigned_integer_ptype!(indices.ptype(), |I| {
             for (idx, value) in indices
                 .as_slice::<I>()

vortex-array/src/arrays/decimal/compute/fill_null.rs

@@ -36,7 +36,7 @@ impl FillNullKernel for Decimal {
                 let is_invalid = is_valid
                     .clone()
                     .execute::<BoolArray>(ctx)?
-                    .to_bit_buffer()
+                    .into_bit_buffer()
                     .not();
                 let decimal_scalar = fill_value.as_decimal();
                 let decimal_value = decimal_scalar

vortex-array/src/arrays/primitive/compute/fill_null.rs

@@ -31,7 +31,7 @@ impl FillNullKernel for Primitive {
                 let is_invalid = is_valid
                     .clone()
                     .execute::<BoolArray>(ctx)?
-                    .to_bit_buffer()
+                    .into_bit_buffer()
                     .not();
                 match_each_native_ptype!(array.ptype(), |T| {
                     let mut buffer = array.to_buffer::<T>().into_mut();

vortex-buffer/public-api.lock

@@ -302,8 +302,6 @@ impl vortex_buffer::BitBuffer
 
 pub fn vortex_buffer::BitBuffer::into_inner(self) -> (usize, usize, vortex_buffer::ByteBuffer)
 
-pub fn vortex_buffer::BitBuffer::into_mut(self) -> vortex_buffer::BitBufferMut
-
 pub fn vortex_buffer::BitBuffer::try_into_mut(self) -> core::result::Result<vortex_buffer::BitBufferMut, Self>
 
 impl core::clone::Clone for vortex_buffer::BitBuffer

vortex-buffer/src/bit/buf.rs

@@ -26,7 +26,6 @@ use crate::bit::get_bit_unchecked;
 use crate::bit::ops::bitwise_binary_op;
 use crate::bit::ops::bitwise_unary_op;
 use crate::buffer;
-use crate::trusted_len::TrustedLenExt;
 
 /// An immutable bitset stored as a packed byte buffer.
 #[derive(Debug, Clone, Eq)]
@@ -348,12 +347,8 @@ impl BitBuffer {
                 self.len,
             );
         }
-        // Use Chunk iterator here to reset offset to 0
-        let iter = self.chunks().iter_padded();
-        let iter = unsafe { iter.trusted_len() };
-        let result = Buffer::<u64>::from_trusted_len_iter(iter).into_byte_buffer();
 
-        BitBuffer::new(result, self.len())
+        bitwise_unary_op(self.clone(), |a| a)
     }
 }
 
@@ -372,15 +367,6 @@ impl BitBuffer {
             Err(buffer) => Err(BitBuffer::new_with_offset(buffer, self.len, self.offset)),
         }
     }
-
-    /// Get a mutable version of this `BitBuffer` along with bit offset in the first byte.
-    ///
-    /// If the caller doesn't hold only reference to the underlying buffer, a copy is created.
-    /// The second value of the tuple is a bit_offset of the first value in the first byte
-    pub fn into_mut(self) -> BitBufferMut {
-        let (offset, len, inner) = self.into_inner();
-        BitBufferMut::from_buffer(inner.into_mut(), offset, len)
-    }
 }
 
 impl From<&[bool]> for BitBuffer {
@@ -469,7 +455,7 @@ impl Not for &BitBuffer {
 
     #[inline]
     fn not(self) -> Self::Output {
-        bitwise_unary_op(self, |a| !a)
+        !self.clone()
     }
 }
 
@@ -478,7 +464,7 @@ impl Not for BitBuffer {
 
     #[inline]
     fn not(self) -> Self::Output {
-        (&self).not()
+        bitwise_unary_op(self, |a| !a)
     }
 }
 

vortex-buffer/src/bit/buf_mut.rs

@@ -656,7 +656,7 @@ impl Not for BitBufferMut {
 
 impl From<&[bool]> for BitBufferMut {
     fn from(value: &[bool]) -> Self {
-        BitBuffer::collect_bool(value.len(), |i| value[i]).into_mut()
+        BitBufferMut::collect_bool(value.len(), |i| value[i])
     }
 }
 
@@ -1183,7 +1183,7 @@ mod tests {
         for i in 0..10 {
             let buf = bitbuffer![0 1 0 1 0 1 0 1 0 1];
 
-            let mut buf_mut = buf.clone().into_mut();
+            let mut buf_mut = BitBufferMut::copy_from(&buf);
             assert_eq!(buf_mut.len(), 10);
 
             let tail = buf_mut.split_off(i);
@@ -1201,7 +1201,7 @@ mod tests {
         for i in 0..10 {
             let buf = bitbuffer![0 1 0 1 0 1 0 1 0 1 0 1].slice(2..);
 
-            let mut buf_mut = buf.clone().into_mut();
+            let mut buf_mut = BitBufferMut::copy_from(&buf);
             assert_eq!(buf_mut.len(), 10);
 
             let tail = buf_mut.split_off(i);

vortex-buffer/src/bit/ops.rs

@@ -4,13 +4,54 @@
 use crate::BitBuffer;
 use crate::BitBufferMut;
 use crate::Buffer;
+use crate::ByteBufferMut;
 use crate::trusted_len::TrustedLenExt;
 
 #[inline]
-pub(super) fn bitwise_unary_op<F: FnMut(u64) -> u64>(buffer: &BitBuffer, op: F) -> BitBuffer {
-    let mut buf = buffer.clone().into_mut();
-    bitwise_unary_op_mut(&mut buf, op);
-    buf.freeze()
+pub(super) fn bitwise_unary_op<F: FnMut(u64) -> u64>(buffer: BitBuffer, mut op: F) -> BitBuffer {
+    match buffer.try_into_mut() {
+        Ok(mut buf) => {
+            bitwise_unary_op_mut(&mut buf, op);
+            buf.freeze()
+        }
+        Err(buffer) => {
+            let len = buffer.len();
+            let offset = buffer.offset();
+            let src = buffer.inner().as_slice();
+
+            let mut dst = ByteBufferMut::with_capacity(src.len());
+            let u64_len = src.len() / 8;
+            let remainder = src.len() % 8;
+
+            let mut src_ptr = src.as_ptr() as *const u64;
+            let mut dst_ptr = dst.spare_capacity_mut().as_mut_ptr() as *mut u64;
+            for _ in 0..u64_len {
+                let value = unsafe { src_ptr.read_unaligned() };
+                unsafe { dst_ptr.write_unaligned(op(value)) };
+                src_ptr = unsafe { src_ptr.add(1) };
+                dst_ptr = unsafe { dst_ptr.add(1) };
+            }
+
+            if remainder > 0 {
+                let mut remainder_u64 = 0u64;
+                let src_bytes = src_ptr as *const u8;
+                let dst_bytes = dst_ptr as *mut u8;
+                for i in 0..remainder {
+                    let byte = unsafe { src_bytes.add(i).read() };
+                    remainder_u64 |= (byte as u64) << (i * 8);
+                }
+                let remainder_u64 = op(remainder_u64);
+                for i in 0..remainder {
+                    let byte = ((remainder_u64 >> (i * 8)) & 0xFF) as u8;
+                    unsafe { dst_bytes.add(i).write(byte) };
+                }
+            }
+
+            // SAFETY: we wrote exactly src.len() bytes into the spare capacity.
+            unsafe { dst.set_len(src.len()) };
+            BitBuffer::new_with_offset(dst.freeze(), len, offset)
+        }
+    }
 }
 
 #[inline]
@@ -95,7 +136,7 @@ mod tests {
     #[test]
     fn test_bitwise_unary_not() {
         let buffer = BitBuffer::new(buffer![0b10101010u8], 4);
-        let result = bitwise_unary_op(&buffer, |x| !x);
+        let result = bitwise_unary_op(buffer, |x| !x);
         assert_eq!(result, bitbuffer![true, false, true, false]);
     }
 

vortex-mask/src/mask_mut.rs

@@ -498,10 +498,10 @@ impl Mask {
             Mask::AllFalse(len) => MaskMut::new_false(len),
             Mask::Values(values) => {
                 let bit_buffer_mut = match Arc::try_unwrap(values) {
-                    Ok(mask_values) => {
-                        let bit_buffer = mask_values.into_buffer();
-                        bit_buffer.into_mut()
-                    }
+                    Ok(mask_values) => mask_values
+                        .into_buffer()
+                        .try_into_mut()
+                        .unwrap_or_else(|bb| BitBufferMut::copy_from(&bb)),
                     Err(arc_mask_values) => {
                         let bit_buffer = arc_mask_values.bit_buffer();
                         BitBufferMut::copy_from(bit_buffer)