Feature: Struct Vectors (#5094)

Tracking Issue: #5028

Adds `StructVector` and `StructVectorMut`.

Note: It would be really nice to have something like serde values that
we can iterate and take as an iterator to construct these struct
vectors, otherwise we are left with what I currently have on this PR,
which imo is very unwieldy. This can come in a later PR though.

The second commit fixes some errors and also removes all of the `new`
functions for vectors (as per an offline discussion with @gatesn).

---------

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

Author: connortsui20

vortex-array/src/arrays/bool/vtable/operator.rs

@@ -27,7 +27,7 @@ impl OperatorVTable<BoolVTable> for BoolVTable {
             // Note that validity already has the mask applied so we only need to apply it to bits.
             let bits = bits.filter(&mask);
 
-            Ok(BoolVector::new(bits, validity).into())
+            Ok(BoolVector::try_new(bits, validity)?.into())
         }))
     }
 }

vortex-array/src/arrays/primitive/vtable/operator.rs

@@ -30,7 +30,7 @@ impl OperatorVTable<PrimitiveVTable> for PrimitiveVTable {
                 // the elements.
                 let elements = elements.filter(&mask);
 
-                Ok(PVector::new(elements, validity).into())
+                Ok(PVector::try_new(elements, validity)?.into())
             }))
         })
     }

vortex-compute/src/filter/bool.rs

@@ -8,6 +8,12 @@ use crate::filter::Filter;
 
 impl Filter for BoolVector {
     fn filter(&self, mask: &Mask) -> Self {
-        Self::new(self.bits().filter(mask), self.validity().filter(mask))
+        let filtered_bits = self.bits().filter(mask);
+        let filtered_validity = self.validity().filter(mask);
+
+        // SAFETY: We filter the bits and validity with the same mask, and since they came from an
+        // existing and valid `BoolVector`, we know that the filtered output must have the same
+        // length.
+        unsafe { Self::new_unchecked(filtered_bits, filtered_validity) }
     }
 }

vortex-vector/src/bool/iter.rs

@@ -85,6 +85,7 @@ impl FromIterator<bool> for BoolVectorMut {
 ///
 /// It consumes the mutable vector and iterates over the elements, yielding `None` for null values
 /// and `Some(value)` for valid values.
+#[derive(Debug)]
 pub struct BoolVectorMutIterator {
     /// The vector being iterated over.
     vector: BoolVectorMut,

vortex-vector/src/bool/vector.rs

@@ -30,8 +30,7 @@ impl BoolVector {
     ///
     /// Panics if the length of the validity mask does not match the length of the bits.
     pub fn new(bits: BitBuffer, validity: Mask) -> Self {
-        Self::try_new(bits, validity)
-            .vortex_expect("`BoolVector` validity mask must have the same length as bits")
+        Self::try_new(bits, validity).vortex_expect("Failed to create `BoolVector`")
     }
 
     /// Tries to create a new [`BoolVector`] from the given bits and validity mask.
@@ -53,17 +52,15 @@ impl BoolVector {
     /// # Safety
     ///
     /// The caller must ensure that the validity mask has the same length as the bits.
-    pub fn new_unchecked(bits: BitBuffer, validity: Mask) -> Self {
-        debug_assert_eq!(
-            validity.len(),
-            bits.len(),
-            "`BoolVector` validity mask must have the same length as bits"
-        );
-
-        Self { bits, validity }
+    pub unsafe fn new_unchecked(bits: BitBuffer, validity: Mask) -> Self {
+        if cfg!(debug_assertions) {
+            Self::new(bits, validity)
+        } else {
+            Self { bits, validity }
+        }
     }
 
-    /// Decomposes the boolean vector into its constituent parts.
+    /// Decomposes the boolean vector into its constituent parts (bit buffer and validity).
     pub fn into_parts(self) -> (BitBuffer, Mask) {
         (self.bits, self.validity)
     }

vortex-vector/src/bool/vector_mut.rs

@@ -77,8 +77,7 @@ impl BoolVectorMut {
     ///
     /// Panics if the length of the validity mask does not match the length of the bits.
     pub fn new(bits: BitBufferMut, validity: MaskMut) -> Self {
-        Self::try_new(bits, validity)
-            .vortex_expect("`BoolVector` validity mask must have the same length as bits")
+        Self::try_new(bits, validity).vortex_expect("Failed to create `BoolVectorMut`")
     }
 
     /// Tries to create a new [`BoolVectorMut`] from the given bits and validity mask.
@@ -104,13 +103,11 @@ impl BoolVectorMut {
     /// Ideally, they are taken from `into_parts`, mutated in a way that doesn't re-allocate, and
     /// then passed back to this function.
     pub unsafe fn new_unchecked(bits: BitBufferMut, validity: MaskMut) -> Self {
-        debug_assert_eq!(
-            bits.len(),
-            validity.len(),
-            "`BoolVector` validity mask must have the same length as bits"
-        );
-
-        Self { bits, validity }
+        if cfg!(debug_assertions) {
+            Self::new(bits, validity)
+        } else {
+            Self { bits, validity }
+        }
     }
 
     /// Creates a new mutable boolean vector with the given `capacity`.
@@ -121,7 +118,7 @@ impl BoolVectorMut {
         }
     }
 
-    /// Returns the parts of the mutable vector.
+    /// Decomposes the boolean vector into its constituent parts (bit buffer and validity).
     pub fn into_parts(self) -> (BitBufferMut, MaskMut) {
         (self.bits, self.validity)
     }
@@ -151,7 +148,7 @@ impl VectorMutOps for BoolVectorMut {
     }
 
     fn append_nulls(&mut self, n: usize) {
-        self.bits.append_n(false, n);
+        self.bits.append_n(false, n); // Note that the value we push doesn't actually matter.
         self.validity.append_n(false, n);
     }
 

vortex-vector/src/lib.rs

@@ -12,17 +12,22 @@
 #![deny(clippy::missing_safety_doc)]
 
 mod bool;
-mod macros;
 mod null;
-mod ops;
 mod primitive;
-mod private;
-mod vector;
-mod vector_mut;
+mod struct_;
 
 pub use bool::{BoolVector, BoolVectorMut};
 pub use null::{NullVector, NullVectorMut};
-pub use ops::{VectorMutOps, VectorOps};
 pub use primitive::{PVector, PVectorMut, PrimitiveVector, PrimitiveVectorMut};
+pub use struct_::{StructVector, StructVectorMut};
+
+mod ops;
+mod vector;
+mod vector_mut;
+
+pub use ops::{VectorMutOps, VectorOps};
 pub use vector::Vector;
 pub use vector_mut::VectorMut;
+
+mod macros;
+mod private;

vortex-vector/src/macros.rs

@@ -48,6 +48,10 @@ macro_rules! match_each_vector {
                 let $vec = v;
                 $body
             }
+            $crate::Vector::Struct(v) => {
+                let $vec = v;
+                $body
+            }
         }
     }};
 }
@@ -98,6 +102,10 @@ macro_rules! match_each_vector_mut {
                 let $vec = v;
                 $body
             }
+            $crate::VectorMut::Struct(v) => {
+                let $vec = v;
+                $body
+            }
         }
     }};
 }

vortex-vector/src/ops.rs

@@ -70,6 +70,11 @@ pub trait VectorMutOps: private::Sealed + Into<VectorMut> {
     fn reserve(&mut self, additional: usize);
 
     /// Extends the vector by appending elements from another vector.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the `other` vector has the wrong type (for example, a
+    /// [`StructVector`](crate::StructVector) might have incorrect fields).
     fn extend_from_vector(&mut self, other: &Self::Immutable);
 
     /// Appends `n` null elements to the vector.

vortex-vector/src/primitive/generic.rs

@@ -34,8 +34,7 @@ impl<T: NativePType> PVector<T> {
     ///
     /// Panics if the length of the validity mask does not match the length of the elements buffer.
     pub fn new(elements: Buffer<T>, validity: Mask) -> Self {
-        Self::try_new(elements, validity)
-            .vortex_expect("`PVector` validity mask must have the same length as elements")
+        Self::try_new(elements, validity).vortex_expect("Failed to create `PVector`")
     }
 
     /// Tries to create a new [`PVector<T>`] from the given elements buffer and validity mask.
@@ -59,17 +58,15 @@ impl<T: NativePType> PVector<T> {
     /// # Safety
     ///
     /// The caller must ensure that the validity mask has the same length as the elements buffer.
-    pub fn new_unchecked(elements: Buffer<T>, validity: Mask) -> Self {
-        debug_assert_eq!(
-            validity.len(),
-            elements.len(),
-            "`PVector` validity mask must have the same length as elements"
-        );
-
-        Self { elements, validity }
+    pub unsafe fn new_unchecked(elements: Buffer<T>, validity: Mask) -> Self {
+        if cfg!(debug_assertions) {
+            Self::new(elements, validity)
+        } else {
+            Self { elements, validity }
+        }
     }
 
-    /// Decomposes the primitive vector into its constituent parts.
+    /// Decomposes the primitive vector into its constituent parts (buffer and validity).
     pub fn into_parts(self) -> (Buffer<T>, Mask) {
         (self.elements, self.validity)
     }