Try and improve allocation-happy code paths

vortex-array/src/array/mod.rs

@@ -633,23 +633,27 @@ impl<V: VTable> ArrayVisitor for ArrayAdapter<V> {
             children: Vec<ArrayRef>,
         }
 
-        impl ArrayChildVisitor for ChildrenCollector {
-            fn visit_child(&mut self, _name: &str, array: &ArrayRef) {
+        impl ArrayChildVisitorUnnamed for ChildrenCollector {
+            fn visit_child(&mut self, array: &ArrayRef) {
                 self.children.push(array.clone());
             }
         }
 
         let mut collector = ChildrenCollector {
             children: Vec::new(),
         };
-        <V::VisitorVTable as VisitorVTable<V>>::visit_children(&self.0, &mut collector);
+        <V::VisitorVTable as VisitorVTable<V>>::visit_children_unnamed(&self.0, &mut collector);
         collector.children
     }
 
     fn nchildren(&self) -> usize {
         <V::VisitorVTable as VisitorVTable<V>>::nchildren(&self.0)
     }
 
+    fn nth_child(&self, idx: usize) -> Option<ArrayRef> {
+        <V::VisitorVTable as VisitorVTable<V>>::nth_child(&self.0, idx)
+    }
+
     fn children_names(&self) -> Vec<String> {
         struct ChildNameCollector {
             names: Vec<String>,

vortex-array/src/array/visitor.rs

@@ -22,6 +22,11 @@ pub trait ArrayVisitor {
     /// Returns the number of children of the array.
     fn nchildren(&self) -> usize;
 
+    /// Returns the nth child of the array without allocating a Vec.
+    ///
+    /// Returns `None` if the index is out of bounds.
+    fn nth_child(&self, idx: usize) -> Option<ArrayRef>;
+
     /// Returns the names of the children of the array.
     fn children_names(&self) -> Vec<String>;
 
@@ -65,6 +70,10 @@ impl ArrayVisitor for Arc<dyn Array> {
         self.as_ref().nchildren()
     }
 
+    fn nth_child(&self, idx: usize) -> Option<ArrayRef> {
+        self.as_ref().nth_child(idx)
+    }
+
     fn children_names(&self) -> Vec<String> {
         self.as_ref().children_names()
     }
@@ -147,6 +156,40 @@ pub trait ArrayBufferVisitor {
     fn visit_buffer_handle(&mut self, _name: &str, handle: &BufferHandle);
 }
 
+/// A visitor for array children that does not require names.
+///
+/// This is more efficient than [`ArrayChildVisitor`] when you only need to
+/// iterate over children without accessing their names (e.g., for counting
+/// or accessing by index).
+pub trait ArrayChildVisitorUnnamed {
+    /// Visit a child of this array.
+    fn visit_child(&mut self, array: &ArrayRef);
+
+    /// Utility for visiting Array validity.
+    fn visit_validity(&mut self, validity: &Validity, len: usize) {
+        if let Some(vlen) = validity.maybe_len() {
+            assert_eq!(vlen, len, "Validity length mismatch");
+        }
+
+        match validity {
+            Validity::NonNullable | Validity::AllValid => {}
+            Validity::AllInvalid => self.visit_child(&ConstantArray::new(false, len).into_array()),
+            Validity::Array(array) => {
+                self.visit_child(array);
+            }
+        }
+    }
+
+    /// Utility for visiting Array patches.
+    fn visit_patches(&mut self, patches: &Patches) {
+        self.visit_child(patches.indices());
+        self.visit_child(patches.values());
+        if let Some(chunk_offsets) = patches.chunk_offsets() {
+            self.visit_child(chunk_offsets);
+        }
+    }
+}
+
 pub trait ArrayChildVisitor {
     /// Visit a child of this array.
     fn visit_child(&mut self, _name: &str, _array: &ArrayRef);

vortex-array/src/arrays/bool/vtable/visitor.rs

@@ -3,9 +3,12 @@
 
 use crate::ArrayBufferVisitor;
 use crate::ArrayChildVisitor;
+use crate::ArrayRef;
 use crate::arrays::BoolArray;
 use crate::arrays::BoolVTable;
 use crate::vtable::VisitorVTable;
+use crate::vtable::validity_nchildren;
+use crate::vtable::validity_to_child;
 
 impl VisitorVTable<BoolVTable> for BoolVTable {
     fn visit_buffers(array: &BoolArray, visitor: &mut dyn ArrayBufferVisitor) {
@@ -15,4 +18,15 @@ impl VisitorVTable<BoolVTable> for BoolVTable {
     fn visit_children(array: &BoolArray, visitor: &mut dyn ArrayChildVisitor) {
         visitor.visit_validity(&array.validity, array.len());
     }
+
+    fn nchildren(array: &BoolArray) -> usize {
+        validity_nchildren(&array.validity)
+    }
+
+    fn nth_child(array: &BoolArray, idx: usize) -> Option<ArrayRef> {
+        match idx {
+            0 => validity_to_child(&array.validity, array.len()),
+            _ => None,
+        }
+    }
 }

vortex-array/src/arrays/chunked/vtable/visitor.rs

@@ -3,6 +3,7 @@
 
 use crate::ArrayBufferVisitor;
 use crate::ArrayChildVisitor;
+use crate::ArrayChildVisitorUnnamed;
 use crate::arrays::ChunkedArray;
 use crate::arrays::ChunkedVTable;
 use crate::vtable::VisitorVTable;
@@ -17,4 +18,23 @@ impl VisitorVTable<ChunkedVTable> for ChunkedVTable {
             visitor.visit_child(format!("chunks[{idx}]").as_str(), chunk);
         }
     }
+
+    fn visit_children_unnamed(array: &ChunkedArray, visitor: &mut dyn ArrayChildVisitorUnnamed) {
+        visitor.visit_child(&array.chunk_offsets.to_array());
+
+        for chunk in array.chunks().iter() {
+            visitor.visit_child(chunk);
+        }
+    }
+
+    fn nchildren(array: &ChunkedArray) -> usize {
+        1 + array.chunks().len()
+    }
+
+    fn nth_child(array: &ChunkedArray, idx: usize) -> Option<crate::ArrayRef> {
+        match idx {
+            0 => Some(array.chunk_offsets.to_array()),
+            n => array.chunks().get(n - 1).cloned(),
+        }
+    }
 }

vortex-array/src/arrays/constant/vtable/visitor.rs

@@ -5,6 +5,7 @@ use vortex_buffer::ByteBufferMut;
 
 use crate::ArrayBufferVisitor;
 use crate::ArrayChildVisitor;
+use crate::ArrayRef;
 use crate::arrays::ConstantArray;
 use crate::arrays::ConstantVTable;
 use crate::buffer::BufferHandle;
@@ -21,4 +22,12 @@ impl VisitorVTable<ConstantVTable> for ConstantVTable {
     }
 
     fn visit_children(_array: &ConstantArray, _visitor: &mut dyn ArrayChildVisitor) {}
+
+    fn nchildren(_array: &ConstantArray) -> usize {
+        0
+    }
+
+    fn nth_child(_array: &ConstantArray, _idx: usize) -> Option<ArrayRef> {
+        None
+    }
 }

vortex-array/src/arrays/decimal/vtable/visitor.rs

@@ -3,10 +3,13 @@
 
 use crate::ArrayBufferVisitor;
 use crate::ArrayChildVisitor;
+use crate::ArrayRef;
 use crate::arrays::DecimalArray;
 use crate::arrays::DecimalVTable;
 use crate::vtable::ValidityHelper;
 use crate::vtable::VisitorVTable;
+use crate::vtable::validity_nchildren;
+use crate::vtable::validity_to_child;
 
 impl VisitorVTable<DecimalVTable> for DecimalVTable {
     fn visit_buffers(array: &DecimalArray, visitor: &mut dyn ArrayBufferVisitor) {
@@ -16,4 +19,15 @@ impl VisitorVTable<DecimalVTable> for DecimalVTable {
     fn visit_children(array: &DecimalArray, visitor: &mut dyn ArrayChildVisitor) {
         visitor.visit_validity(array.validity(), array.len())
     }
+
+    fn nchildren(array: &DecimalArray) -> usize {
+        validity_nchildren(array.validity())
+    }
+
+    fn nth_child(array: &DecimalArray, idx: usize) -> Option<ArrayRef> {
+        match idx {
+            0 => validity_to_child(array.validity(), array.len()),
+            _ => None,
+        }
+    }
 }

vortex-array/src/arrays/dict/vtable/visitor.rs

@@ -4,6 +4,7 @@
 use super::DictVTable;
 use crate::ArrayBufferVisitor;
 use crate::ArrayChildVisitor;
+use crate::ArrayRef;
 use crate::arrays::dict::DictArray;
 use crate::vtable::VisitorVTable;
 
@@ -14,4 +15,16 @@ impl VisitorVTable<DictVTable> for DictVTable {
         visitor.visit_child("codes", array.codes());
         visitor.visit_child("values", array.values());
     }
+
+    fn nchildren(_array: &DictArray) -> usize {
+        2
+    }
+
+    fn nth_child(array: &DictArray, idx: usize) -> Option<ArrayRef> {
+        match idx {
+            0 => Some(array.codes().clone()),
+            1 => Some(array.values().clone()),
+            _ => None,
+        }
+    }
 }

vortex-array/src/arrays/extension/vtable/visitor.rs

@@ -3,6 +3,7 @@
 
 use crate::ArrayBufferVisitor;
 use crate::ArrayChildVisitor;
+use crate::ArrayRef;
 use crate::arrays::extension::ExtensionArray;
 use crate::arrays::extension::ExtensionVTable;
 use crate::vtable::VisitorVTable;
@@ -13,4 +14,15 @@ impl VisitorVTable<ExtensionVTable> for ExtensionVTable {
     fn visit_children(array: &ExtensionArray, visitor: &mut dyn ArrayChildVisitor) {
         visitor.visit_child("storage", &array.storage);
     }
+
+    fn nchildren(_array: &ExtensionArray) -> usize {
+        1
+    }
+
+    fn nth_child(array: &ExtensionArray, idx: usize) -> Option<ArrayRef> {
+        match idx {
+            0 => Some(array.storage.clone()),
+            _ => None,
+        }
+    }
 }

vortex-array/src/arrays/filter/vtable.rs

@@ -176,6 +176,17 @@ impl VisitorVTable<FilterVTable> for FilterVTable {
     fn visit_children(array: &FilterArray, visitor: &mut dyn ArrayChildVisitor) {
         visitor.visit_child("child", &array.child);
     }
+
+    fn nchildren(_array: &FilterArray) -> usize {
+        1
+    }
+
+    fn nth_child(array: &FilterArray, idx: usize) -> Option<ArrayRef> {
+        match idx {
+            0 => Some(array.child.clone()),
+            _ => None,
+        }
+    }
 }
 
 pub struct FilterMetadata(pub(super) Mask);

vortex-array/src/arrays/fixed_size_list/vtable/visitor.rs

@@ -3,10 +3,13 @@
 
 use crate::ArrayBufferVisitor;
 use crate::ArrayChildVisitor;
+use crate::ArrayRef;
 use crate::arrays::FixedSizeListArray;
 use crate::arrays::FixedSizeListVTable;
 use crate::vtable::ValidityHelper;
 use crate::vtable::VisitorVTable;
+use crate::vtable::validity_nchildren;
+use crate::vtable::validity_to_child;
 
 impl VisitorVTable<FixedSizeListVTable> for FixedSizeListVTable {
     fn visit_buffers(_array: &FixedSizeListArray, _visitor: &mut dyn ArrayBufferVisitor) {
@@ -18,4 +21,16 @@ impl VisitorVTable<FixedSizeListVTable> for FixedSizeListVTable {
         visitor.visit_child("elements", array.elements());
         visitor.visit_validity(array.validity(), array.len());
     }
+
+    fn nchildren(array: &FixedSizeListArray) -> usize {
+        1 + validity_nchildren(array.validity())
+    }
+
+    fn nth_child(array: &FixedSizeListArray, idx: usize) -> Option<ArrayRef> {
+        match idx {
+            0 => Some(array.elements().clone()),
+            1 => validity_to_child(array.validity(), array.len()),
+            _ => None,
+        }
+    }
 }

vortex-array/src/arrays/list/vtable/visitor.rs

@@ -3,10 +3,13 @@
 
 use crate::ArrayBufferVisitor;
 use crate::ArrayChildVisitor;
+use crate::ArrayRef;
 use crate::arrays::ListArray;
 use crate::arrays::ListVTable;
 use crate::vtable::ValidityHelper;
 use crate::vtable::VisitorVTable;
+use crate::vtable::validity_nchildren;
+use crate::vtable::validity_to_child;
 
 impl VisitorVTable<ListVTable> for ListVTable {
     fn visit_buffers(_array: &ListArray, _visitor: &mut dyn ArrayBufferVisitor) {}
@@ -16,4 +19,17 @@ impl VisitorVTable<ListVTable> for ListVTable {
         visitor.visit_child("offsets", array.offsets());
         visitor.visit_validity(array.validity(), array.len());
     }
+
+    fn nchildren(array: &ListArray) -> usize {
+        2 + validity_nchildren(array.validity())
+    }
+
+    fn nth_child(array: &ListArray, idx: usize) -> Option<ArrayRef> {
+        match idx {
+            0 => Some(array.elements().clone()),
+            1 => Some(array.offsets().clone()),
+            2 => validity_to_child(array.validity(), array.len()),
+            _ => None,
+        }
+    }
 }

vortex-array/src/arrays/listview/vtable/visitor.rs

@@ -3,10 +3,13 @@
 
 use crate::ArrayBufferVisitor;
 use crate::ArrayChildVisitor;
+use crate::ArrayRef;
 use crate::arrays::ListViewArray;
 use crate::arrays::ListViewVTable;
 use crate::vtable::ValidityHelper;
 use crate::vtable::VisitorVTable;
+use crate::vtable::validity_nchildren;
+use crate::vtable::validity_to_child;
 
 impl VisitorVTable<ListViewVTable> for ListViewVTable {
     fn visit_buffers(_array: &ListViewArray, _visitor: &mut dyn ArrayBufferVisitor) {
@@ -19,4 +22,18 @@ impl VisitorVTable<ListViewVTable> for ListViewVTable {
         visitor.visit_child("sizes", array.sizes());
         visitor.visit_validity(array.validity(), array.len());
     }
+
+    fn nchildren(array: &ListViewArray) -> usize {
+        3 + validity_nchildren(array.validity())
+    }
+
+    fn nth_child(array: &ListViewArray, idx: usize) -> Option<ArrayRef> {
+        match idx {
+            0 => Some(array.elements().clone()),
+            1 => Some(array.offsets().clone()),
+            2 => Some(array.sizes().clone()),
+            3 => validity_to_child(array.validity(), array.len()),
+            _ => None,
+        }
+    }
 }

vortex-array/src/arrays/masked/vtable/mod.rs

@@ -33,6 +33,8 @@ use crate::vtable::ArrayId;
 use crate::vtable::VTable;
 use crate::vtable::ValidityVTableFromValidityHelper;
 use crate::vtable::VisitorVTable;
+use crate::vtable::validity_nchildren;
+use crate::vtable::validity_to_child;
 
 vtable!(Masked);
 
@@ -50,6 +52,18 @@ impl VisitorVTable<MaskedVTable> for MaskedVTable {
         visitor.visit_child("child", &array.child);
         visitor.visit_validity(&array.validity, array.child.len());
     }
+
+    fn nchildren(array: &MaskedArray) -> usize {
+        1 + validity_nchildren(&array.validity)
+    }
+
+    fn nth_child(array: &MaskedArray, idx: usize) -> Option<ArrayRef> {
+        match idx {
+            0 => Some(array.child.clone()),
+            1 => validity_to_child(&array.validity, array.child.len()),
+            _ => None,
+        }
+    }
 }
 
 impl VTable for MaskedVTable {