add constructors and validation to list vector

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

Author: connortsui20

vortex-buffer/src/bit/buf_mut.rs

@@ -3,7 +3,7 @@
 
 use std::ops::Not;
 
-use arrow_buffer::bit_chunk_iterator::BitChunks;
+use arrow_buffer::bit_chunk_iterator::{BitChunks, UnalignedBitChunk};
 use bitvec::view::BitView;
 
 use crate::bit::{get_bit_unchecked, ops, set_bit_unchecked, unset_bit_unchecked};
@@ -495,6 +495,21 @@ impl BitBufferMut {
     pub fn as_mut_ptr(&mut self) -> *mut u8 {
         self.buffer.as_mut_ptr()
     }
+
+    /// Access chunks of the buffer aligned to 8 byte boundary as [prefix, \<full chunks\>, suffix]
+    pub fn unaligned_chunks(&self) -> UnalignedBitChunk<'_> {
+        UnalignedBitChunk::new(self.buffer.as_slice(), self.offset, self.len)
+    }
+
+    /// Get the number of set bits in the buffer.
+    pub fn true_count(&self) -> usize {
+        self.unaligned_chunks().count_ones()
+    }
+
+    /// Get the number of unset bits in the buffer.
+    pub fn false_count(&self) -> usize {
+        self.len - self.true_count()
+    }
 }
 
 impl Default for BitBufferMut {

vortex-mask/src/lib.rs

@@ -123,93 +123,6 @@ pub struct MaskValues {
     density: f64,
 }
 
-impl MaskValues {
-    /// Returns the length of the mask.
-    #[inline]
-    pub fn len(&self) -> usize {
-        self.buffer.len()
-    }
-
-    /// Returns true if the mask is empty i.e., it's length is 0.
-    #[inline]
-    pub fn is_empty(&self) -> bool {
-        self.buffer.is_empty()
-    }
-
-    /// Returns the true count of the mask.
-    #[inline]
-    pub fn true_count(&self) -> usize {
-        self.true_count
-    }
-
-    /// Returns the boolean buffer representation of the mask.
-    #[inline]
-    pub fn bit_buffer(&self) -> &BitBuffer {
-        &self.buffer
-    }
-
-    /// Returns the boolean value at a given index.
-    #[inline]
-    pub fn value(&self, index: usize) -> bool {
-        self.buffer.value(index)
-    }
-
-    /// Constructs an indices vector from one of the other representations.
-    pub fn indices(&self) -> &[usize] {
-        self.indices.get_or_init(|| {
-            if self.true_count == 0 {
-                return vec![];
-            }
-
-            if self.true_count == self.len() {
-                return (0..self.len()).collect();
-            }
-
-            if let Some(slices) = self.slices.get() {
-                let mut indices = Vec::with_capacity(self.true_count);
-                indices.extend(slices.iter().flat_map(|(start, end)| *start..*end));
-                debug_assert!(indices.is_sorted());
-                assert_eq!(indices.len(), self.true_count);
-                return indices;
-            }
-
-            let mut indices = Vec::with_capacity(self.true_count);
-            indices.extend(self.buffer.set_indices());
-            debug_assert!(indices.is_sorted());
-            assert_eq!(indices.len(), self.true_count);
-            indices
-        })
-    }
-
-    /// Constructs a slices vector from one of the other representations.
-    #[allow(clippy::cast_possible_truncation)]
-    #[inline]
-    pub fn slices(&self) -> &[(usize, usize)] {
-        self.slices.get_or_init(|| {
-            if self.true_count == self.len() {
-                return vec![(0, self.len())];
-            }
-
-            self.buffer.set_slices().collect()
-        })
-    }
-
-    /// Return an iterator over either indices or slices of the mask based on a density threshold.
-    #[inline]
-    pub fn threshold_iter(&self, threshold: f64) -> MaskIter<'_> {
-        if self.density >= threshold {
-            MaskIter::Slices(self.slices())
-        } else {
-            MaskIter::Indices(self.indices())
-        }
-    }
-
-    /// Extracts the internal [`BitBuffer`].
-    pub(crate) fn into_buffer(self) -> BitBuffer {
-        self.buffer
-    }
-}
-
 impl Mask {
     /// Create a new Mask where all values are set.
     #[inline]
@@ -642,6 +555,93 @@ impl Mask {
     }
 }
 
+impl MaskValues {
+    /// Returns the length of the mask.
+    #[inline]
+    pub fn len(&self) -> usize {
+        self.buffer.len()
+    }
+
+    /// Returns true if the mask is empty i.e., it's length is 0.
+    #[inline]
+    pub fn is_empty(&self) -> bool {
+        self.buffer.is_empty()
+    }
+
+    /// Returns the true count of the mask.
+    #[inline]
+    pub fn true_count(&self) -> usize {
+        self.true_count
+    }
+
+    /// Returns the boolean buffer representation of the mask.
+    #[inline]
+    pub fn bit_buffer(&self) -> &BitBuffer {
+        &self.buffer
+    }
+
+    /// Returns the boolean value at a given index.
+    #[inline]
+    pub fn value(&self, index: usize) -> bool {
+        self.buffer.value(index)
+    }
+
+    /// Constructs an indices vector from one of the other representations.
+    pub fn indices(&self) -> &[usize] {
+        self.indices.get_or_init(|| {
+            if self.true_count == 0 {
+                return vec![];
+            }
+
+            if self.true_count == self.len() {
+                return (0..self.len()).collect();
+            }
+
+            if let Some(slices) = self.slices.get() {
+                let mut indices = Vec::with_capacity(self.true_count);
+                indices.extend(slices.iter().flat_map(|(start, end)| *start..*end));
+                debug_assert!(indices.is_sorted());
+                assert_eq!(indices.len(), self.true_count);
+                return indices;
+            }
+
+            let mut indices = Vec::with_capacity(self.true_count);
+            indices.extend(self.buffer.set_indices());
+            debug_assert!(indices.is_sorted());
+            assert_eq!(indices.len(), self.true_count);
+            indices
+        })
+    }
+
+    /// Constructs a slices vector from one of the other representations.
+    #[allow(clippy::cast_possible_truncation)]
+    #[inline]
+    pub fn slices(&self) -> &[(usize, usize)] {
+        self.slices.get_or_init(|| {
+            if self.true_count == self.len() {
+                return vec![(0, self.len())];
+            }
+
+            self.buffer.set_slices().collect()
+        })
+    }
+
+    /// Return an iterator over either indices or slices of the mask based on a density threshold.
+    #[inline]
+    pub fn threshold_iter(&self, threshold: f64) -> MaskIter<'_> {
+        if self.density >= threshold {
+            MaskIter::Slices(self.slices())
+        } else {
+            MaskIter::Indices(self.indices())
+        }
+    }
+
+    /// Extracts the internal [`BitBuffer`].
+    pub(crate) fn into_buffer(self) -> BitBuffer {
+        self.buffer
+    }
+}
+
 /// Iterator over the indices or slices of a mask.
 pub enum MaskIter<'a> {
     /// Slice of pre-cached indices of a mask.

vortex-mask/src/mask_mut.rs

@@ -279,6 +279,24 @@ impl MaskMut {
     pub fn is_empty(&self) -> bool {
         self.len() == 0
     }
+
+    /// Returns true if all values in the mask are true.
+    pub fn all_true(&self) -> bool {
+        match &self.0 {
+            Inner::Empty { .. } => true,
+            Inner::Constant { value, .. } => *value,
+            Inner::Builder(bits) => bits.true_count() == bits.len(),
+        }
+    }
+
+    /// Returns true if all values in the mask are false.
+    pub fn all_false(&self) -> bool {
+        match &self.0 {
+            Inner::Empty { .. } => true,
+            Inner::Constant { value, .. } => !*value,
+            Inner::Builder(bits) => !bits.is_empty() && bits.true_count() == 0,
+        }
+    }
 }
 
 impl Mask {