Update DeltaArray to always pad values to 1024 elements and use bit

transpose functions

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

Author: robert3005

encodings/fastlanes/src/bit_transpose/mod.rs

@@ -25,6 +25,10 @@ mod scalar;
 #[cfg(not(feature = "_test-harness"))]
 mod x86;
 
+mod validity;
+
+pub use validity::*;
+
 /// Base indices for the first 64 output bytes (lanes 0-7).
 /// Each entry indicates the starting input byte index for that output byte group.
 /// Pattern: [0*2, 4*2, 2*2, 6*2, 1*2, 5*2, 3*2, 7*2] = [0, 8, 4, 12, 2, 10, 6, 14]
@@ -39,6 +43,8 @@ const TRANSPOSE_2X2: u64 = 0x00AA_00AA_00AA_00AA;
 const TRANSPOSE_4X4: u64 = 0x0000_CCCC_0000_CCCC;
 const TRANSPOSE_8X8: u64 = 0x0000_0000_F0F0_F0F0;
 
+/// Transpose 1024-bits into FastLanes layout.
+///
 /// Dispatch to the best available implementation at runtime.
 #[inline]
 pub fn transpose_bits(input: &[u8; 128], output: &mut [u8; 128]) {
@@ -64,6 +70,8 @@ pub fn transpose_bits(input: &[u8; 128], output: &mut [u8; 128]) {
     scalar::transpose_bits_scalar(input, output);
 }
 
+/// Untranspose 1024-bits from FastLanes layout.
+///
 /// Dispatch untranspose to the best available implementation at runtime.
 #[inline]
 pub fn untranspose_bits(input: &[u8; 128], output: &mut [u8; 128]) {

encodings/fastlanes/src/bit_transpose/validity.rs

@@ -0,0 +1,97 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::mem;
+use std::mem::MaybeUninit;
+
+use vortex_array::Canonical;
+use vortex_array::ExecutionCtx;
+use vortex_array::IntoArray;
+use vortex_array::arrays::BoolArray;
+use vortex_array::validity::Validity;
+use vortex_buffer::BitBuffer;
+use vortex_buffer::ByteBufferMut;
+use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+
+use crate::bit_transpose::transpose_bits;
+use crate::bit_transpose::untranspose_bits;
+
+pub fn transpose_validity(validity: &Validity, ctx: &mut ExecutionCtx) -> VortexResult<Validity> {
+    match validity {
+        Validity::Array(mask) => {
+            let bools = mask
+                .clone()
+                .execute::<Canonical>(ctx)?
+                .into_bool()
+                .into_bit_buffer();
+
+            Ok(Validity::Array(
+                BoolArray::new(transpose_bitbuffer(bools), Validity::NonNullable).into_array(),
+            ))
+        }
+        v @ Validity::AllValid | v @ Validity::AllInvalid | v @ Validity::NonNullable => {
+            Ok(v.clone())
+        }
+    }
+}
+
+#[inline]
+pub fn transpose_bitbuffer(bits: BitBuffer) -> BitBuffer {
+    fastlanes_layout_apply(bits, transpose_bits)
+}
+
+pub fn untranspose_validity(validity: &Validity, ctx: &mut ExecutionCtx) -> VortexResult<Validity> {
+    match validity {
+        Validity::Array(mask) => {
+            let bools = mask
+                .clone()
+                .execute::<Canonical>(ctx)?
+                .into_bool()
+                .into_bit_buffer();
+
+            Ok(Validity::Array(
+                BoolArray::new(untranspose_bitbuffer(bools), Validity::NonNullable).into_array(),
+            ))
+        }
+        v @ Validity::AllValid | v @ Validity::AllInvalid | v @ Validity::NonNullable => {
+            Ok(v.clone())
+        }
+    }
+}
+
+#[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 {
+    let (offset, len, bytes) = bits.into_inner();
+
+    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>();
+    // We can ignore the spare trailer capacity that can be an artifact of allocator as we requested 128 multiple chunks
+    let (output_chunks, _) = output.spare_capacity_mut().as_chunks_mut::<128>();
+
+    for (input, output) in input_chunks.iter().zip(output_chunks.iter_mut()) {
+        op(input, unsafe {
+            mem::transmute::<&mut [MaybeUninit<u8>; 128], &mut [u8; 128]>(output)
+        });
+    }
+
+    if !input_trailer.is_empty() {
+        let mut padded_input = [0u8; 128];
+        padded_input[0..input_trailer.len()].clone_from_slice(input_trailer);
+        op(&padded_input, unsafe {
+            mem::transmute::<&mut [MaybeUninit<u8>; 128], &mut [u8; 128]>(
+                output_chunks
+                    .last_mut()
+                    .vortex_expect("Output wasn't a multiple of 128 bytes"),
+            )
+        });
+    }
+
+    unsafe { output.set_len(output_len) };
+    BitBuffer::new_with_offset(output.freeze().into_byte_buffer(), len, offset)
+}

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

@@ -1,91 +1,128 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use arrayref::array_mut_ref;
-use arrayref::array_ref;
+use std::mem;
+use std::mem::MaybeUninit;
+
 use fastlanes::Delta;
 use fastlanes::FastLanes;
 use fastlanes::Transpose;
-use num_traits::WrappingSub;
+use vortex_array::Canonical;
+use vortex_array::ExecutionCtx;
+use vortex_array::IntoArray;
+use vortex_array::arrays::BoolArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::dtype::NativePType;
 use vortex_array::match_each_unsigned_integer_ptype;
+use vortex_array::validity::Validity;
 use vortex_array::vtable::ValidityHelper;
+use vortex_buffer::BitBuffer;
 use vortex_buffer::Buffer;
 use vortex_buffer::BufferMut;
 use vortex_error::VortexResult;
 
-pub fn delta_compress(array: &PrimitiveArray) -> VortexResult<(PrimitiveArray, PrimitiveArray)> {
-    // TODO(ngates): fill forward nulls?
-    // let filled = fill_forward(array)?.to_primitive()?;
+use crate::bit_transpose::transpose_bitbuffer;
 
-    // Compress the filled array
+pub fn delta_compress(
+    array: &PrimitiveArray,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<(PrimitiveArray, PrimitiveArray)> {
     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)?;
         (
-            // To preserve nullability, we include Validity
             PrimitiveArray::new(bases, array.dtype().nullability().into()),
-            PrimitiveArray::new(deltas, array.validity().clone()),
+            PrimitiveArray::new(deltas, validity),
         )
     });
 
     Ok((bases, deltas))
 }
 
-fn compress_primitive<T: NativePType + Delta + Transpose + WrappingSub, const LANES: usize>(
+/// 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>) {
-    // How many fastlanes vectors we will process.
-    let num_chunks = array.len() / 1024;
+    let padded_len = array.len().next_multiple_of(1024);
+    let num_chunks = padded_len / 1024;
+    let bases_len = num_chunks * LANES;
+
+    // Split into full 1024-element chunks and a remainder.
+    let (full_chunks, remainder) = array.as_chunks::      <1024>();
 
     // Allocate result arrays.
-    let mut bases = BufferMut::with_capacity(num_chunks * T::LANES + 1);
-    let mut deltas = BufferMut::with_capacity(array.len());
-
-    // Loop over all the 1024-element chunks.
-    if num_chunks > 0 {
-        let mut transposed: [T; 1024] = [T::default(); 1024];
-
-        for i in 0..num_chunks {
-            let start_elem = i * 1024;
-            let chunk: &[T; 1024] = array_ref![array, start_elem, 1024];
-            Transpose::transpose(chunk, &mut transposed);
-
-            // Initialize and store the base vector for each chunk
-            bases.extend_from_slice(&transposed[0..T::LANES]);
-
-            deltas.reserve(1024);
-            let delta_len = deltas.len();
-            unsafe {
-                deltas.set_len(delta_len + 1024);
-                Delta::delta::<LANES>(
-                    &transposed,
-                    &*(transposed[0..T::LANES].as_ptr().cast()),
-                    array_mut_ref![deltas[delta_len..], 0, 1024],
-                );
-            }
+    let mut bases = BufferMut::with_capacity(bases_len);
+    let mut deltas = BufferMut::with_capacity(padded_len);
+    let (output_deltas, _) = deltas.spare_capacity_mut().as_chunks_mut::<1024>();
+
+    // Loop over all full 1024-element chunks.
+    let mut transposed: [T; 1024] = [T::default(); 1024];
+    for (chunk, output) in full_chunks.iter().zip(output_deltas.iter_mut()) {
+        Transpose::transpose(chunk, &mut transposed);
+        bases.extend_from_slice(&transposed[0..T::LANES]);
+
+        unsafe {
+            Delta::delta::<LANES>(
+                &transposed,
+                &*(transposed[0..T::LANES].as_ptr().cast()),
+                mem::transmute::<&mut [MaybeUninit<T>; 1024], &mut [T; 1024]>(output),
+            );
         }
     }
 
-    // To avoid padding, the remainder is encoded with scalar logic.
-    let remainder_size = array.len() % 1024;
-    if remainder_size > 0 {
-        let chunk = &array[array.len() - remainder_size..];
-        let mut base_scalar = chunk[0];
-        bases.push(base_scalar);
-        for next in chunk {
-            let diff = next.wrapping_sub(&base_scalar);
-            deltas.push(diff);
-            base_scalar = *next;
+    // Pad the remainder to 1024 elements and process as a full chunk.
+    if !remainder.is_empty() {
+        let mut padded_chunk = [T::default(); 1024];
+        padded_chunk[..remainder.len()].copy_from_slice(remainder);
+
+        Transpose::transpose(&padded_chunk, &mut transposed);
+        bases.extend_from_slice(&transposed[0..T::LANES]);
+
+        unsafe {
+            Delta::delta::<LANES>(
+                &transposed,
+                &*(transposed[0..T::LANES].as_ptr().cast()),
+                mem::transmute::<&mut [MaybeUninit<T>; 1024], &mut [T; 1024]>(
+                    &mut output_deltas[full_chunks.len()],
+                ),
+            );
         }
     }
 
-    assert_eq!(
-        bases.len(),
-        num_chunks * T::LANES + (if remainder_size > 0 { 1 } else { 0 })
-    );
-    assert_eq!(deltas.len(), array.len());
+    unsafe { deltas.set_len(padded_len) };
+
+    assert_eq!(bases.len(), bases_len);
+    assert_eq!(deltas.len(), padded_len);
 
     (bases.freeze(), deltas.freeze())
 }
@@ -94,6 +131,7 @@ fn compress_primitive<T: NativePType + Delta + Transpose + WrappingSub, const LA
 mod tests {
     use std::sync::LazyLock;
 
+    use rstest::rstest;
     use vortex_array::VortexSessionExecute;
     use vortex_array::arrays::PrimitiveArray;
     use vortex_array::assert_arrays_eq;
@@ -107,28 +145,18 @@ mod tests {
     static SESSION: LazyLock<VortexSession> =
         LazyLock::new(|| VortexSession::empty().with::<ArraySession>());
 
-    #[test]
-    fn test_compress() -> VortexResult<()> {
-        do_roundtrip_test((0u32..10_000).collect())
-    }
-
-    #[test]
-    fn test_compress_nullable() -> VortexResult<()> {
-        do_roundtrip_test(PrimitiveArray::from_option_iter(
+    #[rstest]
+    #[case((0u32..10_000).collect())]
+    #[case((0..10_000).map(|i| (i % (u8::MAX as i32)) as u8).collect())]
+    #[case(PrimitiveArray::from_option_iter(
             (0u32..10_000).map(|i| (i % 2 == 0).then_some(i)),
-        ))
-    }
-
-    #[test]
-    fn test_compress_overflow() -> VortexResult<()> {
-        do_roundtrip_test((0..10_000).map(|i| (i % (u8::MAX as i32)) as u8).collect())
-    }
-
-    fn do_roundtrip_test(input: PrimitiveArray) -> VortexResult<()> {
-        let delta = DeltaArray::try_from_primitive_array(&input)?;
-        assert_eq!(delta.len(), input.len());
+    ))]
+    fn test_compress(#[case] array: PrimitiveArray) -> VortexResult<()> {
+        let delta =
+            DeltaArray::try_from_primitive_array(&array, &mut SESSION.create_execution_ctx())?;
+        assert_eq!(delta.len(), array.len());
         let decompressed = delta_decompress(&delta, &mut SESSION.create_execution_ctx())?;
-        assert_arrays_eq!(decompressed, input);
+        assert_arrays_eq!(decompressed, array);
         Ok(())
     }
 }