inplace

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

Author: robert3005

encodings/fastlanes/src/bit_transpose/validity.rs

@@ -10,6 +10,7 @@ use vortex_array::IntoArray;
 use vortex_array::arrays::BoolArray;
 use vortex_array::validity::Validity;
 use vortex_buffer::BitBuffer;
+use vortex_buffer::ByteBuffer;
 use vortex_buffer::ByteBufferMut;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
@@ -38,7 +39,25 @@ pub fn transpose_validity(validity: &Validity, ctx: &mut ExecutionCtx) -> Vortex
 
 #[inline]
 pub fn transpose_bitbuffer(bits: BitBuffer) -> BitBuffer {
-    fastlanes_layout_apply(bits, transpose_bits)
+    let (offset, len, bytes) = bits.into_inner();
+
+    if bytes.len().is_multiple_of(128) {
+        match bytes.try_into_mut() {
+            Ok(mut bytes_mut) => {
+                // We can ignore the spare trailer capacity that can be an artifact of allocator as we requested 128 multiple chunks
+                let (chunks, _) = bytes_mut.as_chunks_mut::<128>();
+                let mut tmp = [0u8; 128];
+                for chunk in chunks {
+                    transpose_bits(chunk, &mut tmp);
+                    chunk.copy_from_slice(&tmp);
+                }
+                BitBuffer::new_with_offset(bytes_mut.freeze().into_byte_buffer(), len, offset)
+            }
+            Err(bytes) => bits_op_with_copy(bytes, len, offset, transpose_bits),
+        }
+    } else {
+        bits_op_with_copy(bytes, len, offset, transpose_bits)
+    }
 }
 
 pub fn untranspose_validity(validity: &Validity, ctx: &mut ExecutionCtx) -> VortexResult<Validity> {
@@ -62,12 +81,31 @@ pub fn untranspose_validity(validity: &Validity, ctx: &mut ExecutionCtx) -> Vort
 
 #[inline]
 pub fn untranspose_bitbuffer(bits: BitBuffer) -> BitBuffer {
-    fastlanes_layout_apply(bits, untranspose_bits)
-}
-
-fn fastlanes_layout_apply<F: Fn(&[u8; 128], &mut [u8; 128])>(bits: BitBuffer, op: F) -> BitBuffer {
+    assert!(
+        bits.inner().len().is_multiple_of(128),
+        "Transpose BitBuffer must be 128-byte aligned"
+    );
     let (offset, len, bytes) = bits.into_inner();
+    match bytes.try_into_mut() {
+        Ok(mut bytes_mut) => {
+            let (chunks, _) = bytes_mut.as_chunks_mut::<128>();
+            let mut tmp = [0u8; 128];
+            for chunk in chunks {
+                untranspose_bits(chunk, &mut tmp);
+                chunk.copy_from_slice(&tmp);
+            }
+            BitBuffer::new_with_offset(bytes_mut.freeze().into_byte_buffer(), len, offset)
+        }
+        Err(bytes) => bits_op_with_copy(bytes, len, offset, untranspose_bits),
+    }
+}
 
+fn bits_op_with_copy<F: Fn(&[u8; 128], &mut [u8; 128])>(
+    bytes: ByteBuffer,
+    len: usize,
+    offset: usize,
+    op: F,
+) -> BitBuffer {
     let output_len = bytes.len().next_multiple_of(128);
     let mut output = ByteBufferMut::with_capacity(output_len);
     let (input_chunks, input_trailer) = bytes.as_chunks::<128>();

encodings/fastlanes/src/delta/array/delta_compress.rs

@@ -21,7 +21,7 @@ use vortex_buffer::Buffer;
 use vortex_buffer::BufferMut;
 use vortex_error::VortexResult;
 
-use crate::bit_transpose::transpose_bitbuffer;
+use crate::bit_transpose::{transpose_bitbuffer, transpose_validity};
 
 pub fn delta_compress(
     array: &PrimitiveArray,
@@ -30,7 +30,7 @@ pub fn delta_compress(
     let (bases, deltas) = match_each_unsigned_integer_ptype!(array.ptype(), |T| {
         const LANES: usize = T::LANES;
         let (bases, deltas) = compress_primitive::<T, LANES>(array.as_slice::<T>());
-        let validity = transpose_and_pad_validity(array.validity(), deltas.len(), ctx)?;
+        let validity = transpose_validity(array.validity(), ctx)?;
         (
             PrimitiveArray::new(bases, array.dtype().nullability().into()),
             PrimitiveArray::new(deltas, validity),
@@ -40,36 +40,6 @@ pub fn delta_compress(
     Ok((bases, deltas))
 }
 
-/// Transpose and pad validity to match the padded deltas length.
-///
-/// For [`Validity::Array`], the validity bits are transposed into FastLanes order and then
-/// extended to `padded_len`. The underlying byte buffer from transposition is already
-/// padded to 128-byte alignment (1024 bits), which exactly matches our 1024-element chunks.
-fn transpose_and_pad_validity(
-    validity: &Validity,
-    padded_len: usize,
-    ctx: &mut ExecutionCtx,
-) -> VortexResult<Validity> {
-    match validity {
-        Validity::Array(mask) => {
-            let bools = mask
-                .clone()
-                .execute::<Canonical>(ctx)?
-                .into_bool()
-                .into_bit_buffer();
-            let transposed = transpose_bitbuffer(bools);
-            let (offset, _len, bytes) = transposed.into_inner();
-            let padded = BitBuffer::new_with_offset(bytes, padded_len, offset);
-            Ok(Validity::Array(
-                BoolArray::new(padded, Validity::NonNullable).into_array(),
-            ))
-        }
-        v @ Validity::AllValid | v @ Validity::AllInvalid | v @ Validity::NonNullable => {
-            Ok(v.clone())
-        }
-    }
-}
-
 fn compress_primitive<T: NativePType + Delta + Transpose, const LANES: usize>(
     array: &[T],
 ) -> (Buffer<T>, Buffer<T>) {