Add plan optimizer rules and push expressions

vortex-layout/src/plan/mod.rs

@@ -11,6 +11,7 @@ mod children;
 mod display;
 mod lower;
 mod optimize;
+pub mod optimizer;
 mod plans;
 mod typed;
 mod vtable;
@@ -38,12 +39,13 @@ pub use plans::PackPlan;
 pub use plans::RowIdx;
 pub use plans::RowIdxData;
 pub use plans::RowIdxPlan;
-pub use plans::RowIdxPlanMetadata;
 pub use plans::SegmentScan;
 pub use plans::SegmentScanData;
 pub use plans::SegmentScanPlan;
 pub use plans::Take;
 pub use plans::TakePlan;
+pub use plans::plan_row_idx_expression;
+pub use plans::row_idx_dtype;
 pub use typed::DynPlan;
 pub use typed::Plan;
 pub use typed::PlanParts;

vortex-layout/src/plan/optimize.rs

@@ -1,15 +1,35 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-//! Generic bottom-up optimization over physical plans.
+//! Plan optimization.
+//!
+//! The optimizer applies static rewrites top-down, optimizes children, then retries rewrites
+//! exposed by the optimized children.
 
 use vortex_error::VortexResult;
 
-use crate::plan::Eval;
 use crate::plan::PlanRef;
+use crate::plan::optimizer::reduce_parent;
+use crate::plan::optimizer::reduce_plan;
+
+fn reduce(plan: &PlanRef) -> VortexResult<Option<PlanRef>> {
+    if let Some(rewritten) = reduce_plan(plan)? {
+        return Ok(Some(rewritten));
+    }
+    for child_idx in 0..plan.child_count() {
+        if let Some(rewritten) = reduce_parent(plan, child_idx)? {
+            return Ok(Some(rewritten));
+        }
+    }
+    Ok(None)
+}
 
 /// Optimizes `plan`, preserving its dtype and row domain.
 pub fn optimize(plan: PlanRef) -> VortexResult<PlanRef> {
+    if let Some(rewritten) = reduce(&plan)? {
+        return optimize(rewritten);
+    }
+
     let mut children = Vec::with_capacity(plan.child_count());
     let mut changed = false;
     for child in plan.children().iter() {
@@ -19,17 +39,13 @@ pub fn optimize(plan: PlanRef) -> VortexResult<PlanRef> {
         children.push(optimized);
     }
 
-    let plan = if changed {
-        plan.with_children(children)?
-    } else {
-        plan
-    };
-
-    let Some(eval) = plan.as_opt::<Eval>() else {
+    if !changed {
         return Ok(plan);
-    };
-    if eval.expression().is_root() {
-        return eval.child_plan();
+    }
+
+    let plan = plan.with_children(children)?;
+    if let Some(rewritten) = reduce(&plan)? {
+        return optimize(rewritten);
     }
     Ok(plan)
 }

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

@@ -0,0 +1,57 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Static rewrite rules for physical plans.
+
+mod rules;
+
+pub use rules::DynPlanParentReduceRule;
+pub use rules::DynPlanReduceRule;
+pub use rules::PlanParentReduceRule;
+pub use rules::PlanParentReduceRuleAdapter;
+pub use rules::PlanParentRuleSet;
+pub use rules::PlanReduceRule;
+pub use rules::PlanReduceRuleAdapter;
+pub use rules::PlanRuleSet;
+use vortex_error::VortexResult;
+
+use super::Concat;
+use super::Eval;
+use super::Pack;
+use super::PlanRef;
+use super::Take;
+use super::plans::EvalIdentityRule;
+use super::plans::ExpressionConcatRule;
+use super::plans::ExpressionPackRule;
+use super::plans::ExpressionTakeRule;
+
+static EVAL_IDENTITY_RULE: PlanReduceRuleAdapter<Eval, EvalIdentityRule> =
+    PlanReduceRuleAdapter::new(EvalIdentityRule);
+
+static PLAN_RULES: PlanRuleSet = PlanRuleSet::new(&[&EVAL_IDENTITY_RULE]);
+
+static EXPRESSION_CONCAT_RULE: PlanParentReduceRuleAdapter<Concat, ExpressionConcatRule> =
+    PlanParentReduceRuleAdapter::new(ExpressionConcatRule);
+static EXPRESSION_TAKE_RULE: PlanParentReduceRuleAdapter<Take, ExpressionTakeRule> =
+    PlanParentReduceRuleAdapter::new(ExpressionTakeRule);
+static EXPRESSION_PACK_RULE: PlanParentReduceRuleAdapter<Pack, ExpressionPackRule> =
+    PlanParentReduceRuleAdapter::new(ExpressionPackRule);
+
+static PARENT_RULES: PlanParentRuleSet = PlanParentRuleSet::new(&[
+    &EXPRESSION_CONCAT_RULE,
+    &EXPRESSION_TAKE_RULE,
+    &EXPRESSION_PACK_RULE,
+]);
+
+/// Attempts a static rewrite for `plan`.
+pub(crate) fn reduce_plan(plan: &PlanRef) -> VortexResult<Option<PlanRef>> {
+    PLAN_RULES.evaluate(plan)
+}
+
+/// Attempts a static rewrite for `parent` and its child at `child_idx`.
+pub(crate) fn reduce_parent(parent: &PlanRef, child_idx: usize) -> VortexResult<Option<PlanRef>> {
+    let Some(child) = parent.child(child_idx)? else {
+        return Ok(None);
+    };
+    PARENT_RULES.evaluate(&child, parent, child_idx)
+}