fix: Validity export using shifted u64 copy (#4244)

Signed-off-by: Robert Kruszewski <[email protected]>

Author: robert3005

vortex-duckdb/src/exporter/mod.rs

@@ -15,6 +15,8 @@ mod varbinview;
 
 use std::sync::Arc;
 
+use bitvec::prelude::Lsb0;
+use bitvec::view::BitView;
 pub use cache::ConversionCache;
 pub use decimal::precision_to_duckdb_storage_size;
 use itertools::Itertools;
@@ -208,18 +210,186 @@ impl VectorExt for Vector {
                 true
             }
             Mask::Values(arr) => {
-                // TODO(joe): do this MUCH better, with a shifted u64 copy
-                let mut null_count = 0;
-                // SAFETY: Caller guaranteees this.
-                let validity = unsafe { self.ensure_validity_slice(len) };
-                for (idx, v) in arr.boolean_buffer().slice(offset, len).iter().enumerate() {
-                    if !v {
-                        validity.set(idx, false);
-                        null_count += 1;
+                let arr_bits: &[u64] = {
+                    let byte_slice = arr.boolean_buffer().inner().as_slice();
+                    unsafe {
+                        std::slice::from_raw_parts(
+                            byte_slice.as_ptr() as _,
+                            byte_slice.len().div_ceil(size_of::<u64>()),
+                        )
+                    }
+                };
+                let sliced_bits = &arr_bits.view_bits::<Lsb0>()[offset..][..len];
+                let true_count = sliced_bits.count_ones();
+                if true_count == len {
+                    if let Some(validity) = unsafe { self.validity_slice_mut(len) } {
+                        validity.fill(true);
                     }
+                } else if true_count == 0 {
+                    unsafe { self.ensure_validity_slice(len) }.fill(false);
+                } else {
+                    // SAFETY: Caller guaranteees this.
+                    let validity = unsafe { self.ensure_validity_slice(len) };
+                    validity[..len].copy_from_bitslice(sliced_bits);
                 }
-                null_count == len
+                true_count == 0
             }
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use arrow_buffer::buffer::BooleanBuffer;
+    use vortex::mask::Mask;
+
+    use super::VectorExt;
+    use crate::cpp::DUCKDB_TYPE;
+    use crate::duckdb::{LogicalType, Vector};
+
+    #[test]
+    fn test_set_validity_all_true() {
+        let logical_type = LogicalType::new(DUCKDB_TYPE::DUCKDB_TYPE_BIGINT);
+        let mut vector = Vector::with_capacity(logical_type, 100);
+
+        let mask = Mask::AllTrue(10);
+        let all_null = unsafe { vector.set_validity(&mask, 0, 10) };
+
+        assert!(!all_null);
+    }
+
+    #[test]
+    fn test_set_validity_all_false() {
+        let logical_type = LogicalType::new(DUCKDB_TYPE::DUCKDB_TYPE_BIGINT);
+        let mut vector = Vector::with_capacity(logical_type, 100);
+
+        let mask = Mask::AllFalse(10);
+        let all_null = unsafe { vector.set_validity(&mask, 0, 10) };
+
+        assert!(all_null);
+
+        let validity = unsafe { vector.validity_slice_mut(10).unwrap() };
+        assert!(validity.iter().all(|v| !v));
+    }
+
+    #[test]
+    fn test_set_validity_values_all_true() {
+        let logical_type = LogicalType::new(DUCKDB_TYPE::DUCKDB_TYPE_BIGINT);
+        let mut vector = Vector::with_capacity(logical_type, 100);
+
+        let bits = vec![true; 10];
+        let buffer = BooleanBuffer::from(bits.as_slice());
+        let mask = Mask::from(buffer);
+
+        let all_null = unsafe { vector.set_validity(&mask, 0, 10) };
+
+        assert!(!all_null);
+
+        // When all values are true, the mask may be optimized to AllTrue,
+        // so validity_slice_mut may return None (no validity allocated)
+        if let Some(validity) = unsafe { vector.validity_slice_mut(10) } {
+            assert!(validity.iter().all(|v| *v));
+        }
+    }
+
+    #[test]
+    fn test_set_validity_values_all_false() {
+        let logical_type = LogicalType::new(DUCKDB_TYPE::DUCKDB_TYPE_BIGINT);
+        let mut vector = Vector::with_capacity(logical_type, 100);
+
+        let bits = vec![false; 10];
+        let buffer = BooleanBuffer::from(bits.as_slice());
+        let mask = Mask::from(buffer);
+
+        let all_null = unsafe { vector.set_validity(&mask, 0, 10) };
+
+        assert!(all_null);
+
+        let validity = unsafe { vector.validity_slice_mut(10).unwrap() };
+        assert!(validity.iter().all(|v| !v));
+    }
+
+    #[test]
+    fn test_set_validity_values_mixed() {
+        let logical_type = LogicalType::new(DUCKDB_TYPE::DUCKDB_TYPE_BIGINT);
+        let mut vector = Vector::with_capacity(logical_type, 100);
+
+        let bits = vec![
+            true, false, true, true, false, false, true, true, false, true,
+        ];
+        let buffer = BooleanBuffer::from(bits.as_slice());
+        let mask = Mask::from(buffer);
+
+        let all_null = unsafe { vector.set_validity(&mask, 0, 10) };
+
+        assert!(!all_null);
+
+        let validity = unsafe { vector.validity_slice_mut(10).unwrap() };
+        for (i, bit) in bits.iter().enumerate() {
+            assert_eq!(validity[i], *bit);
+        }
+    }
+
+    #[test]
+    fn test_set_validity_values_with_offset() {
+        let logical_type = LogicalType::new(DUCKDB_TYPE::DUCKDB_TYPE_BIGINT);
+        let mut vector = Vector::with_capacity(logical_type, 100);
+
+        let bits = vec![
+            false, false, true, true, false, true, false, true, true, false, true, true, false,
+        ];
+        let buffer = BooleanBuffer::from(bits.as_slice());
+        let mask = Mask::from(buffer);
+
+        let all_null = unsafe { vector.set_validity(&mask, 2, 8) };
+
+        assert!(!all_null);
+
+        let validity = unsafe { vector.validity_slice_mut(8).unwrap() };
+        for i in 0..8 {
+            assert_eq!(validity[i], bits[i + 2]);
+        }
+    }
+
+    #[test]
+    fn test_set_validity_values_with_offset_and_smaller_len() {
+        let logical_type = LogicalType::new(DUCKDB_TYPE::DUCKDB_TYPE_BIGINT);
+        let mut vector = Vector::with_capacity(logical_type, 100);
+
+        let bits = vec![
+            true, false, true, true, false, false, true, true, false, true, true, true, false,
+            true, false,
+        ];
+        let buffer = BooleanBuffer::from(bits.as_slice());
+        let mask = Mask::from(buffer);
+
+        let all_null = unsafe { vector.set_validity(&mask, 3, 5) };
+
+        assert!(!all_null);
+
+        let validity = unsafe { vector.validity_slice_mut(5).unwrap() };
+        for i in 0..5 {
+            assert_eq!(validity[i], bits[i + 3]);
+        }
+    }
+
+    #[test]
+    fn test_set_validity_values_64bit_alignment() {
+        let logical_type = LogicalType::new(DUCKDB_TYPE::DUCKDB_TYPE_BIGINT);
+        let mut vector = Vector::with_capacity(logical_type, 100);
+
+        let bits = (0..70).map(|i| i % 3 == 0).collect::<Vec<_>>();
+
+        let buffer = BooleanBuffer::from(bits.as_slice());
+        let mask = Mask::from(buffer);
+
+        let all_null = unsafe { vector.set_validity(&mask, 5, 60) };
+
+        assert!(!all_null);
+
+        let validity = unsafe { vector.validity_slice_mut(60).unwrap() };
+        for i in 0..60 {
+            assert_eq!(validity[i], bits[i + 5]);
+        }
+    }
+}