Add Eval(Eval()) and Eval(pack()) reduce rules

vortex-layout/src/plan/physical/optimizer/mod.rs

@@ -29,14 +29,20 @@ use super::plans::ExpressionTakeRule;
 use super::plans::ExpressionZonedRule;
 use super::plans::LegacyStatsRule;
 use super::plans::ZonedLegacyStatsRule;
+use crate::plan::physical::plans::EvalEvalRule;
+use crate::plan::physical::plans::EvalPackRule;
 
 static EVAL_IDENTITY_RULE: PlanReduceRuleAdapter<Eval, EvalIdentityRule> =
     PlanReduceRuleAdapter::new(EvalIdentityRule);
 
 static LEGACY_STATS_RULE: PlanReduceRuleAdapter<LegacyStats, LegacyStatsRule> =
     PlanReduceRuleAdapter::new(LegacyStatsRule);
 
-static PLAN_RULES: PlanRuleSet = PlanRuleSet::new(&[&EVAL_IDENTITY_RULE, &LEGACY_STATS_RULE]);
+static PLAN_RULES: PlanRuleSet = PlanRuleSet::new(&[
+    &EVAL_IDENTITY_RULE,
+    &PlanReduceRuleAdapter::new(EvalPackRule),
+    &LEGACY_STATS_RULE,
+]);
 
 static EXPRESSION_CONCAT_RULE: PlanParentReduceRuleAdapter<Concat, ExpressionConcatRule> =
     PlanParentReduceRuleAdapter::new(ExpressionConcatRule);
@@ -51,6 +57,7 @@ static ZONED_LEGACY_STATS_RULE: PlanParentReduceRuleAdapter<LegacyStats, ZonedLe
     PlanParentReduceRuleAdapter::new(ZonedLegacyStatsRule);
 
 static PARENT_RULES: PlanParentRuleSet = PlanParentRuleSet::new(&[
+    &PlanParentReduceRuleAdapter::new(EvalEvalRule),
     &EXPRESSION_CONCAT_RULE,
     &EXPRESSION_TAKE_RULE,
     &EXPRESSION_PACK_RULE,

vortex-layout/src/plan/physical/plans/eval.rs

@@ -6,12 +6,20 @@ use std::fmt;
 use std::ops::Range;
 
 use vortex_array::EmptyMetadata;
+use vortex_array::dtype::Nullability;
+use vortex_array::dtype::StructFields;
 use vortex_array::expr::BoundExpression;
+use vortex_array::expr::traversal::NodeExt;
+use vortex_array::expr::traversal::Transformed;
+use vortex_array::expr::traversal::TraversalOrder;
+use vortex_array::scalar_fn::fns::pack::Pack;
+use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_mask::Mask;
 use vortex_session::registry::CachedId;
 
+use crate::plan::physical::PackPlan;
 use crate::plan::physical::Plan;
 use crate::plan::physical::PlanChildren;
 use crate::plan::physical::PlanId;
@@ -20,6 +28,7 @@ use crate::plan::physical::PlanRef;
 use crate::plan::physical::PlanVTable;
 use crate::plan::physical::check_child_count;
 use crate::plan::physical::operators::EvalStage;
+use crate::plan::physical::optimizer::PlanParentReduceRule;
 use crate::plan::physical::optimizer::PlanReduceRule;
 use crate::plan::physical::pipeline::Chain;
 use crate::plan::physical::pipeline::Compiler;
@@ -153,3 +162,62 @@ impl PlanReduceRule<Eval> for EvalIdentityRule {
         }
     }
 }
+
+/// Eval(expr1, Eval(expr2)) -> Eval(expr1 . expr2)
+#[derive(Debug)]
+pub(crate) struct EvalEvalRule;
+impl PlanParentReduceRule<Eval> for EvalEvalRule {
+    type Parent = Eval;
+    fn reduce_parent(
+        &self,
+        child: &Plan<Eval>,
+        parent: &Plan<Eval>,
+        _: usize,
+    ) -> VortexResult<Option<PlanRef>> {
+        let child_expr = child.expression();
+        let expr = parent.expression().clone().transform_down(|expr| {
+            if !expr.is_root() {
+                return Ok(Transformed::no(expr));
+            }
+            Ok(Transformed {
+                value: child_expr.clone(),
+                order: TraversalOrder::Skip,
+                changed: true,
+            })
+        })?;
+        // We want to fold repeated expressions like ($.a).b -> $.a.b.
+        let expr = expr.into_inner().optimize()?;
+        let plan = EvalPlan::try_new(expr, child.child_plan()?)?;
+        Ok(Some(plan.into_plan()))
+    }
+}
+
+/// Eval(pack({0: expr})) -> PackPlan(expr) for a non-nullable pack
+#[derive(Debug)]
+pub(crate) struct EvalPackRule;
+impl PlanReduceRule<Eval> for EvalPackRule {
+    fn reduce(&self, plan: &Plan<Eval>) -> VortexResult<Option<PlanRef>> {
+        let Some(expr) = plan.expression.as_opt::<Pack>() else {
+            return Ok(None);
+        };
+        if expr.nullability != Nullability::NonNullable {
+            return Ok(None);
+        }
+        let [child_expr] = plan.expression.children() else {
+            return Ok(None);
+        };
+        let fields = StructFields::new(expr.names.clone(), vec![child_expr.dtype().clone()]);
+
+        let child_plan = plan.child(0)?.vortex_expect("no child plan");
+        let child_plan = EvalPlan::try_new(child_expr.clone(), child_plan)?.into_plan();
+
+        let plan = PackPlan::try_new(
+            fields,
+            Nullability::NonNullable,
+            plan.row_count(),
+            vec![child_plan],
+            None,
+        )?;
+        Ok(Some(plan.into_plan()))
+    }
+}

vortex-layout/src/plan/physical/plans/mod.rs

@@ -19,7 +19,9 @@ pub use concat::ConcatPlan;
 pub(crate) use concat::ExpressionConcatRule;
 pub use eval::Eval;
 pub use eval::EvalData;
+pub(crate) use eval::EvalEvalRule;
 pub(crate) use eval::EvalIdentityRule;
+pub(crate) use eval::EvalPackRule;
 pub use eval::EvalPlan;
 pub use filter::Filter;
 pub use filter::FilterData;

vortex-layout/src/plan/physical/tests.rs

@@ -1498,7 +1498,7 @@ fn legacy_stats_rewrite_to_zoned_over_converted_zones() -> VortexResult<()> {
         vec![max.to_string()]
     );
     let zones = zoned.zones_plan()?;
-    assert!(zones.is::<Eval>());
+    assert!(zones.is::<Pack>());
     assert_eq!(zones.dtype(), &zones_dtype);
 
     let falsifier = gt(root(), lit(5_i32))
@@ -1510,7 +1510,7 @@ fn legacy_stats_rewrite_to_zoned_over_converted_zones() -> VortexResult<()> {
         .as_opt::<Zoned>()
         .ok_or_else(|| vortex_err!("legacy stats predicate did not become a pruning plan"))?;
     assert!(pruning.is_pruning());
-    assert!(pruning.zones_plan()?.is::<Eval>());
+    assert!(pruning.zones_plan()?.is::<Pack>());
     Ok(())
 }
 
@@ -1722,3 +1722,36 @@ fn zoned_pruning_plan_with_new_zones_proves_them_again() -> VortexResult<()> {
     assert_eq!(proof.get(&falsifier, &SESSION)?, Some(original));
     Ok(())
 }
+
+/// Eval($ > 5, Eval($.a, (...))) -> Eval($.a > 5, (...))
+#[test]
+fn test_eval_of_eval() -> VortexResult<()> {
+    let dtype = primitive(PType::I32, Nullability::NonNullable);
+    let values = vec![flat(3, dtype.clone(), 2)];
+    let fields = StructFields::from_iter([("a", dtype)]);
+    let struct_dtype = DType::Struct(fields, Nullability::NonNullable);
+    let layout = StructLayout::new(3, struct_dtype, values).into_layout();
+    let plan = make_eval(get_item("a", root()), make_plan(layout)?)?.into_plan();
+    let plan = make_eval(gt(root(), lit(5_i32)), plan)?.into_plan();
+    insta::assert_snapshot!(optimize(plan)?.display_tree(), @"
+    root: vortex.plan.eval(bool, rows=3) expr=($ > 5i32)
+      child: vortex.plan.filter(i32, rows=3)
+        child: vortex.plan.segment_scan(i32, rows=3)
+    ");
+    Ok(())
+}
+
+/// Eval(pack({0: $ > 0}, (...))) -> Pack({0: $ > 0}, (...))
+#[test]
+fn test_eval_of_pack() -> VortexResult<()> {
+    let expr = pack([("0", gt(root(), lit(0_i32)))], Nullability::NonNullable);
+    let plan = make_plan(flat(3, primitive(PType::I32, Nullability::NonNullable), 0))?;
+    let plan = make_eval(expr, plan)?.into_plan();
+    insta::assert_snapshot!(optimize(plan)?.display_tree(), @"
+    root: vortex.plan.pack({0=bool}, rows=3)
+      0: vortex.plan.eval(bool, rows=3) expr=($ > 0i32)
+        child: vortex.plan.filter(i32, rows=3)
+          child: vortex.plan.segment_scan(i32, rows=3)
+    ");
+    Ok(())
+}