ValidityKernel wrapper

Signed-off-by: Andrew Duffy <[email protected]>

Author: a10y

encodings/fastlanes/src/bitpacking/array/bitpack_pipeline.rs

@@ -3,13 +3,14 @@
 
 use fastlanes::{BitPacking, FastLanes};
 use static_assertions::const_assert_eq;
+use vortex_array::pipeline::validity::ValidityKernel;
 use vortex_array::pipeline::{
     BindContext, BitView, Kernel, KernelCtx, N, PipelineInputs, PipelinedNode,
 };
+use vortex_array::validity::Validity;
 use vortex_buffer::Buffer;
 use vortex_dtype::{PTypeDowncastExt, PhysicalPType, match_each_integer_ptype};
 use vortex_error::VortexResult;
-use vortex_mask::Mask;
 use vortex_vector::primitive::PVectorMut;
 use vortex_vector::{VectorMut, VectorMutOps};
 
@@ -53,11 +54,21 @@ impl PipelinedNode for BitPackedArray {
                 )
             }
 
-            Ok(Box::new(AlignedBitPackedKernel::<T>::new(
-                packed_bit_width,
-                packed_buffer,
-                self.validity.to_mask(self.len()),
-            )) as Box<dyn Kernel>)
+            match self.validity {
+                Validity::NonNullable | Validity::AllValid => Ok(Box::new(
+                    AlignedBitPackedKernel::<T>::new(packed_bit_width, packed_buffer),
+                )
+                    as Box<dyn Kernel>),
+                Validity::AllInvalid => {
+                    todo!("make a kernel that returns constant null");
+                }
+                Validity::Array(_) => {
+                    let inner = AlignedBitPackedKernel::<T>::new(packed_bit_width, packed_buffer);
+                    let mask = self.validity_mask();
+
+                    Ok(Box::new(ValidityKernel::new(inner, mask)) as Box<dyn Kernel>)
+                }
+            }
         })
     }
 }
@@ -84,19 +95,12 @@ pub struct AlignedBitPackedKernel<BP: PhysicalPType<Physical: BitPacking>> {
     /// The buffer containing the bitpacked values.
     packed_buffer: Buffer<BP::Physical>,
 
-    /// The validity mask for the bitpacked array.
-    validity: Mask,
-
     /// The total number of bitpacked chunks we have unpacked.
     num_chunks_unpacked: usize,
 }
 
 impl<BP: PhysicalPType<Physical: BitPacking>> AlignedBitPackedKernel<BP> {
-    pub fn new(
-        packed_bit_width: usize,
-        packed_buffer: Buffer<BP::Physical>,
-        validity: Mask,
-    ) -> Self {
+    pub fn new(packed_bit_width: usize, packed_buffer: Buffer<BP::Physical>) -> Self {
         let packed_stride =
             packed_bit_width * <<BP as PhysicalPType>::Physical as FastLanes>::LANES;
 
@@ -110,7 +114,6 @@ impl<BP: PhysicalPType<Physical: BitPacking>> AlignedBitPackedKernel<BP> {
             packed_bit_width,
             packed_stride,
             packed_buffer,
-            validity,
             num_chunks_unpacked: 0,
         }
     }
@@ -130,9 +133,6 @@ impl<BP: PhysicalPType<Physical: BitPacking>> Kernel for AlignedBitPackedKernel<
         let not_yet_unpacked_values = &self.packed_buffer.as_slice()[packed_offset..];
 
         let true_count = selection.true_count();
-        let chunk_offset = self.num_chunks_unpacked * N;
-        let array_len = self.validity.len();
-        debug_assert!(chunk_offset < array_len);
 
         // If the true count is very small (the selection is sparse), we can unpack individual
         // elements directly into the output vector.
@@ -141,26 +141,21 @@ impl<BP: PhysicalPType<Physical: BitPacking>> Kernel for AlignedBitPackedKernel<
             debug_assert!(true_count <= output_vector.capacity());
 
             selection.iter_ones(|idx| {
-                let absolute_idx = chunk_offset + idx;
-                if self.validity.value(absolute_idx) {
-                    // SAFETY:
-                    // - The documentation for `packed_bit_width` explains that the size is valid.
-                    // - We know that the size of the `next_packed_chunk` we provide is equal to
-                    //   `self.packed_stride`, and we explain why this is correct in its
-                    //   documentation.
-                    let unpacked_value = unsafe {
-                        BitPacking::unchecked_unpack_single(
-                            self.packed_bit_width,
-                            not_yet_unpacked_values,
-                            idx,
-                        )
-                    };
-
-                    // SAFETY: We just reserved enough capacity to push these values.
-                    unsafe { output_vector.push_unchecked(unpacked_value) };
-                } else {
-                    output_vector.append_nulls(1);
-                }
+                // SAFETY:
+                // - The documentation for `packed_bit_width` explains that the size is valid.
+                // - We know that the size of the `next_packed_chunk` we provide is equal to
+                //   `self.packed_stride`, and we explain why this is correct in its
+                //   documentation.
+                let unpacked_value = unsafe {
+                    BitPacking::unchecked_unpack_single(
+                        self.packed_bit_width,
+                        not_yet_unpacked_values,
+                        idx,
+                    )
+                };
+
+                // SAFETY: We just reserved enough capacity to push these values.
+                unsafe { output_vector.elements_mut().push_unchecked(unpacked_value) };
             });
         } else {
             // Otherwise if the mask is dense, it is faster to fully unpack the entire 1024
@@ -192,26 +187,6 @@ impl<BP: PhysicalPType<Physical: BitPacking>> Kernel for AlignedBitPackedKernel<
                     output_vector.as_mut(),
                 );
             }
-
-            if array_len < chunk_offset + N {
-                let vector_len = array_len - chunk_offset;
-                debug_assert!(vector_len < N, "math is broken");
-
-                // SAFETY: This must be less than `N` so this is just a truncate.
-                unsafe { output_vector.elements_mut().set_len(vector_len) };
-
-                let chunk_mask = self.validity.slice(chunk_offset..array_len);
-
-                // SAFETY: We have just set the elements length to N, and the validity buffer has
-                // capacity for N elements.
-                unsafe { output_vector.validity_mut() }.append_mask(&chunk_mask);
-            } else {
-                let chunk_mask = self.validity.slice(chunk_offset..chunk_offset + N);
-
-                // SAFETY: We have just set the elements length to N, and the validity buffer has
-                // capacity for N elements.
-                unsafe { output_vector.validity_mut() }.append_mask(&chunk_mask);
-            }
         }
 
         self.num_chunks_unpacked += 1;

vortex-array/src/pipeline/mod.rs

@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 pub mod driver;
+pub mod validity;
 
 use vortex_error::{VortexExpect, VortexResult};
 use vortex_vector::{Vector, VectorMut};

vortex-array/src/pipeline/validity.rs

@@ -0,0 +1,52 @@
+use vortex_error::VortexResult;
+use vortex_mask::Mask;
+use vortex_vector::{VectorMut, VectorMutOps};
+
+use crate::pipeline::{BitView, Kernel, KernelCtx};
+
+/// `ValidityKernel` wraps a child kernel, passing a validity mask through to the output.
+pub struct ValidityKernel<K> {
+    child: K,
+    validity: Mask,
+    position: usize,
+}
+
+impl<K> ValidityKernel<K> {
+    pub fn new(inner: K, validity: Mask) -> Self {
+        Self {
+            child: inner,
+            validity,
+            position: 0,
+        }
+    }
+}
+
+impl<K: Kernel> Kernel for ValidityKernel<K> {
+    fn step(
+        &mut self,
+        ctx: &KernelCtx,
+        selection: &BitView,
+        out: &mut VectorMut,
+    ) -> VortexResult<()> {
+        // execute the child kernel
+        self.child.step(ctx, selection, out)?;
+
+        debug_assert_eq!(
+            out.validity().len(),
+            self.position,
+            "child kernel should not step validity when wrapped with ValidityKernel"
+        );
+
+        let new_position = self.position + out.len();
+
+        let slice = self.validity.slice(self.position..new_position);
+
+        // SAFETY: the child kernel must extend elements in its step function.
+        unsafe { out.validity_mut().append_mask(&slice) };
+
+        // Advance the position in the kernel here.
+        self.position = new_position;
+
+        Ok(())
+    }
+}

vortex-compute/src/filter/vector/binaryview.rs

@@ -3,10 +3,10 @@
 
 use vortex_buffer::{Buffer, BufferMut};
 use vortex_mask::{Mask, MaskMut};
-use vortex_vector::VectorOps;
 use vortex_vector::binaryview::{
     BinaryView, BinaryViewType, BinaryViewVector, BinaryViewVectorMut,
 };
+use vortex_vector::{VectorMutOps, VectorOps};
 
 use crate::filter::Filter;
 

vortex-compute/src/filter/vector/bool.rs

@@ -3,8 +3,8 @@
 
 use vortex_buffer::BitBuffer;
 use vortex_mask::{Mask, MaskMut};
-use vortex_vector::VectorOps;
 use vortex_vector::bool::{BoolVector, BoolVectorMut};
+use vortex_vector::{VectorMutOps, VectorOps};
 
 use crate::filter::Filter;
 

vortex-compute/src/filter/vector/dvector.rs

@@ -4,8 +4,8 @@
 use vortex_buffer::{Buffer, BufferMut};
 use vortex_dtype::NativeDecimalType;
 use vortex_mask::{Mask, MaskMut};
-use vortex_vector::VectorOps;
 use vortex_vector::decimal::{DVector, DVectorMut};
+use vortex_vector::{VectorMutOps, VectorOps};
 
 use crate::filter::Filter;
 

vortex-compute/src/filter/vector/list.rs

@@ -4,9 +4,9 @@
 use std::sync::Arc;
 
 use vortex_mask::{Mask, MaskMut};
-use vortex_vector::VectorOps;
 use vortex_vector::listview::{ListViewVector, ListViewVectorMut};
 use vortex_vector::primitive::{PrimitiveVector, PrimitiveVectorMut};
+use vortex_vector::{VectorMutOps, VectorOps};
 
 use crate::filter::Filter;
 

vortex-compute/src/filter/vector/struct_.rs

@@ -5,7 +5,7 @@ use std::sync::Arc;
 
 use vortex_mask::{Mask, MaskMut};
 use vortex_vector::struct_::{StructVector, StructVectorMut};
-use vortex_vector::{Vector, VectorMut, VectorOps};
+use vortex_vector::{Vector, VectorMut, VectorMutOps, VectorOps};
 
 use crate::filter::Filter;
 

vortex-vector/src/binaryview/vector_mut.rs

@@ -122,16 +122,6 @@ impl<T: BinaryViewType> BinaryViewVectorMut<T> {
         &mut self.views
     }
 
-    /// Get a mutable handle to the validity mask of the vector.
-    ///
-    /// # Safety
-    ///
-    /// Caller must make sure that the length of the validity mask
-    /// always matches the length of the views
-    pub unsafe fn validity_mut(&mut self) -> &mut MaskMut {
-        &mut self.validity
-    }
-
     /// Get a mutable handle to the vector of buffers backing the string data of the vector.
     pub fn buffers(&mut self) -> &mut Vec<ByteBuffer> {
         &mut self.buffers
@@ -216,6 +206,10 @@ impl<T: BinaryViewType> VectorMutOps for BinaryViewVectorMut<T> {
         &self.validity
     }
 
+    unsafe fn validity_mut(&mut self) -> &mut MaskMut {
+        &mut self.validity
+    }
+
     fn capacity(&self) -> usize {
         self.views.capacity()
     }

vortex-vector/src/bool/vector_mut.rs

@@ -93,16 +93,6 @@ impl BoolVectorMut {
     pub unsafe fn bits_mut(&mut self) -> &mut BitBufferMut {
         &mut self.bits
     }
-
-    /// Get a mutable handle to the validity mask of the vector.
-    ///
-    /// # Safety
-    ///
-    /// Caller must ensure that length of the validity always matches
-    /// length of the bits.
-    pub unsafe fn validity_mut(&mut self) -> &mut MaskMut {
-        &mut self.validity
-    }
 }
 
 impl VectorMutOps for BoolVectorMut {
@@ -118,6 +108,10 @@ impl VectorMutOps for BoolVectorMut {
         &self.validity
     }
 
+    unsafe fn validity_mut(&mut self) -> &mut MaskMut {
+        &mut self.validity
+    }
+
     fn capacity(&self) -> usize {
         self.bits.capacity()
     }