diff --git a/vortex-layout/src/plan/physical/optimizer/mod.rs b/vortex-layout/src/plan/physical/optimizer/mod.rs index c4841740572..9b928c80daf 100644 --- a/vortex-layout/src/plan/physical/optimizer/mod.rs +++ b/vortex-layout/src/plan/physical/optimizer/mod.rs @@ -29,6 +29,8 @@ 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 = PlanReduceRuleAdapter::new(EvalIdentityRule); @@ -36,7 +38,11 @@ static EVAL_IDENTITY_RULE: PlanReduceRuleAdapter = static LEGACY_STATS_RULE: PlanReduceRuleAdapter = 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 = PlanParentReduceRuleAdapter::new(ExpressionConcatRule); @@ -51,6 +57,7 @@ static ZONED_LEGACY_STATS_RULE: PlanParentReduceRuleAdapter for EvalIdentityRule { } } } + +/// Eval(expr1, Eval(expr2)) -> Eval(expr1 . expr2) +#[derive(Debug)] +pub(crate) struct EvalEvalRule; +impl PlanParentReduceRule for EvalEvalRule { + type Parent = Eval; + fn reduce_parent( + &self, + child: &Plan, + parent: &Plan, + _: usize, + ) -> VortexResult> { + 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 for EvalPackRule { + fn reduce(&self, plan: &Plan) -> VortexResult> { + let Some(expr) = plan.expression.as_opt::() 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())) + } +} diff --git a/vortex-layout/src/plan/physical/plans/mod.rs b/vortex-layout/src/plan/physical/plans/mod.rs index fe5a4eae225..8dba359d635 100644 --- a/vortex-layout/src/plan/physical/plans/mod.rs +++ b/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; diff --git a/vortex-layout/src/plan/physical/tests.rs b/vortex-layout/src/plan/physical/tests.rs index 44705753b3e..7ee6647bd47 100644 --- a/vortex-layout/src/plan/physical/tests.rs +++ b/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::()); + assert!(zones.is::()); 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::() .ok_or_else(|| vortex_err!("legacy stats predicate did not become a pruning plan"))?; assert!(pruning.is_pruning()); - assert!(pruning.zones_plan()?.is::()); + assert!(pruning.zones_plan()?.is::()); 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(()) +}