Move rules onto array vtable so we can optimize during construction without a session

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

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

@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::any::Any;
+use std::sync::Arc;
 
 use vortex_dtype::DType;
 use vortex_error::VortexExpect;
@@ -12,6 +13,7 @@ use vortex_vector::Datum;
 use vortex_vector::VectorOps;
 use vortex_vector::datum_matches_dtype;
 
+use crate::Array;
 use crate::ArrayRef;
 use crate::ArrayVisitor;
 use crate::IntoArray;
@@ -22,6 +24,7 @@ use crate::arrays::ScalarFnVTable;
 use crate::expr::ExecutionArgs;
 use crate::expr::ReduceCtx;
 use crate::expr::ReduceNode;
+use crate::expr::ReduceNodeRef;
 use crate::expr::ScalarFn;
 use crate::optimizer::rules::ArrayReduceRule;
 use crate::optimizer::rules::ReduceRuleSet;
@@ -109,8 +112,12 @@ impl ReduceNode for ArrayRef {
         self.as_opt::<ScalarFnVTable>().map(|a| a.scalar_fn())
     }
 
-    fn child(&self, idx: usize) -> Box<dyn ReduceNode> {
-        Box::new(self.children()[idx].clone())
+    fn child(&self, idx: usize) -> ReduceNodeRef {
+        Arc::new(<dyn Array>::children(self)[idx].clone())
+    }
+
+    fn child_count(&self) -> usize {
+        self.nchildren()
     }
 }
 
@@ -119,12 +126,12 @@ struct ArrayReduceCtx {
     len: usize,
 }
 impl ReduceCtx for ArrayReduceCtx {
-    fn create_node(
+    fn new_node(
         &self,
         scalar_fn: ScalarFn,
-        children: &[Box<dyn ReduceNode>],
-    ) -> VortexResult<Box<dyn ReduceNode>> {
-        Ok(Box::new(
+        children: &[ReduceNodeRef],
+    ) -> VortexResult<ReduceNodeRef> {
+        Ok(Arc::new(
             ScalarFnArray::try_new(
                 scalar_fn,
                 children

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

@@ -32,6 +32,7 @@ use crate::expr::Mask;
 use crate::expr::Pack;
 use crate::expr::ReduceCtx;
 use crate::expr::ReduceNode;
+use crate::expr::ReduceNodeRef;
 use crate::expr::StatsCatalog;
 use crate::expr::VTable;
 use crate::expr::VTableExt;
@@ -145,7 +146,7 @@ impl VTable for GetItem {
         field_name: &FieldName,
         node: &dyn ReduceNode,
         ctx: &dyn ReduceCtx,
-    ) -> VortexResult<Option<Box<dyn ReduceNode>>> {
+    ) -> VortexResult<Option<ReduceNodeRef>> {
         let child = node.child(0);
         if let Some(child_fn) = child.scalar_fn()
             && let Some(pack) = child_fn.as_opt::<Pack>()
@@ -155,9 +156,9 @@ impl VTable for GetItem {
 
             // Possibly mask the field if the pack is nullable
             if pack.nullability.is_nullable() {
-                field = ctx.create_node(
+                field = ctx.new_node(
                     Mask.bind(EmptyOptions),
-                    &[field, ctx.create_node(Literal.bind(true.into()), &[])?],
+                    &[field, ctx.new_node(Literal.bind(true.into()), &[])?],
                 )?;
             }
 

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

@@ -27,11 +27,14 @@ use crate::expr::ChildName;
 use crate::expr::ExecutionArgs;
 use crate::expr::ExprId;
 use crate::expr::Expression;
-use crate::expr::SimplifyCtx;
+use crate::expr::GetItem;
+use crate::expr::Pack;
+use crate::expr::PackOptions;
+use crate::expr::ReduceCtx;
+use crate::expr::ReduceNode;
+use crate::expr::ReduceNodeRef;
 use crate::expr::VTable;
 use crate::expr::VTableExt;
-use crate::expr::get_item;
-use crate::expr::pack;
 use crate::validity::Validity;
 
 /// Merge zero or more expressions that ALL return structs.
@@ -185,19 +188,19 @@ impl VTable for Merge {
         todo!()
     }
 
-    fn simplify(
+    fn reduce(
         &self,
         options: &Self::Options,
-        expr: &Expression,
-        ctx: &dyn SimplifyCtx,
-    ) -> VortexResult<Option<Expression>> {
-        let merge_dtype = ctx.return_dtype(expr)?;
-        let mut names = Vec::with_capacity(expr.children().len() * 2);
-        let mut children = Vec::with_capacity(expr.children().len() * 2);
+        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 expr.children().iter() {
-            let child_dtype = ctx.return_dtype(child)?;
+        for child in node.children() {
+            let child_dtype = child.node_dtype()?;
             if !child_dtype.is_struct() {
                 vortex_bail!(
                     "Merge child must return a non-nullable struct dtype, got {}",
@@ -227,15 +230,21 @@ impl VTable for Merge {
             }
         }
 
-        let expr = pack(
-            names
-                .into_iter()
-                .zip(children)
-                .map(|(name, child)| (name.clone(), get_item(name, child))),
-            merge_dtype.nullability(),
-        );
-
-        Ok(Some(expr))
+        let pack_children: Vec<_> = names
+            .iter()
+            .zip(children)
+            .map(|(name, child)| ctx.new_node(GetItem.bind(name.clone()), &[child]))
+            .try_collect()?;
+
+        let pack_expr = ctx.new_node(
+            Pack.bind(PackOptions {
+                names: FieldNames::from(names),
+                nullability: merge_dtype.nullability(),
+            }),
+            &pack_children,
+        )?;
+
+        Ok(Some(pack_expr))
     }
 
     fn is_null_sensitive(&self, _instance: &Self::Options) -> bool {
@@ -575,7 +584,7 @@ mod tests {
             DuplicateHandling::RightMost,
         );
 
-        let result = e.simplify(&dtype).unwrap();
+        let result = e.optimize_root(&dtype).unwrap();
 
         assert!(result.is::<Pack>());
         assert_eq!(

vortex-array/src/expr/mod.rs

@@ -31,13 +31,13 @@ mod expression;
 mod exprs;
 mod field;
 pub mod forms;
+mod optimize;
 mod options;
 pub mod proto;
 pub mod pruning;
 mod scalar_fn;
 pub mod session;
 mod signature;
-mod simplify;
 pub mod stats;
 pub mod transform;
 pub mod traversal;