clean up struct vector

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

Author: connortsui20

vortex-vector/src/struct_/vector.rs

@@ -137,6 +137,7 @@ impl VectorOps for StructVector {
         Self: Sized,
     {
         let len = self.len;
+
         let fields = match Arc::try_unwrap(self.fields) {
             Ok(fields) => fields,
             Err(fields) => return Err(StructVector { fields, ..self }),

vortex-vector/src/struct_/vector_mut.rs

@@ -5,6 +5,7 @@
 
 use std::sync::Arc;
 
+use vortex_dtype::StructFields;
 use vortex_error::{VortexExpect, VortexResult, vortex_ensure};
 use vortex_mask::MaskMut;
 
@@ -174,6 +175,31 @@ impl StructVectorMut {
         }
     }
 
+    /// Creates a new [`StructVectorMut`] with the given fields and capacity.
+    pub fn with_capacity(struct_fields: &StructFields, capacity: usize) -> Self {
+        let fields: Vec<VectorMut> = struct_fields
+            .fields()
+            .map(|dtype| VectorMut::with_capacity(&dtype, capacity))
+            .collect();
+
+        let validity = MaskMut::with_capacity(capacity);
+        let len = validity.len();
+
+        #[cfg(debug_assertions)]
+        {
+            for field in &fields {
+                debug_assert_eq!(field.len(), 0);
+            }
+            debug_assert_eq!(validity.len(), 0);
+        }
+
+        Self {
+            fields: fields.into_boxed_slice(),
+            validity,
+            len,
+        }
+    }
+
     /// Decomposes the struct vector into its constituent parts (fields, validity, and length).
     pub fn into_parts(self) -> (Box<[VectorMut]>, MaskMut, usize) {
         (self.fields, self.validity, self.len)
@@ -662,7 +688,7 @@ mod tests {
         );
 
         // Create a VectorMut with capacity using the struct dtype.
-        let vector_mut = VectorMut::with_capacity(100, &struct_dtype);
+        let vector_mut = VectorMut::with_capacity(&struct_dtype, 100);
 
         // Verify it's a struct vector.
         match vector_mut {

vortex-vector/src/vector_mut.rs

@@ -8,7 +8,6 @@
 
 use vortex_dtype::DType;
 use vortex_error::vortex_panic;
-use vortex_mask::MaskMut;
 
 use crate::varbin::{BinaryVectorMut, StringVectorMut};
 use crate::{
@@ -46,34 +45,22 @@ pub enum VectorMut {
 
 impl VectorMut {
     /// Create a new mutable vector with the given capacity and dtype.
-    pub fn with_capacity(capacity: usize, dtype: &DType) -> Self {
+    pub fn with_capacity(dtype: &DType, capacity: usize) -> Self {
         match dtype {
-            DType::Null => NullVectorMut::new(0).into(), // `NullVector` has `usize::MAX` capacity.
+            DType::Null => NullVectorMut::new(0).into(),
             DType::Bool(_) => BoolVectorMut::with_capacity(capacity).into(),
             DType::Primitive(ptype, _) => {
                 PrimitiveVectorMut::with_capacity(*ptype, capacity).into()
             }
             DType::Struct(struct_fields, _) => {
-                let fields: Vec<VectorMut> = struct_fields
-                    .fields()
-                    .map(|dtype| Self::with_capacity(capacity, &dtype))
-                    .collect();
-                let validity = MaskMut::with_capacity(capacity);
-
-                #[cfg(debug_assertions)]
-                {
-                    for field in &fields {
-                        debug_assert_eq!(field.len(), 0);
-                    }
-                    debug_assert_eq!(validity.len(), 0);
-                }
-
-                // SAFETY: All fields and validity have length 0, so they all have the same length.
-                Self::Struct(unsafe {
-                    StructVectorMut::new_unchecked(fields.into_boxed_slice(), validity)
-                })
+                StructVectorMut::with_capacity(struct_fields, capacity).into()
             }
-            _ => vortex_panic!("Unsupported dtype for VectorMut"),
+            DType::Decimal(..)
+            | DType::Utf8(_)
+            | DType::Binary(_)
+            | DType::List(..)
+            | DType::FixedSizeList(..)
+            | DType::Extension(_) => vortex_panic!("Unsupported dtype for VectorMut"),
         }
     }
 }
@@ -226,24 +213,24 @@ mod tests {
     #[test]
     fn test_with_capacity() {
         // Test capacity allocation for different types.
-        let null_vec = VectorMut::with_capacity(10, &DType::Null);
+        let null_vec = VectorMut::with_capacity(&DType::Null, 10);
         assert_eq!(null_vec.capacity(), usize::MAX); // Null vectors have unlimited capacity.
 
-        let bool_vec = VectorMut::with_capacity(100, &DType::Bool(Nullability::Nullable));
+        let bool_vec = VectorMut::with_capacity(&DType::Bool(Nullability::Nullable), 100);
         assert!(bool_vec.capacity() >= 100);
 
         let prim_vec =
-            VectorMut::with_capacity(50, &DType::Primitive(PType::I32, Nullability::Nullable));
+            VectorMut::with_capacity(&DType::Primitive(PType::I32, Nullability::Nullable), 50);
         assert!(prim_vec.capacity() >= 50);
     }
 
     #[test]
     #[should_panic(expected = "Mismatched vector types")]
     fn test_type_mismatch_panics() {
         // Test that operations between mismatched types panic.
-        let mut vec1 = VectorMut::with_capacity(10, &DType::Bool(Nullability::Nullable));
+        let mut vec1 = VectorMut::with_capacity(&DType::Bool(Nullability::Nullable), 10);
         let vec2 =
-            VectorMut::with_capacity(10, &DType::Primitive(PType::I32, Nullability::Nullable));
+            VectorMut::with_capacity(&DType::Primitive(PType::I32, Nullability::Nullable), 10);
 
         vec1.unsplit(vec2); // Should panic.
     }