Expression::optimize_recursive is idempotent

vortex-array/src/expr/optimize.rs

@@ -134,30 +134,47 @@ impl Expression {
     ) -> VortexResult<Option<Expression>> {
         // First optimize the root
         let mut current = self.try_optimize(cache)?;
+        let mut loop_counter = 0;
 
-        // Then recursively optimize children. The new children vector is only allocated once a
-        // child actually changes, so fully-optimized subtrees cost no allocations.
-        let expr = current.as_ref().unwrap_or(self);
-        let children = expr.children();
-        let mut new_children: Option<Vec<Expression>> = None;
-        for (idx, child) in children.iter().enumerate() {
-            if let Some(optimized) = child.try_optimize_recursive_inner(cache)? {
-                new_children
-                    .get_or_insert_with(|| children[..idx].to_vec())
-                    .push(optimized);
-            } else if let Some(new_children) = new_children.as_mut() {
-                new_children.push(child.clone());
+        // A root rewrite that runs after the children were optimized can introduce fresh,
+        // unoptimized children (for example `select(pack)` becomes `pack(get_item(pack))`).
+        // Loop until the root stops changing so those children get optimized as well. This makes
+        // the pass idempotent: a second `optimize_recursive` finds nothing left to do.
+        loop {
+            if loop_counter > 100 {
+                vortex_error::vortex_bail!(
+                    "Exceeded maximum recursive optimization iterations (possible infinite loop)"
+                );
             }
-        }
+            loop_counter += 1;
+
+            // Recursively optimize children. The new children vector is only allocated once a
+            // child actually changes, so fully-optimized subtrees cost no allocations.
+            let expr = current.as_ref().unwrap_or(self);
+            let children = expr.children();
+            let mut new_children: Option<Vec<Expression>> = None;
+            for (idx, child) in children.iter().enumerate() {
+                if let Some(optimized) = child.try_optimize_recursive_inner(cache)? {
+                    new_children
+                        .get_or_insert_with(|| children[..idx].to_vec())
+                        .push(optimized);
+                } else if let Some(new_children) = new_children.as_mut() {
+                    new_children.push(child.clone());
+                }
+            }
+
+            let Some(new_children) = new_children else {
+                return Ok(current);
+            };
 
-        if let Some(new_children) = new_children {
             let updated = expr.clone().with_children(new_children)?;
 
-            // After updating children, try to optimize root again
-            current = Some(updated.try_optimize(cache)?.unwrap_or(updated));
+            // After updating children, try to optimize root again.
+            match updated.try_optimize(cache)? {
+                None => return Ok(Some(updated)),
+                Some(reoptimized) => current = Some(reoptimized),
+            }
         }
-
-        Ok(current)
     }
 }
 
@@ -210,6 +227,7 @@ mod tests {
     use vortex_error::vortex_err;
 
     use crate::dtype::DType;
+    use crate::dtype::FieldName;
     use crate::dtype::Nullability;
     use crate::dtype::PType;
     use crate::dtype::StructFields;
@@ -218,8 +236,13 @@ mod tests {
     use crate::expr::get_item;
     use crate::expr::lit;
     use crate::expr::lt_eq;
+    use crate::expr::merge;
     use crate::expr::or;
+    use crate::expr::pack;
     use crate::expr::root;
+    use crate::expr::select;
+    use crate::expr::transform::replace;
+    use crate::expr::transform::replace_root_fields;
     use crate::scalar::Scalar;
     use crate::scalar_fn::fns::literal::Literal;
 
@@ -272,4 +295,42 @@ mod tests {
         assert_eq!(rhs, &Scalar::primitive(3.0f64, Nullability::NonNullable));
         Ok(())
     }
+
+    /// A root rewrite that fires only after the children were optimized (`select(pack)` becomes
+    /// `pack(get_item(pack))`) must not leave its fresh children unoptimized.
+    #[test]
+    fn optimize_recursive_is_idempotent() -> VortexResult<()> {
+        let field_names = (0..4)
+            .map(|idx| FieldName::from(format!("field_{idx}")))
+            .collect::<Vec<_>>();
+        let fields = StructFields::new(
+            field_names.clone().into(),
+            vec![DType::Primitive(PType::U64, Nullability::NonNullable); field_names.len()],
+        );
+        let scope = DType::Struct(fields.clone(), Nullability::NonNullable);
+
+        let expr = select(
+            field_names
+                .into_iter()
+                .chain(["input_row_idx".into(), "input_file_idx".into()])
+                .collect::<Vec<_>>(),
+            merge([
+                root(),
+                pack(
+                    [
+                        ("input_row_idx", lit(0_u64)),
+                        ("input_file_idx", lit(1_u64)),
+                    ],
+                    Nullability::NonNullable,
+                ),
+            ]),
+        );
+        let expr = replace(expr, &root(), replace_root_fields(root(), &fields));
+
+        let once = expr.optimize_recursive(&scope)?;
+        let twice = once.optimize_recursive(&scope)?;
+        assert_eq!(once, twice, "optimize is not idempotent:\n{once}\n{twice}");
+
+        Ok(())
+    }
 }