diff --git a/vortex-vector/src/struct_/vector.rs b/vortex-vector/src/struct_/vector.rs index b69b3daef1b..6290bce1e3d 100644 --- a/vortex-vector/src/struct_/vector.rs +++ b/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 }), diff --git a/vortex-vector/src/struct_/vector_mut.rs b/vortex-vector/src/struct_/vector_mut.rs index f4c64b2c9a2..e7a868aa5b0 100644 --- a/vortex-vector/src/struct_/vector_mut.rs +++ b/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,23 @@ 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 = struct_fields + .fields() + .map(|dtype| VectorMut::with_capacity(&dtype, capacity)) + .collect(); + + let validity = MaskMut::with_capacity(capacity); + let len = validity.len(); + + 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 +680,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 { diff --git a/vortex-vector/src/vector_mut.rs b/vortex-vector/src/vector_mut.rs index 22ec47ba583..6dbe0fa5519 100644 --- a/vortex-vector/src/vector_mut.rs +++ b/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 = 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,14 +213,14 @@ 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); } @@ -241,9 +228,9 @@ mod tests { #[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. }