handle validity in bitpack pipeline kernel

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

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

@@ -8,6 +8,7 @@ use vortex_array::pipeline::{BindContext, Kernel, KernelCtx, N, PipelineInputs,
 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};
 
@@ -54,6 +55,7 @@ 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>)
         })
     }
@@ -81,12 +83,19 @@ 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>) -> Self {
+    pub fn new(
+        packed_bit_width: usize,
+        packed_buffer: Buffer<BP::Physical>,
+        validity: Mask,
+    ) -> Self {
         let packed_stride =
             packed_bit_width * <<BP as PhysicalPType>::Physical as FastLanes>::LANES;
 
@@ -100,6 +109,7 @@ impl<BP: PhysicalPType<Physical: BitPacking>> AlignedBitPackedKernel<BP> {
             packed_bit_width,
             packed_stride,
             packed_buffer,
+            validity,
             num_chunks_unpacked: 0,
         }
     }
@@ -119,6 +129,9 @@ 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.
@@ -127,20 +140,26 @@ impl<BP: PhysicalPType<Physical: BitPacking>> Kernel for AlignedBitPackedKernel<
             debug_assert!(true_count <= output_vector.capacity());
 
             selection.iter_ones(|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) };
+                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);
+                }
             });
         } else {
             // Otherwise if the mask is dense, it is faster to fully unpack the entire 1024
@@ -149,16 +168,17 @@ impl<BP: PhysicalPType<Physical: BitPacking>> Kernel for AlignedBitPackedKernel<
             output_vector.reserve(N);
             debug_assert!(N <= output_vector.capacity());
 
-            // SAFETY: We have just reserved enough capacity for the vector to set the length, and
-            // we also are about to initialize all of the values **without** reading the memory.
-            unsafe { output_vector.set_len(N) };
-
             let next_packed_chunk = &not_yet_unpacked_values[..self.packed_stride];
             debug_assert_eq!(
                 next_packed_chunk.len(),
                 FL_VECTOR_SIZE * self.packed_bit_width / BP::Physical::T
             );
 
+            // SAFETY: We have just reserved enough capacity for the elements buffer to set the
+            // length, and we are about to initialize all of the values **without** reading the
+            // memory.
+            unsafe { output_vector.elements_mut().set_len(N) };
+
             // 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
@@ -168,9 +188,29 @@ impl<BP: PhysicalPType<Physical: BitPacking>> Kernel for AlignedBitPackedKernel<
                 BitPacking::unchecked_unpack(
                     self.packed_bit_width,
                     next_packed_chunk,
-                    &mut output_vector.as_mut()[..FL_VECTOR_SIZE],
+                    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;
@@ -285,4 +325,104 @@ mod tests {
         assert_eq!(elements[3], 512);
         assert_eq!(elements[4], 1000);
     }
+
+    #[test]
+    fn test_bitpack_pipeline_sparse_with_nulls() {
+        // Create 1024 elements with some nulls.
+        let values: Vec<Option<u32>> = (0..1024)
+            .map(|i| if i % 100 == 0 { None } else { Some(i as u32) })
+            .collect();
+        let primitive = PrimitiveArray::from_option_iter(values).to_array();
+
+        // Encode with 10-bit width.
+        let bitpacked = BitPackedArray::encode(&primitive, 10).unwrap();
+        assert_eq!(bitpacked.bit_width(), 10, "Bit width should be 10");
+
+        // Select only 5 elements at specific indices, including a null value at index 100.
+        let indices = vec![10, 100, 256, 512, 1000];
+        let mask = Mask::from_indices(1024, indices);
+
+        // This should use the sparse path since true_count < 7.
+        let result = bitpacked.to_array().execute_with_selection(&mask).unwrap();
+        assert_eq!(result.len(), 5, "Result should have 5 elements");
+
+        let pvector_u32 = result.as_primitive().into_u32();
+        let elements = pvector_u32.elements().as_slice();
+
+        // Verify the values and validity.
+        assert_eq!(elements[0], 10);
+        assert!(
+            pvector_u32.validity().value(0),
+            "Element at index 0 should be valid"
+        );
+
+        // Index 100 should be null.
+        assert!(
+            !pvector_u32.validity().value(1),
+            "Element at index 1 (original index 100) should be null"
+        );
+
+        assert_eq!(elements[2], 256);
+        assert!(
+            pvector_u32.validity().value(2),
+            "Element at index 2 should be valid"
+        );
+
+        assert_eq!(elements[3], 512);
+        assert!(
+            pvector_u32.validity().value(3),
+            "Element at index 3 should be valid"
+        );
+
+        // Index 1000 should be null.
+        assert!(
+            !pvector_u32.validity().value(4),
+            "Element at index 4 (original index 1000) should be null"
+        );
+    }
+
+    #[test]
+    fn test_bitpack_pipeline_dense_with_nulls() {
+        // Create 1024 elements with some nulls.
+        let values: Vec<Option<u32>> = (0..1024)
+            .map(|i| if i % 100 == 0 { None } else { Some(i as u32) })
+            .collect();
+        let primitive = PrimitiveArray::from_option_iter(values).to_array();
+
+        // Encode with 10-bit width.
+        let bitpacked = BitPackedArray::encode(&primitive, 10).unwrap();
+        assert_eq!(bitpacked.bit_width(), 10, "Bit width should be 10");
+
+        // Select all elements (dense path).
+        let mask = Mask::new_true(1024);
+
+        // This should use the dense path since true_count >= 7.
+        let result = bitpacked.to_array().execute_with_selection(&mask).unwrap();
+        assert_eq!(result.len(), 1024, "Result should have 1024 elements");
+
+        let pvector_u32 = result.as_primitive().into_u32();
+        let elements = pvector_u32.elements().as_slice();
+
+        // Verify the values and validity.
+        for i in 0..1024 {
+            if i % 100 == 0 {
+                assert!(
+                    !pvector_u32.validity().value(i),
+                    "Element at index {} should be null",
+                    i
+                );
+            } else {
+                assert_eq!(
+                    elements[i], i as u32,
+                    "Element at index {} should be {}",
+                    i, i
+                );
+                assert!(
+                    pvector_u32.validity().value(i),
+                    "Element at index {} should be valid",
+                    i
+                );
+            }
+        }
+    }
 }

vortex-vector/src/primitive/generic_mut_impl.rs

@@ -27,15 +27,17 @@ impl<T: NativePType> PVectorMut<T> {
         self.validity.value(index).then(|| self.elements[index])
     }
 
-    /// Appends an element to the back of the vector.
+    /// Pushes an element to the back of the vector.
     ///
-    /// The element is treated as valid.
+    /// The element is treated as non-null.
     pub fn push(&mut self, value: T) {
         self.elements.push(value);
         self.validity.append_n(true, 1);
     }
 
-    /// Pushes a value without bounds checking or validity updates.
+    /// Pushes an element without bounds checking.
+    ///
+    /// The element is treated as non-null.
     ///
     /// # Safety
     ///