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/expr/exprs/get_item.rs

@@ -329,7 +329,7 @@ mod tests {
         let get_item_expr = get_item("b", pack_expr);
 
         let result = get_item_expr
-            .simplify(&DType::Struct(StructFields::empty(), NonNullable))
+            .optimize_recursive(&DType::Struct(StructFields::empty(), NonNullable))
             .unwrap();
 
         assert_eq!(result, lit(2));
@@ -345,7 +345,7 @@ mod tests {
 
         let dtype = DType::Primitive(PType::I32, NonNullable);
 
-        let result = get_z.simplify(&dtype).unwrap();
+        let result = get_z.optimize_recursive(&dtype).unwrap();
         assert_eq!(result, lit(4));
     }
 
@@ -365,7 +365,7 @@ mod tests {
 
         let dtype = DType::Primitive(PType::I32, NonNullable);
 
-        let result = get_final.simplify(&dtype).unwrap();
+        let result = get_final.optimize_recursive(&dtype).unwrap();
         assert_eq!(result, lit(42));
     }
 
@@ -380,7 +380,7 @@ mod tests {
 
         let dtype = DType::Primitive(PType::I32, NonNullable);
 
-        let result = get_result.simplify(&dtype).unwrap();
+        let result = get_result.optimize_recursive(&dtype).unwrap();
         let expected = checked_add(lit(1), lit(10));
         assert_eq!(&result, &expected);
     }

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

@@ -584,7 +584,7 @@ mod tests {
             DuplicateHandling::RightMost,
         );
 
-        let result = e.optimize_root(&dtype).unwrap();
+        let result = e.optimize(&dtype).unwrap();
 
         assert!(result.is::<Pack>());
         assert_eq!(

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

@@ -461,7 +461,7 @@ mod tests {
         );
         let e = select(["a", "b"], root());
 
-        let result = e.simplify(&dtype).unwrap();
+        let result = e.optimize_recursive(&dtype).unwrap();
 
         assert!(result.is::<Pack>());
         assert!(result.return_dtype(&dtype).unwrap().is_nullable());
@@ -480,7 +480,7 @@ mod tests {
         );
         let e = select_exclude(["c"], root());
 
-        let result = e.simplify(&dtype).unwrap();
+        let result = e.optimize_recursive(&dtype).unwrap();
 
         assert!(result.is::<Pack>());
 

vortex-array/src/expr/optimize.rs

@@ -28,15 +28,15 @@ impl Expression {
     /// 1. `simplify_untyped` - type-independent simplifications
     /// 2. `simplify` - type-aware simplifications
     /// 3. `reduce` - abstract reduction rules via `ReduceNode`/`ReduceCtx`
-    pub fn optimize_root(&self, scope: &DType) -> VortexResult<Expression> {
+    pub fn optimize(&self, scope: &DType) -> VortexResult<Expression> {
         Ok(self
             .clone()
-            .try_optimize_root(scope)?
+            .try_optimize(scope)?
             .unwrap_or_else(|| self.clone()))
     }
 
     /// Try to optimize the root expression node only, returning None if no optimizations applied.
-    pub fn try_optimize_root(&self, scope: &DType) -> VortexResult<Option<Expression>> {
+    pub fn try_optimize(&self, scope: &DType) -> VortexResult<Option<Expression>> {
         let cache = SimplifyCache {
             scope,
             dtype_cache: RefCell::new(HashMap::new()),
@@ -118,7 +118,7 @@ impl Expression {
         let mut any_optimizations = false;
 
         // First optimize the root
-        if let Some(optimized) = current.clone().try_optimize_root(scope)? {
+        if let Some(optimized) = current.clone().try_optimize(scope)? {
             current = optimized;
             any_optimizations = true;
         }
@@ -140,7 +140,7 @@ impl Expression {
             any_optimizations = true;
 
             // After updating children, try to optimize root again
-            if let Some(optimized) = current.clone().try_optimize_root(scope)? {
+            if let Some(optimized) = current.clone().try_optimize(scope)? {
                 current = optimized;
             }
         }
@@ -282,8 +282,9 @@ impl ReduceCtx for ExpressionReduceCtx {
                 .iter()
                 .map(|c| {
                     c.as_any()
-                        .downcast_ref::<Expression>()
-                        .vortex_expect("ReduceNode not an Expression")
+                        .downcast_ref::<ExpressionReduceNode>()
+                        .vortex_expect("ReduceNode not an ExpressionReduceNode")
+                        .expression
                         .clone()
                 })
                 .collect::<Vec<_>>(),

vortex-array/src/expr/transform/partition.rs

@@ -284,7 +284,7 @@ mod tests {
 
         let split_a = partitioned.find_partition(&"a".into()).unwrap();
         assert_eq!(
-            &split_a.simplify(&dtype).unwrap(),
+            &split_a.optimize_recursive(&dtype).unwrap(),
             &pack(
                 [
                     ("a_0", get_item("x", get_item("a", root()))),
@@ -326,7 +326,7 @@ mod tests {
             get_item("y", get_item("a", root())),
             select(["a", "b"], root()),
         );
-        let expr = expr.simplify(&dtype).unwrap();
+        let expr = expr.optimize_recursive(&dtype).unwrap();
         let partitioned = partition(expr, &dtype, annotate_scope_access(fields)).unwrap();
 
         // One for id.a and id.b

vortex-array/src/optimizer/mod.rs

@@ -28,7 +28,6 @@ impl ArrayOptimizer for ArrayRef {
         Ok(self.clone().try_optimize()?.unwrap_or_else(|| self.clone()))
     }
 
-    #[expect(clippy::cognitive_complexity)]
     fn try_optimize(&self) -> VortexResult<Option<ArrayRef>> {
         let mut current_array = self.clone();
         let mut any_optimizations = false;