add StructVector implementation

connortsui20 · connortsui20 · commit 1eed89b7e1c2 · 2025-10-27T17:14:50.000-04:00
Signed-off-by: Connor Tsui &lt;connor.tsui20@gmail.com&gt;
diff --git a/vortex-vector/src/bool/iter.rs b/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,
diff --git a/vortex-vector/src/bool/vector.rs b/vortex-vector/src/bool/vector.rs
@@ -63,7 +63,7 @@ impl BoolVector {
         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)
     }
diff --git a/vortex-vector/src/bool/vector_mut.rs b/vortex-vector/src/bool/vector_mut.rs
@@ -121,7 +121,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 +151,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);
     }
 
diff --git a/vortex-vector/src/lib.rs b/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;
diff --git a/vortex-vector/src/macros.rs b/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
+            }
         }
     }};
 }
diff --git a/vortex-vector/src/ops.rs b/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.
diff --git a/vortex-vector/src/primitive/generic.rs b/vortex-vector/src/primitive/generic.rs
@@ -69,7 +69,7 @@ impl<T: NativePType> PVector<T> {
         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)
     }
diff --git a/vortex-vector/src/primitive/generic_mut.rs b/vortex-vector/src/primitive/generic_mut.rs
@@ -160,7 +160,7 @@ impl<T> PVectorMut<T> {
         }
     }
 
-    /// Decomposes the primitive vector into its constituent parts.
+    /// Decomposes the primitive vector into its constituent parts (buffer and validity).
     pub fn into_parts(self) -> (BufferMut<T>, MaskMut) {
         (self.elements, self.validity)
     }
@@ -189,7 +189,7 @@ impl<T: NativePType> VectorMutOps for PVectorMut<T> {
     }
 
     fn append_nulls(&mut self, n: usize) {
-        self.elements.push_n(T::zero(), n);
+        self.elements.push_n(T::zero(), n); // Note that the value we push doesn't actually matter.
         self.validity.append_n(false, n);
     }
 
diff --git a/vortex-vector/src/primitive/iter.rs b/vortex-vector/src/primitive/iter.rs
@@ -116,6 +116,7 @@ impl<T: NativePType> FromIterator<T> for PVectorMut<T> {
 ///
 /// 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 PVectorMutIterator<T: NativePType> {
     /// The vector being iterated over.
     vector: PVectorMut<T>,
diff --git a/vortex-vector/src/private.rs b/vortex-vector/src/private.rs
@@ -28,3 +28,6 @@ impl Sealed for PrimitiveVector {}
 impl Sealed for PrimitiveVectorMut {}
 impl<T: NativePType> Sealed for PVector<T> {}
 impl<T: NativePType> Sealed for PVectorMut<T> {}
+
+impl Sealed for StructVector {}
+impl Sealed for StructVectorMut {}
diff --git a/vortex-vector/src/struct_/mod.rs b/vortex-vector/src/struct_/mod.rs
@@ -0,0 +1,24 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Definition and implementation of [`StructVector`] and [`StructVectorMut`].
+
+mod vector;
+pub use vector::StructVector;
+
+mod vector_mut;
+pub use vector_mut::StructVectorMut;
+
+use crate::{Vector, VectorMut};
+
+impl From<StructVector> for Vector {
+    fn from(v: StructVector) -> Self {
+        Self::Struct(v)
+    }
+}
+
+impl From<StructVectorMut> for VectorMut {
+    fn from(v: StructVectorMut) -> Self {
+        Self::Struct(v)
+    }
+}
diff --git a/vortex-vector/src/struct_/vector.rs b/vortex-vector/src/struct_/vector.rs
@@ -0,0 +1,199 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Definition and implementation of [`StructVector`].
+
+use vortex_error::{VortexExpect, VortexResult, vortex_ensure};
+use vortex_mask::Mask;
+
+use crate::{StructVectorMut, Vector, VectorMutOps, VectorOps};
+
+/// An immutable vector of boolean values.
+///
+/// `StructVector` can be considered a borrowed / frozen version of [`StructVectorMut`], which is
+/// created via the [`freeze`](crate::VectorMutOps::freeze) method.
+///
+/// See the documentation for [`StructVectorMut`] for more information.
+#[derive(Debug, Clone)]
+pub struct StructVector {
+    /// The fields of the `StructVector`, each stored column-wise as a [`Vector`].
+    pub(super) fields: Box<[Vector]>,
+
+    /// The validity mask (where `true` represents an element is **not** null).
+    pub(super) validity: Mask,
+
+    /// The length of the vector (which is the same as all field vectors).
+    ///
+    /// This is stored here as a convenience, and also helps in the case that the `StructVector` has
+    /// no fields.
+    pub(super) len: usize,
+}
+
+impl StructVector {
+    /// Creates a new [`StructVector`] from the given fields and validity mask.
+    ///
+    /// # Panics
+    ///
+    /// Panics if:
+    ///
+    /// - Any field vector has a length that does not match the length of other fields.
+    /// - The validity mask length does not match the field length.
+    pub fn new(fields: Box<[Vector]>, validity: Mask) -> Self {
+        Self::try_new(fields, validity).vortex_expect(
+            "`StructVector` fields must have matching length and validity constraints",
+        )
+    }
+
+    /// Tries to create a new [`StructVector`] from the given fields and validity mask.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if:
+    ///
+    /// - Any field vector has a length that does not match the length of other fields.
+    /// - The validity mask length does not match the field length.
+    pub fn try_new(fields: Box<[Vector]>, validity: Mask) -> VortexResult<Self> {
+        let len = if fields.is_empty() {
+            validity.len()
+        } else {
+            fields[0].len()
+        };
+
+        Self::validate(&fields, len, &validity)?;
+
+        Ok(Self {
+            fields,
+            validity,
+            len,
+        })
+    }
+
+    /// Creates a new [`StructVector`] from the given fields and validity mask without validation.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that:
+    ///
+    /// - All field vectors have the same length.
+    /// - The validity mask has a length equal to the field length.
+    pub unsafe fn new_unchecked(fields: Box<[Vector]>, validity: Mask) -> Self {
+        let len = if fields.is_empty() {
+            validity.len()
+        } else {
+            fields[0].len()
+        };
+
+        debug_assert!(
+            Self::validate(&fields, len, &validity).is_ok(),
+            "`StructVector` fields must have matching length and validity constraints"
+        );
+
+        Self {
+            fields,
+            validity,
+            len,
+        }
+    }
+
+    /// Validates the fields and validity mask for a [`StructVector`].
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if:
+    ///
+    /// - Any field vector has a length that does not match the length of other fields.
+    /// - The validity mask length does not match the field length.
+    fn validate(fields: &[Vector], len: usize, validity: &Mask) -> VortexResult<()> {
+        // Validate that the validity mask has the correct length.
+        vortex_ensure!(
+            validity.len() == len,
+            "Validity mask length ({}) does not match expected length ({})",
+            validity.len(),
+            len
+        );
+
+        // Validate that all fields have the correct length.
+        for (i, field) in fields.iter().enumerate() {
+            vortex_ensure!(
+                field.len() == len,
+                "Field {} has length {} but expected length {}",
+                i,
+                field.len(),
+                len
+            );
+        }
+
+        Ok(())
+    }
+
+    /// Decomposes the struct vector into its constituent parts (fields, validity, and length).
+    pub fn into_parts(self) -> (Box<[Vector]>, Mask, usize) {
+        (self.fields, self.validity, self.len)
+    }
+
+    /// Returns the fields of the `StructVector`, each stored column-wise as a [`Vector`].
+    pub fn fields(&self) -> &[Vector] {
+        self.fields.as_ref()
+    }
+}
+
+impl VectorOps for StructVector {
+    type Mutable = StructVectorMut;
+
+    fn len(&self) -> usize {
+        self.len
+    }
+
+    fn validity(&self) -> &Mask {
+        &self.validity
+    }
+
+    fn try_into_mut(self) -> Result<Self::Mutable, Self>
+    where
+        Self: Sized,
+    {
+        let validity = match self.validity.try_into_mut() {
+            Ok(validity) => validity,
+            Err(validity) => {
+                return Err(StructVector { validity, ..self });
+            }
+        };
+
+        // Convert all of the remaining fields to mutable, if possible.
+        let mut mutable_fields = Vec::with_capacity(self.fields.len());
+        let mut fields_iter = self.fields.into_iter();
+
+        while let Some(field) = fields_iter.next() {
+            match field.try_into_mut() {
+                Ok(mutable_field) => {
+                    // We were able to take ownership of the field vector, so add it and keep going.
+                    mutable_fields.push(mutable_field);
+                }
+                Err(immutable_field) => {
+                    // We were unable to take ownership, so we must re-freeze all of the fields
+                    // vectors we took ownership over and reconstruct the original `StructVector`.
+
+                    let mut all_fields: Vec<Vector> = mutable_fields
+                        .into_iter()
+                        .map(|mut_field| mut_field.freeze())
+                        .collect();
+
+                    all_fields.push(immutable_field);
+                    all_fields.extend(fields_iter);
+
+                    return Err(StructVector {
+                        fields: all_fields.into_boxed_slice(),
+                        len: self.len,
+                        validity: validity.freeze(),
+                    });
+                }
+            }
+        }
+
+        Ok(StructVectorMut {
+            fields: mutable_fields,
+            len: self.len,
+            validity,
+        })
+    }
+}
diff --git a/vortex-vector/src/struct_/vector_mut.rs b/vortex-vector/src/struct_/vector_mut.rs
diff --git a/vortex-vector/src/vector.rs b/vortex-vector/src/vector.rs
diff --git a/vortex-vector/src/vector_mut.rs b/vortex-vector/src/vector_mut.rs

Original file line number	Diff line number	Diff line change
`@@ -63,7 +63,7 @@ impl BoolVector {`
`63`	`63`	`Self { bits, validity }`
`64`	`64`	`}`
`65`	`65`
`66`		`- /// Decomposes the boolean vector into its constituent parts.`
	`66`	`+ /// Decomposes the boolean vector into its constituent parts (bit buffer and validity).`
`67`	`67`	`pub fn into_parts(self) -> (BitBuffer, Mask) {`
`68`	`68`	`(self.bits, self.validity)`
`69`	`69`	`}`
Original file line number	Diff line number	Diff line change
`@@ -121,7 +121,7 @@ impl BoolVectorMut {`
`121`	`121`	`}`
`122`	`122`	`}`
`123`	`123`
`124`		`- /// Returns the parts of the mutable vector.`
	`124`	`+ /// Decomposes the boolean vector into its constituent parts (bit buffer and validity).`
`125`	`125`	`pub fn into_parts(self) -> (BitBufferMut, MaskMut) {`
`126`	`126`	`(self.bits, self.validity)`
`127`	`127`	`}`
`@@ -151,7 +151,7 @@ impl VectorMutOps for BoolVectorMut {`
`151`	`151`	`}`
`152`	`152`
`153`	`153`	`fn append_nulls(&mut self, n: usize) {`
`154`		`- self.bits.append_n(false, n);`
	`154`	`+ self.bits.append_n(false, n); // Note that the value we push doesn't actually matter.`
`155`	`155`	`self.validity.append_n(false, n);`
`156`	`156`	`}`
`157`	`157`
Original file line number	Diff line number	Diff line change
`@@ -48,6 +48,10 @@ macro_rules! match_each_vector {`
`48`	`48`	`let $vec = v;`
`49`	`49`	`$body`
`50`	`50`	`}`
	`51`	`+ $crate::Vector::Struct(v) => {`
	`52`	`+ let $vec = v;`
	`53`	`+ $body`
	`54`	`+ }`
`51`	`55`	`}`
`52`	`56`	`}};`
`53`	`57`	`}`
`@@ -98,6 +102,10 @@ macro_rules! match_each_vector_mut {`
`98`	`102`	`let $vec = v;`
`99`	`103`	`$body`
`100`	`104`	`}`
	`105`	`+ $crate::VectorMut::Struct(v) => {`
	`106`	`+ let $vec = v;`
	`107`	`+ $body`
	`108`	`+ }`
`101`	`109`	`}`
`102`	`110`	`}};`
`103`	`111`	`}`
Original file line number	Diff line number	Diff line change
`@@ -69,7 +69,7 @@ impl<T: NativePType> PVector<T> {`
`69`	`69`	`Self { elements, validity }`
`70`	`70`	`}`
`71`	`71`
`72`		`- /// Decomposes the primitive vector into its constituent parts.`
	`72`	`+ /// Decomposes the primitive vector into its constituent parts (buffer and validity).`
`73`	`73`	`pub fn into_parts(self) -> (Buffer<T>, Mask) {`
`74`	`74`	`(self.elements, self.validity)`
`75`	`75`	`}`
Original file line number	Diff line number	Diff line change
`@@ -160,7 +160,7 @@ impl<T> PVectorMut<T> {`
`160`	`160`	`}`
`161`	`161`	`}`
`162`	`162`
`163`		`- /// Decomposes the primitive vector into its constituent parts.`
	`163`	`+ /// Decomposes the primitive vector into its constituent parts (buffer and validity).`
`164`	`164`	`pub fn into_parts(self) -> (BufferMut<T>, MaskMut) {`
`165`	`165`	`(self.elements, self.validity)`
`166`	`166`	`}`
`@@ -189,7 +189,7 @@ impl<T: NativePType> VectorMutOps for PVectorMut<T> {`
`189`	`189`	`}`
`190`	`190`
`191`	`191`	`fn append_nulls(&mut self, n: usize) {`
`192`		`- self.elements.push_n(T::zero(), n);`
	`192`	`+ self.elements.push_n(T::zero(), n); // Note that the value we push doesn't actually matter.`
`193`	`193`	`self.validity.append_n(false, n);`
`194`	`194`	`}`
`195`	`195`