Route StructArray through the array_slots accessors and drop unmasked_fields()

encodings/sparse/src/canonical.rs

@@ -475,7 +475,6 @@ fn execute_sparse_struct(
         ),
     };
     let patch_values_as_struct = unresolved_patches.values().as_::<Struct>().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 {

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,

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<bool> {
-    for field in s.unmasked_fields().iter() {
+    for field in s.iter_unmasked_fields() {
         if !is_constant(field, ctx)? {
             return Ok(false);
         }

vortex-array/src/arrays/chunked/tests.rs

@@ -376,10 +376,12 @@ pub fn pack_nested_structs() -> VortexResult<()> {
     )?
     .into_array();
     let canonical_struct = chunked.execute::<StructArray>(&mut ctx)?;
-    let canonical_varbin = canonical_struct.unmasked_fields()[0]
+    let canonical_varbin = canonical_struct
+        .unmasked_field(0)
         .clone()
         .execute::<VarBinViewArray>(&mut ctx)?;
-    let original_varbin = struct_array.unmasked_fields()[0]
+    let original_varbin = struct_array
+        .unmasked_field(0)
         .clone()
         .execute::<VarBinViewArray>(&mut ctx)?;
     assert_arrays_eq!(original_varbin, canonical_varbin, &mut ctx);

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);

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<StructArray> {
     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<UnionArray> {
@@ -155,10 +160,9 @@ fn mask_validity_union(array: UnionArray, validity: Validity) -> VortexResult<Un
         .clone()
         .mask(validity.to_array(array.len()))?;
     let variants = array.variants().clone();
-    let children = array.children().to_vec();
 
     // SAFETY: We're only changing validity, not the data structure.
-    Ok(unsafe { UnionArray::new_unchecked(type_ids, variants, children) })
+    Ok(unsafe { UnionArray::new_unchecked(type_ids, variants, array.iter_children().cloned()) })
 }
 
 fn mask_validity_extension(

vortex-array/src/arrays/struct_/array.rs

@@ -169,20 +169,39 @@ pub struct StructDataParts {
     pub validity: Validity,
 }
 
-pub(super) fn make_struct_slots(
-    fields: &[ArrayRef],
+/// Allocate the slots of a [`Struct`] array holding only its validity, with room for `nfields`
+/// field slots to be pushed.
+///
+/// Callers that build their field arrays fallibly use this directly, so they don't have to stage
+/// them in a `Vec` just to hand them over as an infallible iterator.
+pub(super) fn struct_slots_with_capacity(
     validity: &Validity,
     length: usize,
+    nfields: usize,
 ) -> 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<Struct> {
+pub(super) fn make_struct_slots(
+    fields: impl IntoIterator<Item = ArrayRef>,
+    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<Struct> {
     }
 
     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<Item = &ArrayRef> + '_ {
-        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<Item = &ArrayRef> + '_ {
+        self.fields().iter()
     }
 
-    fn unmasked_fields(&self) -> Vec<ArrayRef> {
-        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<Struct> {
         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<Struct> {
         validity: Validity,
     ) -> VortexResult<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);
         Array::try_from_parts(
             ArrayParts::new(Struct, outer_dtype, length, EmptyArrayData).with_slots(slots),
         )
@@ -404,7 +419,7 @@ impl Array<Struct> {
             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<Struct> {
 
     // TODO(ngates): remove this... it doesn't help to consume self.
     pub fn into_data_parts(self) -> StructDataParts {
-        let fields: Vec<ArrayRef> = 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<Struct> {
         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<Struct> {
         let types = struct_dtype.fields().chain(once(array.dtype().clone()));
         let new_fields = StructFields::new(names.collect(), types.collect());
 
-        let children: Vec<ArrayRef> = 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> {
             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,
             );

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<Option<ArrayRef>> {
         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()))?,

vortex-array/src/arrays/struct_/compute/rules.rs

@@ -333,7 +333,7 @@ mod tests {
             .unwrap()
             .execute::<StructArray>(&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(),

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::<Vec<_>>(),
+                array.iter_unmasked_fields().cloned(),
                 array.struct_fields().clone(),
                 indices.len(),
                 Validity::AllInvalid,

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;

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))
     }
 

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()))
     }

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::<StructArray>(&mut ctx)?
-            .unmasked_fields()
-            .get(index)
+        let struct_array = array.clone().execute::<StructArray>(&mut ctx)?;
+        let field_array = struct_array
+            .unmasked_field_opt(index)
             .ok_or_else(|| vortex_err!("Field index out of bounds"))?
             .clone();