diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs index 7af207736ad..1bf09be641a 100644 --- a/vortex-array/src/array/mod.rs +++ b/vortex-array/src/array/mod.rs @@ -633,8 +633,8 @@ impl ArrayVisitor for ArrayAdapter { children: Vec, } - 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()); } } @@ -642,7 +642,7 @@ impl ArrayVisitor for ArrayAdapter { let mut collector = ChildrenCollector { children: Vec::new(), }; - >::visit_children(&self.0, &mut collector); + >::visit_children_unnamed(&self.0, &mut collector); collector.children } @@ -650,6 +650,10 @@ impl ArrayVisitor for ArrayAdapter { >::nchildren(&self.0) } + fn nth_child(&self, idx: usize) -> Option { + >::nth_child(&self.0, idx) + } + fn children_names(&self) -> Vec { struct ChildNameCollector { names: Vec, diff --git a/vortex-array/src/array/visitor.rs b/vortex-array/src/array/visitor.rs index 61c1167daad..8067fafc3f4 100644 --- a/vortex-array/src/array/visitor.rs +++ b/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; + /// Returns the names of the children of the array. fn children_names(&self) -> Vec; @@ -65,6 +70,10 @@ impl ArrayVisitor for Arc { self.as_ref().nchildren() } + fn nth_child(&self, idx: usize) -> Option { + self.as_ref().nth_child(idx) + } + fn children_names(&self) -> Vec { 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); diff --git a/vortex-array/src/arrays/bool/vtable/visitor.rs b/vortex-array/src/arrays/bool/vtable/visitor.rs index 4e29761ae7a..26e18d9ad3d 100644 --- a/vortex-array/src/arrays/bool/vtable/visitor.rs +++ b/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 for BoolVTable { fn visit_buffers(array: &BoolArray, visitor: &mut dyn ArrayBufferVisitor) { @@ -15,4 +18,15 @@ impl VisitorVTable 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 { + match idx { + 0 => validity_to_child(&array.validity, array.len()), + _ => None, + } + } } diff --git a/vortex-array/src/arrays/chunked/vtable/visitor.rs b/vortex-array/src/arrays/chunked/vtable/visitor.rs index 398a9110b5b..a8da03eb7ac 100644 --- a/vortex-array/src/arrays/chunked/vtable/visitor.rs +++ b/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 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 { + match idx { + 0 => Some(array.chunk_offsets.to_array()), + n => array.chunks().get(n - 1).cloned(), + } + } } diff --git a/vortex-array/src/arrays/constant/vtable/visitor.rs b/vortex-array/src/arrays/constant/vtable/visitor.rs index 28e613ba121..7e26682a248 100644 --- a/vortex-array/src/arrays/constant/vtable/visitor.rs +++ b/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 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 { + None + } } diff --git a/vortex-array/src/arrays/decimal/vtable/visitor.rs b/vortex-array/src/arrays/decimal/vtable/visitor.rs index 74fd694b3bd..e45e7c0c33d 100644 --- a/vortex-array/src/arrays/decimal/vtable/visitor.rs +++ b/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 for DecimalVTable { fn visit_buffers(array: &DecimalArray, visitor: &mut dyn ArrayBufferVisitor) { @@ -16,4 +19,15 @@ impl VisitorVTable 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 { + match idx { + 0 => validity_to_child(array.validity(), array.len()), + _ => None, + } + } } diff --git a/vortex-array/src/arrays/dict/vtable/visitor.rs b/vortex-array/src/arrays/dict/vtable/visitor.rs index 584cda832d9..3d233d7d180 100644 --- a/vortex-array/src/arrays/dict/vtable/visitor.rs +++ b/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 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 { + match idx { + 0 => Some(array.codes().clone()), + 1 => Some(array.values().clone()), + _ => None, + } + } } diff --git a/vortex-array/src/arrays/extension/vtable/visitor.rs b/vortex-array/src/arrays/extension/vtable/visitor.rs index 8ebf8404150..68216350a92 100644 --- a/vortex-array/src/arrays/extension/vtable/visitor.rs +++ b/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 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 { + match idx { + 0 => Some(array.storage.clone()), + _ => None, + } + } } diff --git a/vortex-array/src/arrays/filter/vtable.rs b/vortex-array/src/arrays/filter/vtable.rs index aa83561e3cd..d2876e9e197 100644 --- a/vortex-array/src/arrays/filter/vtable.rs +++ b/vortex-array/src/arrays/filter/vtable.rs @@ -176,6 +176,17 @@ impl VisitorVTable 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 { + match idx { + 0 => Some(array.child.clone()), + _ => None, + } + } } pub struct FilterMetadata(pub(super) Mask); diff --git a/vortex-array/src/arrays/fixed_size_list/vtable/visitor.rs b/vortex-array/src/arrays/fixed_size_list/vtable/visitor.rs index a2c525df56f..85f32a5aa6f 100644 --- a/vortex-array/src/arrays/fixed_size_list/vtable/visitor.rs +++ b/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 for FixedSizeListVTable { fn visit_buffers(_array: &FixedSizeListArray, _visitor: &mut dyn ArrayBufferVisitor) { @@ -18,4 +21,16 @@ impl VisitorVTable 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 { + match idx { + 0 => Some(array.elements().clone()), + 1 => validity_to_child(array.validity(), array.len()), + _ => None, + } + } } diff --git a/vortex-array/src/arrays/list/vtable/visitor.rs b/vortex-array/src/arrays/list/vtable/visitor.rs index 9eeb112fb40..b5b4685be45 100644 --- a/vortex-array/src/arrays/list/vtable/visitor.rs +++ b/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 for ListVTable { fn visit_buffers(_array: &ListArray, _visitor: &mut dyn ArrayBufferVisitor) {} @@ -16,4 +19,17 @@ impl VisitorVTable 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 { + match idx { + 0 => Some(array.elements().clone()), + 1 => Some(array.offsets().clone()), + 2 => validity_to_child(array.validity(), array.len()), + _ => None, + } + } } diff --git a/vortex-array/src/arrays/listview/vtable/visitor.rs b/vortex-array/src/arrays/listview/vtable/visitor.rs index d1aeabd03d6..8d4fa32a4a8 100644 --- a/vortex-array/src/arrays/listview/vtable/visitor.rs +++ b/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 for ListViewVTable { fn visit_buffers(_array: &ListViewArray, _visitor: &mut dyn ArrayBufferVisitor) { @@ -19,4 +22,18 @@ impl VisitorVTable 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 { + 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, + } + } } diff --git a/vortex-array/src/arrays/masked/vtable/mod.rs b/vortex-array/src/arrays/masked/vtable/mod.rs index 9b0fa70c6d5..d63c9c8c7d2 100644 --- a/vortex-array/src/arrays/masked/vtable/mod.rs +++ b/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 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 { + match idx { + 0 => Some(array.child.clone()), + 1 => validity_to_child(&array.validity, array.child.len()), + _ => None, + } + } } impl VTable for MaskedVTable { diff --git a/vortex-array/src/arrays/null/mod.rs b/vortex-array/src/arrays/null/mod.rs index 44ccba7bf4e..93d796fce5e 100644 --- a/vortex-array/src/arrays/null/mod.rs +++ b/vortex-array/src/arrays/null/mod.rs @@ -171,6 +171,14 @@ impl VisitorVTable for NullVTable { fn visit_buffers(_array: &NullArray, _visitor: &mut dyn ArrayBufferVisitor) {} fn visit_children(_array: &NullArray, _visitor: &mut dyn ArrayChildVisitor) {} + + fn nchildren(_array: &NullArray) -> usize { + 0 + } + + fn nth_child(_array: &NullArray, _idx: usize) -> Option { + None + } } impl OperationsVTable for NullVTable { diff --git a/vortex-array/src/arrays/primitive/vtable/visitor.rs b/vortex-array/src/arrays/primitive/vtable/visitor.rs index 0ae80a26f38..69823c8530e 100644 --- a/vortex-array/src/arrays/primitive/vtable/visitor.rs +++ b/vortex-array/src/arrays/primitive/vtable/visitor.rs @@ -3,10 +3,13 @@ use crate::ArrayBufferVisitor; use crate::ArrayChildVisitor; +use crate::ArrayRef; use crate::arrays::PrimitiveArray; use crate::arrays::PrimitiveVTable; use crate::vtable::ValidityHelper; use crate::vtable::VisitorVTable; +use crate::vtable::validity_nchildren; +use crate::vtable::validity_to_child; impl VisitorVTable for PrimitiveVTable { fn visit_buffers(array: &PrimitiveArray, visitor: &mut dyn ArrayBufferVisitor) { @@ -16,4 +19,15 @@ impl VisitorVTable for PrimitiveVTable { fn visit_children(array: &PrimitiveArray, visitor: &mut dyn ArrayChildVisitor) { visitor.visit_validity(array.validity(), array.len()); } + + fn nchildren(array: &PrimitiveArray) -> usize { + validity_nchildren(array.validity()) + } + + fn nth_child(array: &PrimitiveArray, idx: usize) -> Option { + match idx { + 0 => validity_to_child(array.validity(), array.len()), + _ => None, + } + } } diff --git a/vortex-array/src/arrays/scalar_fn/rules.rs b/vortex-array/src/arrays/scalar_fn/rules.rs index 4cc3b5b79bd..e7b08d88fcf 100644 --- a/vortex-array/src/arrays/scalar_fn/rules.rs +++ b/vortex-array/src/arrays/scalar_fn/rules.rs @@ -121,7 +121,10 @@ impl ReduceNode for ArrayRef { } fn child(&self, idx: usize) -> ReduceNodeRef { - Arc::new(::children(self)[idx].clone()) + Arc::new( + self.nth_child(idx) + .vortex_expect("child index out of bounds"), + ) } fn child_count(&self) -> usize { diff --git a/vortex-array/src/arrays/scalar_fn/vtable/visitor.rs b/vortex-array/src/arrays/scalar_fn/vtable/visitor.rs index fa5e8607380..69fc6bea0d3 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/visitor.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/visitor.rs @@ -3,6 +3,8 @@ use crate::ArrayBufferVisitor; use crate::ArrayChildVisitor; +use crate::ArrayChildVisitorUnnamed; +use crate::ArrayRef; use crate::arrays::scalar_fn::array::ScalarFnArray; use crate::arrays::scalar_fn::vtable::ScalarFnVTable; use crate::vtable::VisitorVTable; @@ -16,4 +18,18 @@ impl VisitorVTable for ScalarFnVTable { visitor.visit_child(name.as_ref(), child) } } + + fn visit_children_unnamed(array: &ScalarFnArray, visitor: &mut dyn ArrayChildVisitorUnnamed) { + for child in array.children.iter() { + visitor.visit_child(child); + } + } + + fn nchildren(array: &ScalarFnArray) -> usize { + array.children.len() + } + + fn nth_child(array: &ScalarFnArray, idx: usize) -> Option { + array.children.get(idx).cloned() + } } diff --git a/vortex-array/src/arrays/shared/vtable.rs b/vortex-array/src/arrays/shared/vtable.rs index 510bf256fd7..5ed623d9f31 100644 --- a/vortex-array/src/arrays/shared/vtable.rs +++ b/vortex-array/src/arrays/shared/vtable.rs @@ -144,4 +144,15 @@ impl VisitorVTable for SharedVTable { fn visit_children(array: &SharedArray, visitor: &mut dyn ArrayChildVisitor) { visitor.visit_child("source", &array.current_array_ref()); } + + fn nchildren(_array: &SharedArray) -> usize { + 1 + } + + fn nth_child(array: &SharedArray, idx: usize) -> Option { + match idx { + 0 => Some(array.current_array_ref()), + _ => None, + } + } } diff --git a/vortex-array/src/arrays/slice/vtable.rs b/vortex-array/src/arrays/slice/vtable.rs index 5fbaa21d7df..f0c542fe6d3 100644 --- a/vortex-array/src/arrays/slice/vtable.rs +++ b/vortex-array/src/arrays/slice/vtable.rs @@ -175,6 +175,17 @@ impl VisitorVTable for SliceVTable { fn visit_children(array: &SliceArray, visitor: &mut dyn ArrayChildVisitor) { visitor.visit_child("child", &array.child); } + + fn nchildren(_array: &SliceArray) -> usize { + 1 + } + + fn nth_child(array: &SliceArray, idx: usize) -> Option { + match idx { + 0 => Some(array.child.clone()), + _ => None, + } + } } pub struct SliceMetadata(pub(super) Range); diff --git a/vortex-array/src/arrays/struct_/vtable/visitor.rs b/vortex-array/src/arrays/struct_/vtable/visitor.rs index 83caebc3f55..bd3c0d199f8 100644 --- a/vortex-array/src/arrays/struct_/vtable/visitor.rs +++ b/vortex-array/src/arrays/struct_/vtable/visitor.rs @@ -5,10 +5,14 @@ use itertools::Itertools; use crate::ArrayBufferVisitor; use crate::ArrayChildVisitor; +use crate::ArrayChildVisitorUnnamed; +use crate::ArrayRef; use crate::arrays::struct_::StructArray; use crate::arrays::struct_::StructVTable; use crate::vtable::ValidityHelper; use crate::vtable::VisitorVTable; +use crate::vtable::validity_nchildren; +use crate::vtable::validity_to_child; impl VisitorVTable for StructVTable { fn visit_buffers(_array: &StructArray, _visitor: &mut dyn ArrayBufferVisitor) {} @@ -19,4 +23,27 @@ impl VisitorVTable for StructVTable { visitor.visit_child(name.as_ref(), field); } } + + fn visit_children_unnamed(array: &StructArray, visitor: &mut dyn ArrayChildVisitorUnnamed) { + visitor.visit_validity(array.validity(), array.len()); + for field in array.unmasked_fields().iter() { + visitor.visit_child(field); + } + } + + fn nchildren(array: &StructArray) -> usize { + validity_nchildren(array.validity()) + array.unmasked_fields().len() + } + + fn nth_child(array: &StructArray, idx: usize) -> Option { + let validity_children = validity_nchildren(array.validity()); + if idx < validity_children { + validity_to_child(array.validity(), array.len()) + } else { + array + .unmasked_fields() + .get(idx - validity_children) + .cloned() + } + } } diff --git a/vortex-array/src/arrays/varbin/vtable/visitor.rs b/vortex-array/src/arrays/varbin/vtable/visitor.rs index 7845f53e17a..4cf76382615 100644 --- a/vortex-array/src/arrays/varbin/vtable/visitor.rs +++ b/vortex-array/src/arrays/varbin/vtable/visitor.rs @@ -3,10 +3,13 @@ use crate::ArrayBufferVisitor; use crate::ArrayChildVisitor; +use crate::ArrayRef; use crate::arrays::VarBinArray; use crate::arrays::VarBinVTable; use crate::vtable::ValidityHelper; use crate::vtable::VisitorVTable; +use crate::vtable::validity_nchildren; +use crate::vtable::validity_to_child; impl VisitorVTable for VarBinVTable { fn visit_buffers(array: &VarBinArray, visitor: &mut dyn ArrayBufferVisitor) { @@ -18,4 +21,16 @@ impl VisitorVTable for VarBinVTable { visitor.visit_child("offsets", array.offsets()); visitor.visit_validity(array.validity(), array.len()); } + + fn nchildren(array: &VarBinArray) -> usize { + 1 + validity_nchildren(array.validity()) + } + + fn nth_child(array: &VarBinArray, idx: usize) -> Option { + match idx { + 0 => Some(array.offsets().clone()), + 1 => validity_to_child(array.validity(), array.len()), + _ => None, + } + } } diff --git a/vortex-array/src/arrays/varbinview/vtable/visitor.rs b/vortex-array/src/arrays/varbinview/vtable/visitor.rs index a5321d4cc6f..187a673b0fa 100644 --- a/vortex-array/src/arrays/varbinview/vtable/visitor.rs +++ b/vortex-array/src/arrays/varbinview/vtable/visitor.rs @@ -4,9 +4,12 @@ use super::VarBinViewVTable; use crate::ArrayBufferVisitor; use crate::ArrayChildVisitor; +use crate::ArrayRef; use crate::arrays::VarBinViewArray; use crate::vtable::ValidityHelper; use crate::vtable::VisitorVTable; +use crate::vtable::validity_nchildren; +use crate::vtable::validity_to_child; impl VisitorVTable for VarBinViewVTable { fn visit_buffers(array: &VarBinViewArray, visitor: &mut dyn ArrayBufferVisitor) { @@ -19,4 +22,15 @@ impl VisitorVTable for VarBinViewVTable { fn visit_children(array: &VarBinViewArray, visitor: &mut dyn ArrayChildVisitor) { visitor.visit_validity(array.validity(), array.len()) } + + fn nchildren(array: &VarBinViewArray) -> usize { + validity_nchildren(array.validity()) + } + + fn nth_child(array: &VarBinViewArray, idx: usize) -> Option { + match idx { + 0 => validity_to_child(array.validity(), array.len()), + _ => None, + } + } } diff --git a/vortex-array/src/arrow/array.rs b/vortex-array/src/arrow/array.rs index 49000059da2..0255b6b1e87 100644 --- a/vortex-array/src/arrow/array.rs +++ b/vortex-array/src/arrow/array.rs @@ -178,4 +178,12 @@ impl VisitorVTable for ArrowVTable { fn visit_buffers(_array: &ArrowArray, _visitor: &mut dyn ArrayBufferVisitor) {} fn visit_children(_array: &ArrowArray, _visitor: &mut dyn ArrayChildVisitor) {} + + fn nchildren(_array: &ArrowArray) -> usize { + 0 + } + + fn nth_child(_array: &ArrowArray, _idx: usize) -> Option { + None + } } diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index 2695303c8dc..12a98c1851b 100644 --- a/vortex-array/src/executor.rs +++ b/vortex-array/src/executor.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use std::sync::atomic::AtomicUsize; use itertools::Itertools; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -154,8 +155,9 @@ impl Executable for ArrayRef { } // 2. reduce_parent (child-driven metadata-only rewrites) - for (child_idx, child) in array.children().iter().enumerate() { - if let Some(reduced_parent) = child.vtable().reduce_parent(child, &array, child_idx)? { + for child_idx in 0..array.nchildren() { + let child = array.nth_child(child_idx).vortex_expect("checked length"); + if let Some(reduced_parent) = child.vtable().reduce_parent(&child, &array, child_idx)? { ctx.log(format_args!( "reduce_parent: child[{}]({}) rewrote {} -> {}", child_idx, @@ -169,10 +171,11 @@ impl Executable for ArrayRef { } // 3. execute_parent (child-driven optimized execution) - for (child_idx, child) in array.children().iter().enumerate() { + for child_idx in 0..array.nchildren() { + let child = array.nth_child(child_idx).vortex_expect("checked length"); if let Some(executed_parent) = child .vtable() - .execute_parent(child, &array, child_idx, ctx)? + .execute_parent(&child, &array, child_idx, ctx)? { ctx.log(format_args!( "execute_parent: child[{}]({}) rewrote {} -> {}", diff --git a/vortex-array/src/optimizer/mod.rs b/vortex-array/src/optimizer/mod.rs index c086033646f..71d5f1f0df2 100644 --- a/vortex-array/src/optimizer/mod.rs +++ b/vortex-array/src/optimizer/mod.rs @@ -46,6 +46,7 @@ fn try_optimize(array: &ArrayRef) -> VortexResult> { } // Apply parent reduction rules to each child in the context of the current array. + // Its important to take all children here, as `current_array` can change inside the loop. for (idx, child) in current_array.children().iter().enumerate() { if let Some(new_array) = child.vtable().reduce_parent(child, ¤t_array, idx)? { // If the parent was replaced, then we attempt to reduce it again. diff --git a/vortex-array/src/vtable/visitor.rs b/vortex-array/src/vtable/visitor.rs index 47fac40f170..be90ce029c3 100644 --- a/vortex-array/src/vtable/visitor.rs +++ b/vortex-array/src/vtable/visitor.rs @@ -3,10 +3,37 @@ use crate::ArrayBufferVisitor; use crate::ArrayChildVisitor; +use crate::ArrayChildVisitorUnnamed; use crate::ArrayRef; +use crate::IntoArray; +use crate::arrays::ConstantArray; use crate::buffer::BufferHandle; +use crate::validity::Validity; use crate::vtable::VTable; +/// Returns the validity as a child array if it produces one. +/// +/// - `NonNullable` and `AllValid` produce no child (returns `None`) +/// - `AllInvalid` produces a `ConstantArray` of `false` values +/// - `Array` returns the validity array +#[inline] +pub fn validity_to_child(validity: &Validity, len: usize) -> Option { + match validity { + Validity::NonNullable | Validity::AllValid => None, + Validity::AllInvalid => Some(ConstantArray::new(false, len).into_array()), + Validity::Array(array) => Some(array.clone()), + } +} + +/// Returns 1 if validity produces a child, 0 otherwise. +#[inline] +pub fn validity_nchildren(validity: &Validity) -> usize { + match validity { + Validity::NonNullable | Validity::AllValid => 0, + Validity::AllInvalid | Validity::Array(_) => 1, + } +} + pub trait VisitorVTable { /// Visit the buffers of the array. fn visit_buffers(array: &V::Array, visitor: &mut dyn ArrayBufferVisitor); @@ -44,18 +71,64 @@ pub trait VisitorVTable { /// Visit the children of the array. fn visit_children(array: &V::Array, visitor: &mut dyn ArrayChildVisitor); + /// Visit the children of the array without names. + /// + /// This is more efficient than [`Self::visit_children`] when you don't need the + /// child names (e.g., for counting or accessing by index). The default + /// implementation wraps the named visitor, but array types can override + /// this to avoid allocating names. + fn visit_children_unnamed(array: &V::Array, visitor: &mut dyn ArrayChildVisitorUnnamed) { + struct UnnamedWrapper<'a>(&'a mut dyn ArrayChildVisitorUnnamed); + + impl ArrayChildVisitor for UnnamedWrapper<'_> { + fn visit_child(&mut self, _name: &str, array: &ArrayRef) { + self.0.visit_child(array); + } + } + + >::visit_children(array, &mut UnnamedWrapper(visitor)); + } + /// Count the number of children in the array. fn nchildren(array: &V::Array) -> usize { struct NChildren(usize); - impl ArrayChildVisitor for NChildren { - fn visit_child(&mut self, _name: &str, _array: &ArrayRef) { + impl ArrayChildVisitorUnnamed for NChildren { + fn visit_child(&mut self, _array: &ArrayRef) { self.0 += 1; } } let mut visitor = NChildren(0); - >::visit_children(array, &mut visitor); + >::visit_children_unnamed(array, &mut visitor); visitor.0 } + + /// Get the nth child of the array without allocating a Vec. + /// + /// Returns `None` if the index is out of bounds. + fn nth_child(array: &V::Array, idx: usize) -> Option { + struct NthChildVisitor { + target_idx: usize, + current_idx: usize, + result: Option, + } + + impl ArrayChildVisitorUnnamed for NthChildVisitor { + fn visit_child(&mut self, array: &ArrayRef) { + if self.current_idx == self.target_idx && self.result.is_none() { + self.result = Some(array.clone()); + } + self.current_idx += 1; + } + } + + let mut visitor = NthChildVisitor { + target_idx: idx, + current_idx: 0, + result: None, + }; + >::visit_children_unnamed(array, &mut visitor); + visitor.result + } }