Update DeltaArray to always pad values to 1024 elements and use bit transpose functions (#7007)

This pr aims to fix the long standing issue with pushing down computation over
DecimalArrays where their validity wasn't matching the value layout. 
Now the validity is also transposed. At the same time we avoid scalar encoding
and pad values to 1024 elements

fix #4973 

---------

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

Author: robert3005

encodings/fastlanes/public-api.lock

@@ -2,10 +2,18 @@ pub mod vortex_fastlanes
 
 pub mod vortex_fastlanes::bit_transpose
 
+pub fn vortex_fastlanes::bit_transpose::transpose_bitbuffer(bits: vortex_buffer::bit::buf::BitBuffer) -> vortex_buffer::bit::buf::BitBuffer
+
 pub fn vortex_fastlanes::bit_transpose::transpose_bits(input: &[u8; 128], output: &mut [u8; 128])
 
+pub fn vortex_fastlanes::bit_transpose::transpose_validity(validity: &vortex_array::validity::Validity, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::validity::Validity>
+
+pub fn vortex_fastlanes::bit_transpose::untranspose_bitbuffer(bits: vortex_buffer::bit::buf::BitBuffer) -> vortex_buffer::bit::buf::BitBuffer
+
 pub fn vortex_fastlanes::bit_transpose::untranspose_bits(input: &[u8; 128], output: &mut [u8; 128])
 
+pub fn vortex_fastlanes::bit_transpose::untranspose_validity(validity: &vortex_array::validity::Validity, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<vortex_array::validity::Validity>
+
 pub mod vortex_fastlanes::bitpack_compress
 
 pub fn vortex_fastlanes::bitpack_compress::bit_width_histogram(array: &vortex_array::arrays::primitive::array::PrimitiveArray) -> vortex_error::VortexResult<alloc::vec::Vec<usize>>
@@ -296,7 +304,7 @@ pub type vortex_fastlanes::Delta::Metadata = vortex_array::metadata::ProstMetada
 
 pub type vortex_fastlanes::Delta::OperationsVTable = vortex_fastlanes::Delta
 
-pub type vortex_fastlanes::Delta::ValidityVTable = vortex_array::vtable::validity::ValidityVTableFromChildSliceHelper
+pub type vortex_fastlanes::Delta::ValidityVTable = vortex_fastlanes::Delta
 
 pub fn vortex_fastlanes::Delta::array_eq(array: &vortex_fastlanes::DeltaArray, other: &vortex_fastlanes::DeltaArray, precision: vortex_array::hash::Precision) -> bool
 
@@ -340,6 +348,10 @@ impl vortex_array::vtable::operations::OperationsVTable<vortex_fastlanes::Delta>
 
 pub fn vortex_fastlanes::Delta::scalar_at(array: &vortex_fastlanes::DeltaArray, index: usize) -> vortex_error::VortexResult<vortex_array::scalar::Scalar>
 
+impl vortex_array::vtable::validity::ValidityVTable<vortex_fastlanes::Delta> for vortex_fastlanes::Delta
+
+pub fn vortex_fastlanes::Delta::validity(array: &vortex_fastlanes::DeltaArray) -> vortex_error::VortexResult<vortex_array::validity::Validity>
+
 pub struct vortex_fastlanes::DeltaArray
 
 impl vortex_fastlanes::DeltaArray
@@ -358,9 +370,7 @@ pub fn vortex_fastlanes::DeltaArray::offset(&self) -> usize
 
 pub fn vortex_fastlanes::DeltaArray::try_from_delta_compress_parts(bases: vortex_array::array::ArrayRef, deltas: vortex_array::array::ArrayRef) -> vortex_error::VortexResult<Self>
 
-pub fn vortex_fastlanes::DeltaArray::try_from_primitive_array(array: &vortex_array::arrays::primitive::array::PrimitiveArray) -> vortex_error::VortexResult<Self>
-
-pub fn vortex_fastlanes::DeltaArray::try_from_vec<T: vortex_array::dtype::ptype::NativePType>(vec: alloc::vec::Vec<T>) -> vortex_error::VortexResult<Self>
+pub fn vortex_fastlanes::DeltaArray::try_from_primitive_array(array: &vortex_array::arrays::primitive::array::PrimitiveArray, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<Self>
 
 pub fn vortex_fastlanes::DeltaArray::try_new(bases: vortex_array::array::ArrayRef, deltas: vortex_array::array::ArrayRef, offset: usize, logical_len: usize) -> vortex_error::VortexResult<Self>
 
@@ -394,10 +404,6 @@ impl vortex_array::array::IntoArray for vortex_fastlanes::DeltaArray
 
 pub fn vortex_fastlanes::DeltaArray::into_array(self) -> vortex_array::array::ArrayRef
 
-impl vortex_array::vtable::validity::ValidityChildSliceHelper for vortex_fastlanes::DeltaArray
-
-pub fn vortex_fastlanes::DeltaArray::unsliced_child_and_slice(&self) -> (&vortex_array::array::ArrayRef, usize, usize)
-
 pub struct vortex_fastlanes::FoR
 
 impl vortex_fastlanes::FoR
@@ -668,6 +674,6 @@ impl vortex_array::vtable::validity::ValidityChildSliceHelper for vortex_fastlan
 
 pub fn vortex_fastlanes::RLEArray::unsliced_child_and_slice(&self) -> (&vortex_array::array::ArrayRef, usize, usize)
 
-pub fn vortex_fastlanes::delta_compress(array: &vortex_array::arrays::primitive::array::PrimitiveArray) -> vortex_error::VortexResult<(vortex_array::arrays::primitive::array::PrimitiveArray, vortex_array::arrays::primitive::array::PrimitiveArray)>
+pub fn vortex_fastlanes::delta_compress(array: &vortex_array::arrays::primitive::array::PrimitiveArray, ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<(vortex_array::arrays::primitive::array::PrimitiveArray, vortex_array::arrays::primitive::array::PrimitiveArray)>
 
 pub fn vortex_fastlanes::initialize(session: &mut vortex_session::VortexSession)

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,139 @@
+// 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::ByteBuffer;
+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 {
+    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> {
+    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 {
+    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>();
+    // 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.next_multiple_of(1024),
+        offset,
+    )
+}