From adb3fe7995ac77179c9f2b65c30859a65279828e Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Fri, 23 Oct 2026 15:44:14 +0100 Subject: [PATCH 1/3] Add Eval(Eval()) and Eval(pack()) reduce rules Signed-off-by: Mikhail Kot --- .../src/plan/physical/optimizer/mod.rs | 10 ++- vortex-layout/src/plan/physical/plans/eval.rs | 67 +++++++++++++++++++ vortex-layout/src/plan/physical/plans/mod.rs | 4 +- vortex-layout/src/plan/physical/tests.rs | 33 +++++++++ 4 files changed, 112 insertions(+), 2 deletions(-) diff --git a/vortex-layout/src/plan/physical/optimizer/mod.rs b/vortex-layout/src/plan/physical/optimizer/mod.rs index c4841740572..c15636de247 100644 --- a/vortex-layout/src/plan/physical/optimizer/mod.rs +++ b/vortex-layout/src/plan/physical/optimizer/mod.rs @@ -15,6 +15,9 @@ pub use rules::PlanReduceRuleAdapter; pub use rules::PlanRuleSet; use vortex_error::VortexResult; +use crate::plan::physical::plans::EvalEvalRule; +use crate::plan::physical::plans::EvalPackRule; + use super::Concat; use super::Eval; use super::LegacyStats; @@ -36,7 +39,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 +58,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, + }) + })?; + let expr = expr.into_inner().optimize_recursive()?; + 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..46f8428f368 100644 --- a/vortex-layout/src/plan/physical/plans/mod.rs +++ b/vortex-layout/src/plan/physical/plans/mod.rs @@ -16,10 +16,12 @@ mod zoned; pub use concat::Concat; pub use concat::ConcatData; pub use concat::ConcatPlan; -pub(crate) use concat::ExpressionConcatRule; pub use eval::Eval; pub use eval::EvalData; +pub(crate) use concat::ExpressionConcatRule; +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..b374064cb7f 100644 --- a/vortex-layout/src/plan/physical/tests.rs +++ b/vortex-layout/src/plan/physical/tests.rs @@ -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.clone())]); + 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(()) +} From a33d457745c4206db6f31fbdafca627acae2e2ec Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Fri, 23 Oct 2026 15:50:02 +0100 Subject: [PATCH 2/3] fix Signed-off-by: Mikhail Kot --- vortex-layout/src/plan/physical/optimizer/mod.rs | 5 ++--- vortex-layout/src/plan/physical/plans/mod.rs | 2 +- vortex-layout/src/plan/physical/tests.rs | 2 +- 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/vortex-layout/src/plan/physical/optimizer/mod.rs b/vortex-layout/src/plan/physical/optimizer/mod.rs index c15636de247..9b928c80daf 100644 --- a/vortex-layout/src/plan/physical/optimizer/mod.rs +++ b/vortex-layout/src/plan/physical/optimizer/mod.rs @@ -15,9 +15,6 @@ pub use rules::PlanReduceRuleAdapter; pub use rules::PlanRuleSet; use vortex_error::VortexResult; -use crate::plan::physical::plans::EvalEvalRule; -use crate::plan::physical::plans::EvalPackRule; - use super::Concat; use super::Eval; use super::LegacyStats; @@ -32,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); diff --git a/vortex-layout/src/plan/physical/plans/mod.rs b/vortex-layout/src/plan/physical/plans/mod.rs index 46f8428f368..8dba359d635 100644 --- a/vortex-layout/src/plan/physical/plans/mod.rs +++ b/vortex-layout/src/plan/physical/plans/mod.rs @@ -16,9 +16,9 @@ mod zoned; pub use concat::Concat; pub use concat::ConcatData; pub use concat::ConcatPlan; +pub(crate) use concat::ExpressionConcatRule; pub use eval::Eval; pub use eval::EvalData; -pub(crate) use concat::ExpressionConcatRule; pub(crate) use eval::EvalEvalRule; pub(crate) use eval::EvalIdentityRule; pub(crate) use eval::EvalPackRule; diff --git a/vortex-layout/src/plan/physical/tests.rs b/vortex-layout/src/plan/physical/tests.rs index b374064cb7f..b376fd37083 100644 --- a/vortex-layout/src/plan/physical/tests.rs +++ b/vortex-layout/src/plan/physical/tests.rs @@ -1728,7 +1728,7 @@ fn zoned_pruning_plan_with_new_zones_proves_them_again() -> VortexResult<()> { 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.clone())]); + 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(); From d1be7be60ccf26abf7498b8aa6b28eb2a89ae7d1 Mon Sep 17 00:00:00 2001 From: Mikhail Kot Date: Fri, 23 Oct 2026 15:59:50 +0100 Subject: [PATCH 3/3] fix Signed-off-by: Mikhail Kot --- vortex-layout/src/plan/physical/plans/eval.rs | 3 ++- vortex-layout/src/plan/physical/tests.rs | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/vortex-layout/src/plan/physical/plans/eval.rs b/vortex-layout/src/plan/physical/plans/eval.rs index a4eb97be2f0..b6a2f1ff620 100644 --- a/vortex-layout/src/plan/physical/plans/eval.rs +++ b/vortex-layout/src/plan/physical/plans/eval.rs @@ -185,7 +185,8 @@ impl PlanParentReduceRule for EvalEvalRule { changed: true, }) })?; - let expr = expr.into_inner().optimize_recursive()?; + // 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())) } diff --git a/vortex-layout/src/plan/physical/tests.rs b/vortex-layout/src/plan/physical/tests.rs index b376fd37083..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(()) }