add nice-to-have vector methods to PVectorMut

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

Author: connortsui20

vortex-buffer/src/bit/buf_mut.rs

@@ -31,6 +31,9 @@ use crate::{BitBuffer, BufferMut, ByteBufferMut, buffer_mut};
 #[derive(Debug, Clone, Eq)]
 pub struct BitBufferMut {
     buffer: ByteBufferMut,
+    /// Represents the offset of the bit buffer into the first byte.
+    ///
+    /// This is always less than 8 (for when the bit buffer is not aligned to a byte).
     offset: usize,
     len: usize,
 }
@@ -162,6 +165,13 @@ impl BitBufferMut {
         self.buffer.reserve(additional_bytes);
     }
 
+    /// Clears the bit buffer (but keeps any allocated memory).
+    pub fn clear(&mut self) {
+        // Since there are no items we need to drop, we simply set the length to 0.
+        self.len = 0;
+        self.offset = 0;
+    }
+
     /// Set the bit at `index` to the given boolean value.
     ///
     /// This operation is checked so if `index` exceeds the buffer length, this will panic.

vortex-mask/src/mask_mut.rs

@@ -5,6 +5,7 @@ use std::ops::Sub;
 use std::sync::Arc;
 
 use vortex_buffer::BitBufferMut;
+use vortex_error::vortex_panic;
 
 use crate::Mask;
 
@@ -56,6 +57,29 @@ impl MaskMut {
         })
     }
 
+    /// Returns the boolean value at a given index.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the index is out of bounds.
+    pub fn value(&self, index: usize) -> bool {
+        match &self.0 {
+            Inner::Empty { .. } => {
+                vortex_panic!("index out of bounds: the length is 0 but the index is {index}")
+            }
+            Inner::Constant { value, len, .. } => {
+                assert!(
+                    index < *len,
+                    "index out of bounds: the length is {} but the index is {index}",
+                    *len
+                );
+
+                *value
+            }
+            Inner::Builder(bit_buffer) => bit_buffer.value(index),
+        }
+    }
+
     /// Reserve capacity for at least `additional` more values to be appended.
     pub fn reserve(&mut self, additional: usize) {
         match &mut self.0 {
@@ -71,6 +95,39 @@ impl MaskMut {
         }
     }
 
+    /// Clears the mask.
+    ///
+    /// Note that this method has no effect on the allocated capacity of the mask.
+    pub fn clear(&mut self) {
+        match &mut self.0 {
+            Inner::Empty { .. } => {}
+            Inner::Constant { capacity, .. } => {
+                self.0 = Inner::Empty {
+                    capacity: *capacity,
+                }
+            }
+            Inner::Builder(bit_buffer) => bit_buffer.clear(),
+        };
+    }
+
+    /// Shortens the mask, keeping the first `len` bits.
+    ///
+    /// If `len` is greater or equal to the vector’s current length, this has no effect.
+    ///
+    /// Note that this method has no effect on the allocated capacity of the mask.
+    pub fn truncate(&mut self, len: usize) {
+        let truncated_len = len;
+        if truncated_len > self.len() {
+            return;
+        }
+
+        match &mut self.0 {
+            Inner::Empty { .. } => {}
+            Inner::Constant { len, .. } => *len = truncated_len.min(*len),
+            Inner::Builder(bit_buffer) => bit_buffer.truncate(truncated_len),
+        };
+    }
+
     /// Append n values to the mask.
     pub fn append_n(&mut self, new_value: bool, n: usize) {
         match &mut self.0 {

vortex-vector/src/primitive/from_iter.rs

@@ -68,6 +68,76 @@ impl<T: NativePType> PVector<T> {
 
         Self { elements, validity }
     }
+
+    /// Gets a nullable element at the given index, with bounds checking.
+    ///
+    /// If the index is out of bounds, returns `None`. If the element at the given index is null,
+    /// returns `Some(None)`. Otherwise, returns `Some(Some(x))`, where `x: T`.
+    pub fn get_checked(&self, index: usize) -> Option<Option<T>> {
+        (index < self.len()).then(|| {
+            self.validity.value(index).then(|| {
+                self.elements
+                    .get(index)
+                    .copied()
+                    .vortex_expect("length of elements was somehow incorrect")
+            })
+        })
+    }
+
+    /// Gets a nullable element at the given index.
+    ///
+    /// If the element at the given index is null, returns `None`. Otherwise, returns `Some(x)`,
+    /// where `x: T`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the index is out of bounds.
+    pub fn get(&self, index: usize) -> Option<T> {
+        assert!(
+            index < self.len(),
+            "index out of bounds: the length is {} but the index is {index}",
+            self.len()
+        );
+
+        self.validity.value(index).then(|| {
+            self.elements
+                .get(index)
+                .copied()
+                .vortex_expect("length of elements was somehow incorrect")
+        })
+    }
+
+    /// Gets a nullable element at the given index, without checking bounds and without checking
+    /// nullability.
+    ///
+    /// The caller should ensure that the element at the given index is not null (though doing so
+    /// will not cause undefined behavior).
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that the index is within bounds.
+    pub unsafe fn get_unchecked(&self, index: usize) -> T {
+        debug_assert!(
+            index < self.len(),
+            "index out of bounds: the length is {} but the index is {index}",
+            self.len()
+        );
+
+        // SAFETY: The caller ensures that the index is in bounds.
+        unsafe { *self.elements.get_unchecked(index) }
+    }
+
+    /// Returns a slice over the internal buffer with elements of type `T`.
+    ///
+    /// Note that this slice may contain garbage data where the [`validity()`] mask states that an
+    /// element is invalid.
+    ///
+    /// The caller should check the [`validity()`] before performing any operations.
+    ///
+    /// [`validity()`]: Self::validity
+    pub fn as_slice(&self) -> &[T] {
+        self.elements.as_slice()
+    }
 }
 
 impl<T: NativePType> VectorOps for PVector<T> {

vortex-vector/src/primitive/generic.rs

@@ -15,9 +15,9 @@ use crate::{PVector, VectorMutOps, VectorOps};
 /// `T` is expected to be bound by [`NativePType`], which templates an internal [`BufferMut<T>`]
 /// that stores the elements of the vector.
 ///
-/// `PVectorMut<T>` is the primary way to construct primitive vectors. It provides efficient methods
-/// for building vectors incrementally before converting them to an immutable [`PVector<T>`] using
-/// the [`freeze`](crate::VectorMutOps::freeze) method.
+/// [`PVectorMut<T>`] is the primary way to construct primitive vectors. It provides efficient
+/// methods for building vectors incrementally before converting them to an immutable [`PVector<T>`]
+/// using the [`freeze`](crate::VectorMutOps::freeze) method.
 ///
 /// # Examples
 ///

vortex-vector/src/primitive/generic_mut.rs

@@ -23,11 +23,10 @@ pub use generic_mut::PVectorMut;
 mod vector;
 pub use vector::PrimitiveVector;
 
+mod pvector_impl;
 mod vector_mut;
 pub use vector_mut::PrimitiveVectorMut;
 
-mod from_iter;
-
 mod macros;
 
 use vortex_dtype::NativePType;