reset validity instead of filling

Signed-off-by: Mikhail Kot <to@myrrc.dev>

Author: myrrc

vortex-duckdb/cpp/include/duckdb_vx/vector.h

@@ -47,6 +47,10 @@ void duckdb_vx_string_vector_add_vector_data_buffer(duckdb_vector ffi_vector, du
 // valid.
 void duckdb_vx_vector_set_vector_data_buffer(duckdb_vector ffi_vector, duckdb_vx_vector_buffer buffer);
 
+// Reset vector's validity mask to nullptr, making all vector's elements valid.
+// vector must not be a DictionaryVector or a SequenceVector
+void duckdb_vx_vector_set_all_valid(duckdb_vector ffi_vector);
+
 // Set the data pointer for the vector. This is the start of the values array in the vector.
 void duckdb_vx_vector_set_data_ptr(duckdb_vector ffi_vector, void *ptr);
 

vortex-duckdb/cpp/vector.cpp

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+#include "include/duckdb_vx/vector.h"
 #include "duckdb_vx/duckdb_diagnostics.h"
 
 DUCKDB_INCLUDES_BEGIN
@@ -106,3 +107,16 @@ const char *duckdb_vector_to_string(duckdb_vector vector, unsigned long len, duc
         return nullptr;
     }
 }
+
+void duckdb_vx_vector_set_all_valid(duckdb_vector ffi_vector) {
+    using enum VectorType;
+    Vector& vector = *reinterpret_cast<Vector *>(ffi_vector);
+    const VectorType type = vector.GetVectorType();
+    D_ASSERT(type != DICTIONARY_VECTOR && type != SEQUENCE_VECTOR);
+    switch (type) {
+        case CONSTANT_VECTOR: return ConstantVector::Validity(vector).Reset();
+        case FLAT_VECTOR: return FlatVector::Validity(vector).Reset();
+        case FSST_VECTOR: return FSSTVector::Validity(vector).Reset();
+        default: __builtin_unreachable();
+    }
+}

vortex-duckdb/src/duckdb/vector.rs

@@ -207,7 +207,13 @@ impl VectorRef {
     ///
     /// The provided capacity *must* be the actual capacity of this vector.
     pub unsafe fn validity_bitslice_mut(&mut self, capacity: usize) -> Option<&mut BitSlice<u64>> {
-        unsafe { self.validity_slice_mut(capacity) }.map(|slice| slice.view_bits_mut())
+        // from_slice_unchecked_mut implies capacity must be less than BitSlice<u64>::MAX_ELTS
+        // but this is always true since validity is set per VectorRef/vector and its side doesn't
+        // exceed 2048
+        unsafe {
+            self.validity_slice_mut(capacity)
+                .map(|slice| BitSlice::from_slice_unchecked_mut(slice))
+        }
     }
 
     pub fn validity_ref(&self, len: usize) -> ValidityRef<'_> {

vortex-duckdb/src/exporter/mod.rs

@@ -19,8 +19,6 @@ mod validity;
 mod varbinview;
 mod vector;
 
-use bitvec::prelude::Lsb0;
-use bitvec::view::BitView;
 pub use cache::ConversionCache;
 pub use decimal::precision_to_duckdb_storage_size;
 use vortex::array::ArrayRef;
@@ -32,6 +30,7 @@ use vortex::array::arrays::List;
 use vortex::array::arrays::StructArray;
 use vortex::array::arrays::TemporalArray;
 use vortex::array::arrays::struct_::StructArrayExt;
+use vortex::buffer::BitChunks;
 use vortex::encodings::runend::RunEnd;
 use vortex::encodings::sequence::Sequence;
 use vortex::error::VortexResult;
@@ -187,17 +186,24 @@ fn new_array_exporter_with_flatten(
     }
 }
 
-/// Copy the sliced bits from source into target.
+/// Copy the sliced bits from source into target, counting ones.
 ///
 /// Offset and length are a _bit_ offset and a _bit_ length into source.
 ///
 /// `target.len()` must equal `len`.
-fn copy_from_slice(target: &mut [u64], source: &[u8], offset: usize, len: usize) {
-    let (start, middle, end) = unsafe { target.align_to_mut::<u8>() };
-    assert!(start.is_empty());
-    assert!(end.is_empty());
-    let target = &mut middle.view_bits_mut::<Lsb0>()[..len];
-    target.copy_from_bitslice(&source.view_bits()[offset..][..len]);
+fn copy_from_slice(target: &mut [u64], source: &[u8], offset: usize, len: usize) -> usize {
+    let chunks = BitChunks::new(source, offset, len);
+    let mut count = 0usize;
+    for (i, chunk) in chunks.iter().enumerate() {
+        target[i] = chunk;
+        count += chunk.count_ones() as usize;
+    }
+    if chunks.remainder_len() > 0 {
+        let remainder = chunks.remainder_bits();
+        target[chunks.chunk_len()] = remainder;
+        count += remainder.count_ones() as usize;
+    }
+    count
 }
 
 #[cfg(test)]

vortex-duckdb/src/exporter/vector.rs

@@ -2,7 +2,9 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex::mask::Mask;
+use vortex::mask::MaskValues;
 
+use crate::cpp::duckdb_vx_vector_set_all_valid;
 use crate::duckdb::Value;
 use crate::duckdb::VectorRef;
 use crate::exporter::copy_from_slice;
@@ -11,44 +13,44 @@ impl VectorRef {
     pub(super) unsafe fn set_validity(&mut self, mask: &Mask, offset: usize, len: usize) -> bool {
         match mask {
             Mask::AllTrue(_) => {
-                // We only need to blank out validity if there is already a slice allocated.
-                // SAFETY: Caller guarantees this.
-                unsafe { self.set_all_true_validity(len) }
+                self.set_all_true_validity();
                 false
             }
             Mask::AllFalse(_) => {
-                // SAFETY: Caller guarantees this.
                 self.set_all_false_validity();
                 true
             }
-            Mask::Values(arr) => {
-                let true_count = arr.bit_buffer().slice(offset..(offset + len)).true_count();
-                if true_count == len {
-                    unsafe { self.set_all_true_validity(len) }
-                } else if true_count == 0 {
-                    self.set_all_false_validity()
-                } else {
-                    let source = arr.bit_buffer().inner().as_slice();
-                    copy_from_slice(
-                        unsafe { self.ensure_validity_slice(len) },
-                        source,
-                        offset,
-                        len,
-                    );
-                }
-
-                true_count == 0
-            }
+            Mask::Values(arr) => self.set_validity_with_array(&arr, len, offset),
         }
     }
 
-    pub(super) unsafe fn set_all_true_validity(&mut self, len: usize) {
-        if let Some(validity) = unsafe { self.validity_bitslice_mut(len) } {
-            validity.fill(true);
+    fn set_validity_with_array(&mut self, arr: &MaskValues, len: usize, offset: usize) -> bool {
+        let true_count = arr.true_count();
+        if true_count == len {
+            self.set_all_true_validity();
+            return false;
+        } else if true_count == 0 {
+            self.set_all_false_validity();
+            return true;
+        }
+
+        let source = arr.bit_buffer().inner().as_slice();
+        let dest = unsafe { self.ensure_validity_slice(len) };
+        let true_count = copy_from_slice(dest, source, offset, len);
+        if true_count == len {
+            self.set_all_true_validity()
+        } else if true_count == 0 {
+            self.set_all_false_validity()
         }
+
+        true_count == 0
+    }
+
+    fn set_all_true_validity(&mut self) {
+        unsafe { duckdb_vx_vector_set_all_valid(self.as_ptr()) };
     }
 
-    pub(super) fn set_all_false_validity(&mut self) {
+    fn set_all_false_validity(&mut self) {
         self.reference_value(&Value::null(&self.logical_type()));
     }
 }