Implement ReduceNode for Array to not require creating new arc when returning children

vortex-array/public-api.lock

@@ -9172,8 +9172,6 @@ pub fn vortex_array::expr::ReduceNode::child(&self, idx: usize) -> vortex_array:
 
 pub fn vortex_array::expr::ReduceNode::child_count(&self) -> usize
 
-pub fn vortex_array::expr::ReduceNode::children(&self) -> alloc::vec::Vec<vortex_array::expr::ReduceNodeRef>
-
 pub fn vortex_array::expr::ReduceNode::node_dtype(&self) -> vortex_error::VortexResult<vortex_dtype::dtype::DType>
 
 pub fn vortex_array::expr::ReduceNode::scalar_fn(&self) -> core::option::Option<&vortex_array::expr::ScalarFn>
@@ -9202,6 +9200,18 @@ pub fn vortex_array::arrays::ScalarFnArray::node_dtype(&self) -> vortex_error::V
 
 pub fn vortex_array::arrays::ScalarFnArray::scalar_fn(&self) -> core::option::Option<&vortex_array::expr::ScalarFn>
 
+impl<V: vortex_array::vtable::VTable> vortex_array::expr::ReduceNode for vortex_array::ArrayAdapter<V>
+
+pub fn vortex_array::ArrayAdapter<V>::as_any(&self) -> &dyn core::any::Any
+
+pub fn vortex_array::ArrayAdapter<V>::child(&self, idx: usize) -> vortex_array::expr::ReduceNodeRef
+
+pub fn vortex_array::ArrayAdapter<V>::child_count(&self) -> usize
+
+pub fn vortex_array::ArrayAdapter<V>::node_dtype(&self) -> vortex_error::VortexResult<vortex_dtype::dtype::DType>
+
+pub fn vortex_array::ArrayAdapter<V>::scalar_fn(&self) -> core::option::Option<&vortex_array::expr::ScalarFn>
+
 pub trait vortex_array::expr::SimplifyCtx
 
 pub fn vortex_array::expr::SimplifyCtx::return_dtype(&self, expr: &vortex_array::expr::Expression) -> vortex_error::VortexResult<vortex_dtype::dtype::DType>
@@ -12768,6 +12778,18 @@ pub fn vortex_array::ArrayAdapter<V>::nchildren(&self) -> usize
 
 pub fn vortex_array::ArrayAdapter<V>::nth_child(&self, idx: usize) -> core::option::Option<vortex_array::ArrayRef>
 
+impl<V: vortex_array::vtable::VTable> vortex_array::expr::ReduceNode for vortex_array::ArrayAdapter<V>
+
+pub fn vortex_array::ArrayAdapter<V>::as_any(&self) -> &dyn core::any::Any
+
+pub fn vortex_array::ArrayAdapter<V>::child(&self, idx: usize) -> vortex_array::expr::ReduceNodeRef
+
+pub fn vortex_array::ArrayAdapter<V>::child_count(&self) -> usize
+
+pub fn vortex_array::ArrayAdapter<V>::node_dtype(&self) -> vortex_error::VortexResult<vortex_dtype::dtype::DType>
+
+pub fn vortex_array::ArrayAdapter<V>::scalar_fn(&self) -> core::option::Option<&vortex_array::expr::ScalarFn>
+
 pub struct vortex_array::CanonicalValidity(pub vortex_array::Canonical)
 
 impl vortex_array::Executable for vortex_array::CanonicalValidity
@@ -12882,7 +12904,7 @@ pub fn vortex_array::RecursiveCanonical::execute(array: vortex_array::ArrayRef,
 
 pub static vortex_array::LEGACY_SESSION: std::sync::lazy_lock::LazyLock<vortex_session::VortexSession>
 
-pub trait vortex_array::Array: 'static + vortex_array::array::private::Sealed + core::marker::Send + core::marker::Sync + core::fmt::Debug + vortex_array::DynArrayEq + vortex_array::DynArrayHash + vortex_array::ArrayVisitor
+pub trait vortex_array::Array: 'static + vortex_array::array::private::Sealed + core::marker::Send + core::marker::Sync + core::fmt::Debug + vortex_array::DynArrayEq + vortex_array::DynArrayHash + vortex_array::ArrayVisitor + vortex_array::expr::ReduceNode
 
 pub fn vortex_array::Array::all_invalid(&self) -> vortex_error::VortexResult<bool>
 

vortex-array/src/array/mod.rs

@@ -39,12 +39,16 @@ use crate::arrays::DictArray;
 use crate::arrays::FilterArray;
 use crate::arrays::NullVTable;
 use crate::arrays::PrimitiveVTable;
+use crate::arrays::ScalarFnVTable;
 use crate::arrays::SliceArray;
 use crate::arrays::VarBinVTable;
 use crate::arrays::VarBinViewVTable;
 use crate::buffer::BufferHandle;
 use crate::builders::ArrayBuilder;
 use crate::compute;
+use crate::expr::ReduceNode;
+use crate::expr::ReduceNodeRef;
+use crate::expr::ScalarFn;
 use crate::expr::stats::Precision;
 use crate::expr::stats::Stat;
 use crate::expr::stats::StatsProviderExt;
@@ -64,7 +68,15 @@ use crate::vtable::VisitorVTable;
 
 /// The public API trait for all Vortex arrays.
 pub trait Array:
-    'static + private::Sealed + Send + Sync + Debug + DynArrayEq + DynArrayHash + ArrayVisitor
+    'static
+    + private::Sealed
+    + Send
+    + Sync
+    + Debug
+    + DynArrayEq
+    + DynArrayHash
+    + ArrayVisitor
+    + ReduceNode
 {
     /// Returns the array as a reference to a generic [`Any`] trait object.
     fn as_any(&self) -> &dyn Any;
@@ -159,7 +171,7 @@ pub trait Array:
 impl Array for Arc<dyn Array> {
     #[inline]
     fn as_any(&self) -> &dyn Any {
-        self.as_ref().as_any()
+        Array::as_any(self.as_ref())
     }
 
     fn as_any_arc(self: Arc<Self>) -> Arc<dyn Any + Send + Sync> {
@@ -386,6 +398,29 @@ impl<V: VTable> Debug for ArrayAdapter<V> {
     }
 }
 
+impl<V: VTable> ReduceNode for ArrayAdapter<V> {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn node_dtype(&self) -> VortexResult<DType> {
+        Ok(<V::ArrayVTable as BaseArrayVTable<V>>::dtype(&self.0).clone())
+    }
+
+    fn scalar_fn(&self) -> Option<&ScalarFn> {
+        self.0.as_opt::<ScalarFnVTable>().map(|a| a.scalar_fn())
+    }
+
+    fn child(&self, idx: usize) -> ReduceNodeRef {
+        self.nth_child(idx)
+            .unwrap_or_else(|| vortex_panic!("Child index out of bounds: {}", idx))
+    }
+
+    fn child_count(&self) -> usize {
+        self.nchildren()
+    }
+}
+
 impl<V: VTable> Array for ArrayAdapter<V> {
     fn as_any(&self) -> &dyn Any {
         self
@@ -779,12 +814,11 @@ impl<V: VTable> Matcher for V {
     type Match<'a> = &'a V::Array;
 
     fn matches(array: &dyn Array) -> bool {
-        array.as_any().is::<ArrayAdapter<V>>()
+        Array::as_any(array).is::<ArrayAdapter<V>>()
     }
 
     fn try_match<'a>(array: &'a dyn Array) -> Option<Self::Match<'a>> {
-        array
-            .as_any()
+        Array::as_any(array)
             .downcast_ref::<ArrayAdapter<V>>()
             .map(|array_adapter| &array_adapter.0)
     }

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

@@ -46,7 +46,6 @@ impl ScalarFnArray {
     }
 
     /// Get the children arrays of this scalar function array.
-    #[allow(clippy::same_name_method)]
     pub fn children(&self) -> &[ArrayRef] {
         &self.children
     }

vortex-array/src/arrays/scalar_fn/rules.rs

@@ -11,7 +11,6 @@ use vortex_error::VortexResult;
 
 use crate::Array;
 use crate::ArrayRef;
-use crate::ArrayVisitor;
 use crate::Canonical;
 use crate::IntoArray;
 use crate::arrays::ConstantArray;
@@ -135,22 +134,19 @@ impl ReduceNode for ArrayRef {
     }
 
     fn node_dtype(&self) -> VortexResult<DType> {
-        Ok(self.as_ref().dtype().clone())
+        self.as_ref().node_dtype()
     }
 
     fn scalar_fn(&self) -> Option<&ScalarFn> {
-        self.as_opt::<ScalarFnVTable>().map(|a| a.scalar_fn())
+        self.as_ref().scalar_fn()
     }
 
     fn child(&self, idx: usize) -> ReduceNodeRef {
-        Arc::new(
-            self.nth_child(idx)
-                .vortex_expect("child index out of bounds"),
-        )
+        self.as_ref().child(idx)
     }
 
     fn child_count(&self) -> usize {
-        self.nchildren()
+        self.as_ref().child_count()
     }
 }
 

vortex-array/src/expr/exprs/merge.rs

@@ -187,12 +187,11 @@ impl VTable for Merge {
         node: &dyn ReduceNode,
         ctx: &dyn ReduceCtx,
     ) -> VortexResult<Option<ReduceNodeRef>> {
-        let merge_dtype = node.node_dtype()?;
         let mut names = Vec::with_capacity(node.child_count() * 2);
         let mut children = Vec::with_capacity(node.child_count() * 2);
         let mut duplicate_names = HashSet::<_>::new();
 
-        for child in node.children() {
+        for child in (0..node.child_count()).map(|i| node.child(i)) {
             let child_dtype = child.node_dtype()?;
             if !child_dtype.is_struct() {
                 vortex_bail!(
@@ -232,7 +231,7 @@ impl VTable for Merge {
         let pack_expr = ctx.new_node(
             Pack.bind(PackOptions {
                 names: FieldNames::from(names),
-                nullability: merge_dtype.nullability(),
+                nullability: node.node_dtype()?.nullability(),
             }),
             &pack_children,
         )?;

vortex-array/src/expr/vtable.rs

@@ -241,12 +241,6 @@ pub trait ReduceNode {
 
     /// Returns the number of children of this node.
     fn child_count(&self) -> usize;
-
-    /// Returns the children of this node.
-    #[allow(clippy::same_name_method)]
-    fn children(&self) -> Vec<ReduceNodeRef> {
-        (0..self.child_count()).map(|i| self.child(i)).collect()
-    }
 }
 
 /// The arity (number of arguments) of a function.

vortex-array/src/hash.rs

@@ -94,7 +94,7 @@ impl ArrayHash for dyn Array + '_ {
 
 impl ArrayEq for dyn Array + '_ {
     fn array_eq(&self, other: &Self, precision: Precision) -> bool {
-        self.dyn_array_eq(other.as_any(), precision)
+        self.dyn_array_eq(Array::as_any(other), precision)
     }
 }