diff --git a/encodings/sparse/src/canonical.rs b/encodings/sparse/src/canonical.rs index 47a06ceb7e0..a003f0a85ab 100644 --- a/encodings/sparse/src/canonical.rs +++ b/encodings/sparse/src/canonical.rs @@ -475,7 +475,6 @@ fn execute_sparse_struct( ), }; let patch_values_as_struct = unresolved_patches.values().as_::().into_owned(); - let columns_patch_values = patch_values_as_struct.unmasked_fields(); let names = patch_values_as_struct.names(); let validity = top_level_fill_validity.patch( len, @@ -494,8 +493,8 @@ fn execute_sparse_struct( Ok(StructArray::try_from_iter_with_validity( names.iter().zip_eq( - columns_patch_values - .iter() + patch_values_as_struct + .iter_unmasked_fields() .cloned() .zip_eq(fill_values) .map(|(patch_values, fill_value)| unsafe { diff --git a/fuzz/src/array/mask.rs b/fuzz/src/array/mask.rs index d5a125f26a2..64013f336ed 100644 --- a/fuzz/src/array/mask.rs +++ b/fuzz/src/array/mask.rs @@ -131,7 +131,7 @@ pub fn mask_canonical_array( Canonical::Struct(array) => { let new_validity = mask_validity(&array.validity()?, mask, ctx); StructArray::try_new_with_dtype( - array.unmasked_fields(), + array.iter_unmasked_fields().cloned(), array.struct_fields().clone(), array.len(), new_validity, diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/struct_.rs b/vortex-array/src/aggregate_fn/fns/is_constant/struct_.rs index 9376115521f..84c829dfaca 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/struct_.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/struct_.rs @@ -10,7 +10,7 @@ use crate::arrays::struct_::StructArrayExt; /// Check if a struct array is constant by checking each field independently. pub(super) fn check_struct_constant(s: &StructArray, ctx: &mut ExecutionCtx) -> VortexResult { - for field in s.unmasked_fields().iter() { + for field in s.iter_unmasked_fields() { if !is_constant(field, ctx)? { return Ok(false); } diff --git a/vortex-array/src/arrays/chunked/tests.rs b/vortex-array/src/arrays/chunked/tests.rs index dc83b73aa80..d5a993cf04d 100644 --- a/vortex-array/src/arrays/chunked/tests.rs +++ b/vortex-array/src/arrays/chunked/tests.rs @@ -376,10 +376,12 @@ pub fn pack_nested_structs() -> VortexResult<()> { )? .into_array(); let canonical_struct = chunked.execute::(&mut ctx)?; - let canonical_varbin = canonical_struct.unmasked_fields()[0] + let canonical_varbin = canonical_struct + .unmasked_field(0) .clone() .execute::(&mut ctx)?; - let original_varbin = struct_array.unmasked_fields()[0] + let original_varbin = struct_array + .unmasked_field(0) .clone() .execute::(&mut ctx)?; assert_arrays_eq!(original_varbin, canonical_varbin, &mut ctx); diff --git a/vortex-array/src/arrays/listview/conversion.rs b/vortex-array/src/arrays/listview/conversion.rs index 054f3efecb9..f6b30b830c7 100644 --- a/vortex-array/src/arrays/listview/conversion.rs +++ b/vortex-array/src/arrays/listview/conversion.rs @@ -250,11 +250,11 @@ pub fn recursive_list_from_list_view( } } Canonical::Struct(struct_array) => { - let fields = struct_array.unmasked_fields(); - let mut converted_fields = Vec::with_capacity(fields.len()); + let mut converted_fields = + Vec::with_capacity(struct_array.iter_unmasked_fields().len()); let mut any_changed = false; - for field in fields.iter() { + for field in struct_array.iter_unmasked_fields() { let converted_field = recursive_list_from_list_view(field.clone(), ctx)?; // Avoid cloning if elements didn't change. any_changed |= !ArrayRef::ptr_eq(&converted_field, field); diff --git a/vortex-array/src/arrays/masked/execute.rs b/vortex-array/src/arrays/masked/execute.rs index cfeb29696bc..a8d17f15e2d 100644 --- a/vortex-array/src/arrays/masked/execute.rs +++ b/vortex-array/src/arrays/masked/execute.rs @@ -143,10 +143,15 @@ fn mask_validity_fixed_size_list( fn mask_validity_struct(array: StructArray, validity: Validity) -> VortexResult { let len = array.len(); let new_validity = Validity::and(array.validity()?, validity)?; - let fields = array.unmasked_fields(); - let struct_fields = array.struct_fields(); // SAFETY: We're only changing validity, not the data structure. - Ok(unsafe { StructArray::new_unchecked(fields, struct_fields.clone(), len, new_validity) }) + Ok(unsafe { + StructArray::new_unchecked( + array.iter_unmasked_fields().cloned(), + array.struct_fields().clone(), + len, + new_validity, + ) + }) } fn mask_validity_union(array: UnionArray, validity: Validity) -> VortexResult { @@ -155,10 +160,9 @@ fn mask_validity_union(array: UnionArray, validity: Validity) -> VortexResult ArraySlots { - // `fields` is borrowed, so its arrays must be cloned regardless; clone them straight into the - // `SmallVec` rather than through an intermediate `Vec`. - let mut slots = ArraySlots::with_capacity(StructSlots::FIELDS_OFFSET + fields.len()); + let mut slots = ArraySlots::with_capacity(StructSlots::FIELDS_OFFSET + nfields); slots.push(validity_to_child(validity, length)); - slots.extend(fields.iter().cloned().map(Some)); slots } -pub trait StructArrayExt: TypedArrayRef { +pub(super) fn make_struct_slots( + fields: impl IntoIterator, + validity: &Validity, + length: usize, +) -> ArraySlots { + // Take the fields by value so callers that already own them move their `ArrayRef`s straight + // into the `SmallVec` instead of paying a refcount bump per field. + let fields = fields.into_iter(); + let mut slots = struct_slots_with_capacity(validity, length, fields.size_hint().0); + slots.extend(fields.map(Some)); + slots +} + +/// Struct-specific accessors. +/// +/// Slot accessors (`validity`, `fields`, `slots_view`) live on the generated +/// [`StructArraySlotsExt`] supertrait; this trait layers struct-specific lookups on top. +pub trait StructArrayExt: StructArraySlotsExt { fn nullability(&self) -> crate::dtype::Nullability { match self.as_ref().dtype() { DType::Struct(_, nullability) => *nullability, @@ -195,25 +214,29 @@ pub trait StructArrayExt: TypedArrayRef { } fn struct_validity(&self) -> Validity { - child_to_validity( - self.as_ref().slots()[StructSlots::VALIDITY].as_ref(), - self.nullability(), - ) + child_to_validity(self.validity(), self.nullability()) } - fn iter_unmasked_fields(&self) -> impl Iterator + '_ { - self.as_ref().slots()[StructSlots::FIELDS_OFFSET..] - .iter() - .map(|s| s.as_ref().vortex_expect("StructArray field slot")) + /// Iterate over the field arrays in declaration order. + fn iter_unmasked_fields(&self) -> impl ExactSizeIterator + '_ { + self.fields().iter() } - fn unmasked_fields(&self) -> Vec { - self.iter_unmasked_fields().cloned().collect() + /// The field array at `idx`, or `None` if `idx` is out of bounds. + /// + /// Use this over [`unmasked_field`](Self::unmasked_field) when the index comes from outside + /// the library and an out-of-bounds value is an error to report rather than a bug to panic on. + fn unmasked_field_opt(&self, idx: usize) -> Option<&ArrayRef> { + self.fields().get(idx) } + /// The field array at `idx`. + /// + /// # Panics + /// + /// If `idx` is out of bounds. fn unmasked_field(&self, idx: usize) -> &ArrayRef { - self.as_ref().slots()[StructSlots::FIELDS_OFFSET + idx] - .as_ref() + self.unmasked_field_opt(idx) .vortex_expect("StructArray field slot") } @@ -292,11 +315,7 @@ impl Array { validity: Validity, ) -> Self { let outer_dtype = DType::Struct(dtype, validity.nullability()); - let fields = fields.into_iter(); - let (lower, _) = fields.size_hint(); - let mut slots = ArraySlots::with_capacity(StructSlots::FIELDS_OFFSET + lower); - slots.push(validity_to_child(&validity, length)); - slots.extend(fields.map(Some)); + let slots = make_struct_slots(fields, &validity, length); unsafe { Array::from_parts_unchecked( ArrayParts::new(Struct, outer_dtype, length, EmptyArrayData).with_slots(slots), @@ -312,11 +331,7 @@ impl Array { validity: Validity, ) -> VortexResult { let outer_dtype = DType::Struct(dtype, validity.nullability()); - let fields = fields.into_iter(); - let (lower, _) = fields.size_hint(); - let mut slots = ArraySlots::with_capacity(StructSlots::FIELDS_OFFSET + lower); - slots.push(validity_to_child(&validity, length)); - slots.extend(fields.map(Some)); + let slots = make_struct_slots(fields, &validity, length); Array::try_from_parts( ArrayParts::new(Struct, outer_dtype, length, EmptyArrayData).with_slots(slots), ) @@ -404,7 +419,7 @@ impl Array { StructFields::new(FieldNames::default(), Vec::new()), crate::dtype::Nullability::NonNullable, ); - let slots = make_struct_slots(&[], &Validity::NonNullable, len); + let slots = make_struct_slots([], &Validity::NonNullable, len); unsafe { Array::from_parts_unchecked( ArrayParts::new(Struct, dtype, len, EmptyArrayData).with_slots(slots), @@ -414,10 +429,7 @@ impl Array { // TODO(ngates): remove this... it doesn't help to consume self. pub fn into_data_parts(self) -> StructDataParts { - let fields: Vec = self.slots()[StructSlots::FIELDS_OFFSET..] - .iter() - .map(|s| s.as_ref().vortex_expect("StructArray field slot").clone()) - .collect(); + let fields = self.fields().to_vec(); let validity = self.validity().vortex_expect("StructArray validity"); StructDataParts { struct_fields: self.struct_fields().clone(), @@ -434,17 +446,13 @@ impl Array { let position = struct_dtype.find(name.as_ref())?; let slot_position = StructSlots::FIELDS_OFFSET + position; - let field = self.slots()[slot_position] - .as_ref() - .vortex_expect("StructArray field slot") - .clone(); - let new_slots: ArraySlots = self - .slots() - .iter() - .enumerate() - .filter(|(i, _)| *i != slot_position) - .map(|(_, s)| s.clone()) - .collect(); + let field = self.unmasked_field(position).clone(); + // `Filter` has a zero lower bound, so build the slots with an exact capacity instead of + // letting `collect` grow them. + let slots = self.slots(); + let mut new_slots = ArraySlots::with_capacity(slots.len() - 1); + new_slots.extend(slots[..slot_position].iter().cloned()); + new_slots.extend(slots[slot_position + 1..].iter().cloned()); let new_dtype = struct_dtype.without_field(position).ok()?; let new_array = unsafe { @@ -469,11 +477,7 @@ impl Array { let types = struct_dtype.fields().chain(once(array.dtype().clone())); let new_fields = StructFields::new(names.collect(), types.collect()); - let children: Vec = self.slots()[StructSlots::FIELDS_OFFSET..] - .iter() - .map(|s| s.as_ref().vortex_expect("StructArray field slot").clone()) - .chain(once(array)) - .collect(); + let children = self.iter_unmasked_fields().cloned().chain(once(array)); Self::try_new_with_dtype(children, new_fields, self.len(), self.validity()?) } @@ -559,11 +563,12 @@ impl Array { struct_validity.clone() }; - // Nothing to push down. + // Nothing to push down. The fields are unchanged, so reuse the existing `StructFields` + // rather than re-deriving it field by field. if struct_validity.definitely_no_nulls() { - return Self::try_new( - self.names().clone(), - self.unmasked_fields(), + return Self::try_new_with_dtype( + self.iter_unmasked_fields().cloned(), + self.struct_fields().clone(), self.len(), new_validity, ); diff --git a/vortex-array/src/arrays/struct_/compute/mask.rs b/vortex-array/src/arrays/struct_/compute/mask.rs index 0f428284454..1b5b526e69d 100644 --- a/vortex-array/src/arrays/struct_/compute/mask.rs +++ b/vortex-array/src/arrays/struct_/compute/mask.rs @@ -15,7 +15,7 @@ use crate::validity::Validity; impl MaskReduce for Struct { fn mask(array: ArrayView<'_, Struct>, mask: &ArrayRef) -> VortexResult> { StructArray::try_new_with_dtype( - array.unmasked_fields(), + array.iter_unmasked_fields().cloned(), array.struct_fields().clone(), array.len(), array.validity()?.and(Validity::Array(mask.clone()))?, diff --git a/vortex-array/src/arrays/struct_/compute/rules.rs b/vortex-array/src/arrays/struct_/compute/rules.rs index 01981803239..da7fbedb67a 100644 --- a/vortex-array/src/arrays/struct_/compute/rules.rs +++ b/vortex-array/src/arrays/struct_/compute/rules.rs @@ -333,7 +333,7 @@ mod tests { .unwrap() .execute::(&mut SESSION.create_execution_ctx()) .unwrap(); - assert_eq!(result.unmasked_fields().len(), 2); + assert_eq!(result.iter_unmasked_fields().len(), 2); assert_arrays_eq!( result.unmasked_field_by_name("a").unwrap(), buffer![1i32, 2, 3].into_array(), diff --git a/vortex-array/src/arrays/struct_/compute/take.rs b/vortex-array/src/arrays/struct_/compute/take.rs index a63c6c85cf6..b98cbcd2be8 100644 --- a/vortex-array/src/arrays/struct_/compute/take.rs +++ b/vortex-array/src/arrays/struct_/compute/take.rs @@ -20,7 +20,7 @@ impl TakeReduce for Struct { // an out of bounds element. if array.is_empty() { return StructArray::try_new_with_dtype( - array.iter_unmasked_fields().cloned().collect::>(), + array.iter_unmasked_fields().cloned(), array.struct_fields().clone(), indices.len(), Validity::AllInvalid, diff --git a/vortex-array/src/arrays/struct_/mod.rs b/vortex-array/src/arrays/struct_/mod.rs index cdfc632083d..ab533dd8cdb 100644 --- a/vortex-array/src/arrays/struct_/mod.rs +++ b/vortex-array/src/arrays/struct_/mod.rs @@ -3,6 +3,7 @@ mod array; pub use array::StructArrayExt; +pub use array::StructArraySlotsExt; pub use array::StructDataParts; pub use array::StructSlots; pub use array::StructSlotsView; diff --git a/vortex-array/src/arrays/struct_/vtable/mod.rs b/vortex-array/src/arrays/struct_/vtable/mod.rs index 71a49f38ebe..c1639501288 100644 --- a/vortex-array/src/arrays/struct_/vtable/mod.rs +++ b/vortex-array/src/arrays/struct_/vtable/mod.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use itertools::Itertools; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -19,7 +18,7 @@ use crate::array::VTable; use crate::array::child_to_validity; use crate::array::with_empty_buffers; use crate::arrays::struct_::array::StructSlots; -use crate::arrays::struct_::array::make_struct_slots; +use crate::arrays::struct_::array::struct_slots_with_capacity; use crate::arrays::struct_::compute::rules::PARENT_RULES; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; @@ -173,16 +172,18 @@ impl VTable for Struct { ); }; - let field_children: Vec<_> = (0..struct_dtype.nfields()) - .map(|i| { - let child_dtype = struct_dtype - .field_by_index(i) - .vortex_expect("no out of bounds"); - children.get(non_data_children + i, &child_dtype, len) - }) - .try_collect()?; + let mut slots = struct_slots_with_capacity(&validity, len, struct_dtype.nfields()); + for i in 0..struct_dtype.nfields() { + let child_dtype = struct_dtype + .field_by_index(i) + .vortex_expect("no out of bounds"); + slots.push(Some(children.get( + non_data_children + i, + &child_dtype, + len, + )?)); + } - let slots = make_struct_slots(&field_children, &validity, len); Ok(ArrayParts::new(self.clone(), dtype.clone(), len, EmptyArrayData).with_slots(slots)) } diff --git a/vortex-array/src/arrays/union/compute/mask.rs b/vortex-array/src/arrays/union/compute/mask.rs index 03ee226fedf..63bb9ae60f7 100644 --- a/vortex-array/src/arrays/union/compute/mask.rs +++ b/vortex-array/src/arrays/union/compute/mask.rs @@ -18,7 +18,7 @@ impl MaskReduce for Union { UnionArray::try_new( array.type_ids().clone().mask(mask.clone())?, array.variants().clone(), - array.children().to_vec(), + array.iter_children().cloned(), ) .map(|a| Some(a.into_array())) } diff --git a/vortex-ffi/src/array.rs b/vortex-ffi/src/array.rs index 7f8bbaeae56..ae2a2e3d93a 100644 --- a/vortex-ffi/src/array.rs +++ b/vortex-ffi/src/array.rs @@ -234,11 +234,9 @@ pub unsafe extern "C-unwind" fn vx_array_get_field( let array = vx_array::as_ref(array); let mut ctx = legacy_session().create_execution_ctx(); - let field_array = array - .clone() - .execute::(&mut ctx)? - .unmasked_fields() - .get(index) + let struct_array = array.clone().execute::(&mut ctx)?; + let field_array = struct_array + .unmasked_field_opt(index) .ok_or_else(|| vortex_err!("Field index out of bounds"))? .clone(); diff --git a/vortex-tui/src/browse/ui/layouts.rs b/vortex-tui/src/browse/ui/layouts.rs index 93468642db8..bace4ef2c0e 100644 --- a/vortex-tui/src/browse/ui/layouts.rs +++ b/vortex-tui/src/browse/ui/layouts.rs @@ -176,8 +176,7 @@ fn render_array(app: &AppState, area: Rect, buf: &mut Buffer, is_stats_table: bo // Canonicalize each field once so the per-row `execute_scalar` calls below do not // re-decode the field encodings for every row. let field_arrays: Vec = struct_array - .unmasked_fields() - .iter() + .iter_unmasked_fields() .map(|field| { field .clone()