add StructVectorMut ops implementation

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

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/struct_/vector.rs

@@ -3,6 +3,7 @@
 
 //! Definition and implementation of [`StructVector`].
 
+use vortex_error::{VortexExpect, VortexResult, vortex_ensure};
 use vortex_mask::Mask;
 
 use crate::{StructVectorMut, Vector, VectorMutOps, VectorOps};
@@ -18,24 +19,113 @@ 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,
+}
 
-    /// The capacity of the vector (which is the less than or equal to the capacity of all field
-    /// vectors).
+impl StructVector {
+    /// Creates a new [`StructVector`] from the given fields and validity mask.
+    ///
+    /// # Panics
+    ///
+    /// Panics if:
     ///
-    /// This is stored here as a convenience for converting to/from a [`StructVectorMut`], and also
-    /// helps in the case that the `StructVector` has no fields.
-    pub(super) capacity: usize,
+    /// - 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",
+        )
+    }
 
-    /// The validity mask (where `true` represents an element is **not** null).
-    pub(super) validity: Mask,
-}
+    /// 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(())
+    }
 
-impl StructVector {
     /// Returns the fields of the `StructVector`, each stored column-wise as a [`Vector`].
     pub fn fields(&self) -> &[Vector] {
         self.fields.as_ref()
@@ -89,7 +179,6 @@ impl VectorOps for StructVector {
                     return Err(StructVector {
                         fields: all_fields.into_boxed_slice(),
                         len: self.len,
-                        capacity: self.capacity,
                         validity: validity.freeze(),
                     });
                 }
@@ -99,220 +188,7 @@ impl VectorOps for StructVector {
         Ok(StructVectorMut {
             fields: mutable_fields,
             len: self.len,
-            capacity: self.capacity,
             validity,
         })
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use vortex_mask::Mask;
-
-    use super::*;
-    use crate::{BoolVectorMut, NullVector, PVectorMut, VectorMut};
-
-    // TODO(connor): Make sure to test actual logic instead of just lengths and capacity.
-
-    /// Helper function to create a `StructVector` with the given fields.
-    fn create_struct_vector(fields: Vec<Vector>, len: usize, capacity: usize) -> StructVector {
-        StructVector {
-            fields: fields.into_boxed_slice(),
-            len,
-            capacity,
-            validity: Mask::AllTrue(len),
-        }
-    }
-
-    #[test]
-    fn test_try_into_mut_success() {
-        // Create a `StructVector` with 3 different field types (null, bool, primitive).
-        let null_field: Vector = NullVector::new(5).into();
-        let bool_field: Vector = BoolVectorMut::from_iter([true, false, true, false, true])
-            .freeze()
-            .into();
-        let prim_field: Vector = PVectorMut::<i32>::from_iter([10, 20, 30, 40, 50])
-            .freeze()
-            .into();
-
-        let struct_vec = create_struct_vector(vec![null_field, bool_field, prim_field], 5, 5);
-
-        // Attempt to convert to mutable.
-        let result = struct_vec.try_into_mut();
-
-        // Should succeed since all fields have unique ownership.
-        assert!(result.is_ok());
-
-        let mut_struct = result.unwrap();
-
-        // Verify the mutable struct has correct length and capacity.
-        assert_eq!(mut_struct.len(), 5);
-        assert_eq!(mut_struct.capacity(), 5);
-
-        // Verify that we have 3 fields.
-        assert_eq!(mut_struct.fields.len(), 3);
-
-        // Verify each field has the correct type and length.
-        assert!(matches!(mut_struct.fields[0], VectorMut::Null(_)));
-        assert_eq!(mut_struct.fields[0].len(), 5);
-
-        assert!(matches!(mut_struct.fields[1], VectorMut::Bool(_)));
-        assert_eq!(mut_struct.fields[1].len(), 5);
-        assert!(mut_struct.fields[1].capacity() >= 5);
-
-        assert!(matches!(mut_struct.fields[2], VectorMut::Primitive(_)));
-        assert_eq!(mut_struct.fields[2].len(), 5);
-        assert!(mut_struct.fields[2].capacity() >= 5);
-    }
-
-    #[test]
-    fn test_try_into_mut_fails_first_field() {
-        // Create a bool field with shared ownership (cloned to have Arc count > 1).
-        let bool_field: Vector = BoolVectorMut::from_iter([true, false, true])
-            .freeze()
-            .into();
-        let bool_field_clone = bool_field.clone();
-
-        let null_field: Vector = NullVector::new(3).into();
-        let prim_field: Vector = PVectorMut::<i32>::from_iter([1, 2, 3]).freeze().into();
-
-        // Create a struct with the cloned bool field as the first field.
-        let struct_vec = create_struct_vector(vec![bool_field_clone, null_field, prim_field], 3, 3);
-
-        // Attempt to convert to mutable.
-        let result = struct_vec.try_into_mut();
-
-        // Should fail since the first field has shared ownership.
-        assert!(result.is_err());
-
-        // Keep the original alive to ensure shared ownership is maintained.
-        drop(bool_field);
-
-        let recovered_struct = result.unwrap_err();
-
-        // Verify the recovered struct has all 3 fields.
-        assert_eq!(recovered_struct.fields.len(), 3);
-        assert_eq!(recovered_struct.len(), 3);
-        assert_eq!(recovered_struct.capacity, 3);
-
-        // Verify field types are preserved.
-        assert!(matches!(recovered_struct.fields[0], Vector::Bool(_)));
-        assert!(matches!(recovered_struct.fields[1], Vector::Null(_)));
-        assert!(matches!(recovered_struct.fields[2], Vector::Primitive(_)));
-    }
-
-    #[test]
-    fn test_try_into_mut_fails_middle_field() {
-        // Create a primitive field with shared ownership (cloned).
-        let prim_field: Vector = PVectorMut::<i32>::from_iter([100, 200, 300, 400])
-            .freeze()
-            .into();
-        let prim_field_clone = prim_field.clone();
-
-        let bool_field: Vector = BoolVectorMut::from_iter([true, false, true, false])
-            .freeze()
-            .into();
-        let null_field: Vector = NullVector::new(4).into();
-
-        // Create a struct with the cloned primitive field as the middle field.
-        let struct_vec = create_struct_vector(vec![bool_field, prim_field_clone, null_field], 4, 4);
-
-        // Attempt to convert to mutable.
-        let result = struct_vec.try_into_mut();
-
-        // Should fail since the middle field has shared ownership.
-        assert!(result.is_err());
-
-        // Keep the original alive to ensure shared ownership is maintained.
-        drop(prim_field);
-
-        let recovered_struct = result.unwrap_err();
-
-        // Verify all 3 fields are present in the recovered struct.
-        assert_eq!(recovered_struct.fields.len(), 3);
-        assert_eq!(recovered_struct.len(), 4);
-        assert_eq!(recovered_struct.capacity, 4);
-
-        // Verify field types and order are preserved.
-        // The first field was converted to mutable then frozen back, so it should still be bool.
-        assert!(matches!(recovered_struct.fields[0], Vector::Bool(_)));
-        assert!(matches!(recovered_struct.fields[1], Vector::Primitive(_)));
-        assert!(matches!(recovered_struct.fields[2], Vector::Null(_)));
-    }
-
-    #[test]
-    fn test_try_into_mut_fails_last_field() {
-        // Create a null field with "shared ownership" (though NullVector always succeeds, we test
-        // the pattern by cloning a bool vector).
-        let bool_field_last: Vector = BoolVectorMut::from_iter([true, true]).freeze().into();
-        let bool_field_last_clone = bool_field_last.clone();
-
-        let null_field: Vector = NullVector::new(2).into();
-        let prim_field: Vector = PVectorMut::<u64>::from_iter([1000, 2000]).freeze().into();
-
-        // Create a struct with the cloned bool field as the last field.
-        let struct_vec =
-            create_struct_vector(vec![null_field, prim_field, bool_field_last_clone], 2, 2);
-
-        // Attempt to convert to mutable.
-        let result = struct_vec.try_into_mut();
-
-        // Should fail since the last field has shared ownership.
-        assert!(result.is_err());
-
-        // Keep the original alive to ensure shared ownership is maintained.
-        drop(bool_field_last);
-
-        let recovered_struct = result.unwrap_err();
-
-        // Verify all 3 fields are present.
-        assert_eq!(recovered_struct.fields.len(), 3);
-        assert_eq!(recovered_struct.len(), 2);
-        assert_eq!(recovered_struct.capacity, 2);
-
-        // Verify field types are preserved.
-        // The first two fields were converted to mutable then frozen back.
-        assert!(matches!(recovered_struct.fields[0], Vector::Null(_)));
-        assert!(matches!(recovered_struct.fields[1], Vector::Primitive(_)));
-        assert!(matches!(recovered_struct.fields[2], Vector::Bool(_)));
-    }
-
-    #[test]
-    fn test_try_into_mut_nested_struct() {
-        // Create a nested struct: a `StructVector` containing other `StructVector`s as fields.
-        let inner_struct1: Vector = create_struct_vector(
-            vec![
-                NullVector::new(3).into(),
-                BoolVectorMut::from_iter([true, false, true])
-                    .freeze()
-                    .into(),
-            ],
-            3,
-            3,
-        )
-        .into();
-
-        let inner_struct2: Vector = create_struct_vector(
-            vec![PVectorMut::<i32>::from_iter([1, 2, 3]).freeze().into()],
-            3,
-            3,
-        )
-        .into();
-
-        let outer_struct = create_struct_vector(vec![inner_struct1, inner_struct2], 3, 3);
-
-        // Attempt to convert to mutable.
-        let result = outer_struct.try_into_mut();
-
-        // Should succeed once `StructVectorMut` is fully implemented.
-        assert!(result.is_ok());
-
-        let mut_struct = result.unwrap();
-        assert_eq!(mut_struct.len(), 3);
-        assert_eq!(mut_struct.fields.len(), 2);
-
-        // Verify nested structs are converted correctly.
-        assert!(matches!(mut_struct.fields[0], VectorMut::Struct(_)));
-        assert!(matches!(mut_struct.fields[1], VectorMut::Struct(_)));
-    }
-}