reset validity instead of filling

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

Author: myrrc

vortex-buffer/src/bit/buf.rs

@@ -64,6 +64,17 @@ impl PartialEq for BitBuffer {
     }
 }
 
+/// How many bits are set to one in ButBuffer
+#[derive(PartialEq, Eq, Debug)]
+pub enum ZeroesCount {
+    /// All bits are set
+    AllZeroes,
+    /// All bits are unset
+    AllOnes,
+    /// Some bits are set
+    Unknown,
+}
+
 impl BitBuffer {
     /// Create a new `BoolBuffer` backed by a [`ByteBuffer`] with `len` bits in view.
     ///
@@ -319,6 +330,36 @@ impl BitBuffer {
         count_ones(self.buffer.as_slice(), self.offset, self.len)
     }
 
+    /// Copy bits to target slice, returning whether all bits are zero, all are one, or mixed.
+    /// "target" must be a buffer large enough to hold "len" bits where "len" is buffer's len.
+    pub fn copy_to_slice(&self, target: &mut [u64]) -> ZeroesCount {
+        let mut sum_or = 0u64;
+        let mut sum_and = u64::MAX;
+        let chunks = self.chunks();
+        let remainder_bits = chunks.remainder_bits();
+        let remainder_len = chunks.remainder_len();
+        let chunk_len = chunks.chunk_len();
+        for (slot, chunk) in target.iter_mut().zip(chunks) {
+            *slot = chunk;
+            sum_or |= chunk;
+            sum_and &= chunk;
+        }
+        if remainder_len > 0 {
+            target[chunk_len] = remainder_bits;
+            sum_or |= remainder_bits;
+            let mask = (1u64 << remainder_len) - 1;
+            sum_and &= remainder_bits | !mask;
+        }
+
+        if sum_or == 0 {
+            ZeroesCount::AllZeroes
+        } else if sum_and == u64::MAX {
+            ZeroesCount::AllOnes
+        } else {
+            ZeroesCount::Unknown
+        }
+    }
+
     /// Get the number of unset bits in the buffer.
     pub fn false_count(&self) -> usize {
         self.len - self.true_count()
@@ -571,6 +612,7 @@ mod tests {
     use rstest::rstest;
 
     use crate::ByteBuffer;
+    use crate::ZeroesCount;
     use crate::bit::BitBuffer;
     use crate::buffer;
 
@@ -781,4 +823,104 @@ mod tests {
             assert_eq!(mapped.value(i), expected, "Mismatch at index {}", i);
         }
     }
+
+    fn copy_equal(source: BitBuffer, expected_count: ZeroesCount) {
+        let n = source.len().div_ceil(64);
+        let mut target = vec![0u64; n];
+        let result = source.copy_to_slice(&mut target);
+        assert_eq!(result, expected_count);
+        for i in 0..source.len() {
+            let copied = (target[i / 64] >> (i % 64)) & 1 == 1;
+            assert_eq!(copied, source.value(i), "Bit mismatch at index {i}");
+        }
+    }
+
+    #[rstest]
+    #[case(10)]
+    #[case(64)]
+    #[case(65)]
+    #[case(128)]
+    #[case(1024)]
+    fn test_copy_to_slice_all_true(#[case] len_bits: usize) {
+        copy_equal(BitBuffer::new_set(len_bits), ZeroesCount::AllOnes);
+    }
+
+    #[rstest]
+    #[case(10)]
+    #[case(64)]
+    #[case(65)]
+    #[case(128)]
+    #[case(1024)]
+    fn test_copy_to_slice_all_false(#[case] len_bits: usize) {
+        copy_equal(BitBuffer::new_unset(len_bits), ZeroesCount::AllZeroes);
+    }
+
+    #[rstest]
+    #[case(10)]
+    #[case(64)]
+    #[case(65)]
+    #[case(128)]
+    #[case(1024)]
+    fn test_copy_to_slice_all_true_unaligned(#[case] len_bits: usize) {
+        let buffer = buffer![u8::MAX; 8 * 20];
+        copy_equal(
+            BitBuffer::new_with_offset(buffer, len_bits, 1),
+            ZeroesCount::AllOnes,
+        );
+    }
+
+    #[rstest]
+    #[case(10)]
+    #[case(64)]
+    #[case(65)]
+    #[case(128)]
+    #[case(1024)]
+    fn test_copy_to_slice_all_false_unaligned(#[case] len_bits: usize) {
+        let buffer = buffer![0; 8 * 20];
+        copy_equal(
+            BitBuffer::new_with_offset(buffer, len_bits, 1),
+            ZeroesCount::AllZeroes,
+        );
+    }
+
+    #[test]
+    fn test_copy_to_slice_some() {
+        let mut target = vec![0u64];
+        let source = BitBuffer::new(buffer![1u8, 2, 3, 50, 51, 52, 100, 101], 64);
+        let zero_count = source.copy_to_slice(&mut target);
+        assert_eq!(
+            target[0], 0x65_64_34_33_32_03_02_01_u64,
+            "{:#08x} == {:#08x}",
+            target[0], 0x65_64_34_33_32_03_02_01_u64,
+        );
+        assert_eq!(zero_count, ZeroesCount::Unknown);
+    }
+
+    #[test]
+    fn test_copy_to_slice_some_unaligned() {
+        let source = BitBuffer::new_with_offset(buffer![1u8, 2, 3, 50, 51, 52, 100, 101, 0], 64, 1);
+        let mut target = vec![0u64; 1];
+        let zero_count = source.copy_to_slice(&mut target);
+        assert_eq!(zero_count, ZeroesCount::Unknown);
+        for i in 0..64 {
+            let copied = (target[0] >> i) & 1 == 1;
+            assert_eq!(copied, source.value(i), "Bit mismatch at index {i}");
+        }
+    }
+
+    #[test]
+    fn test_copy_to_slice_some_65() {
+        let mut target = vec![0u64; 2];
+        let source = BitBuffer::new(buffer![0u8, 0, 0, 0, 0, 0, 0, 0, 1], 65);
+        let zero_count = source.copy_to_slice(&mut target);
+        assert_eq!(zero_count, ZeroesCount::Unknown);
+    }
+
+    #[test]
+    fn test_copy_to_slice_some_65_unaligned() {
+        let mut target = vec![0u64; 2];
+        let source = BitBuffer::new_with_offset(buffer![0u8, 1, 0, 0, 0, 0, 0, 0, 0], 65, 1);
+        let zero_count = source.copy_to_slice(&mut target);
+        assert_eq!(zero_count, ZeroesCount::Unknown);
+    }
 }

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,20 @@ 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,11 @@ 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())
+        // capacity is always less than BitSlice<u64>::MAX_ELTS
+        unsafe {
+            self.validity_slice_mut(capacity)
+                .map(|slice| BitSlice::from_slice_unchecked_mut(slice))
+        }
     }
 
     pub fn validity_ref(&self, len: usize) -> ValidityRef<'_> {