Add plan optimizer rules and push expressions

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

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;

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,62 @@
+// 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::RowIdx;
+use super::Take;
+use super::plans::EvalIdentityRule;
+use super::plans::ExpressionConcatRule;
+use super::plans::ExpressionPackRule;
+use super::plans::ExpressionRowIdxRule;
+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_ROW_IDX_RULE: PlanParentReduceRuleAdapter<RowIdx, ExpressionRowIdxRule> =
+    PlanParentReduceRuleAdapter::new(ExpressionRowIdxRule);
+static EXPRESSION_PACK_RULE: PlanParentReduceRuleAdapter<Pack, ExpressionPackRule> =
+    PlanParentReduceRuleAdapter::new(ExpressionPackRule);
+
+static PARENT_RULES: PlanParentRuleSet = PlanParentRuleSet::new(&[
+    &EXPRESSION_CONCAT_RULE,
+    &EXPRESSION_TAKE_RULE,
+    &EXPRESSION_ROW_IDX_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)
+}

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

@@ -0,0 +1,257 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Typed and type-erased interfaces for plan rewrites.
+
+use std::any::type_name;
+use std::fmt::Debug;
+use std::marker::PhantomData;
+
+use vortex_error::VortexResult;
+
+use crate::plan::Plan;
+use crate::plan::PlanRef;
+use crate::plan::PlanVTable;
+
+/// A rewrite over one concrete plan operator.
+///
+/// Rules return one rewrite without recursively optimizing the replacement. The plan optimizer
+/// owns traversal and drives further rewrites.
+pub trait PlanReduceRule<P: PlanVTable>: Debug + Send + Sync + 'static {
+    /// Attempts to replace `plan`.
+    fn reduce(&self, plan: &Plan<P>) -> VortexResult<Option<PlanRef>>;
+}
+
+/// Type-erased interface used by [`PlanRuleSet`].
+pub trait DynPlanReduceRule: Debug + Send + Sync + 'static {
+    /// Returns whether this rule supports the concrete plan operator.
+    fn matches(&self, plan: &PlanRef) -> bool;
+
+    /// Attempts to replace `plan`.
+    fn reduce(&self, plan: &PlanRef) -> VortexResult<Option<PlanRef>>;
+}
+
+/// Bridges a typed [`PlanReduceRule`] to a type-erased static registry.
+pub struct PlanReduceRuleAdapter<P, R> {
+    rule: R,
+    _plan: PhantomData<fn() -> P>,
+}
+
+impl<P, R> PlanReduceRuleAdapter<P, R> {
+    /// Creates an adapter for a typed plan rule.
+    pub const fn new(rule: R) -> Self {
+        Self {
+            rule,
+            _plan: PhantomData,
+        }
+    }
+}
+
+impl<P, R> Debug for PlanReduceRuleAdapter<P, R>
+where
+    P: PlanVTable,
+    R: PlanReduceRule<P>,
+{
+    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        formatter
+            .debug_struct("PlanReduceRuleAdapter")
+            .field("plan", &type_name::<P>())
+            .field("rule", &self.rule)
+            .finish()
+    }
+}
+
+impl<P, R> DynPlanReduceRule for PlanReduceRuleAdapter<P, R>
+where
+    P: PlanVTable,
+    R: PlanReduceRule<P>,
+{
+    fn matches(&self, plan: &PlanRef) -> bool {
+        plan.is::<P>()
+    }
+
+    fn reduce(&self, plan: &PlanRef) -> VortexResult<Option<PlanRef>> {
+        let Some(plan) = plan.as_opt::<P>() else {
+            return Ok(None);
+        };
+        self.rule.reduce(plan)
+    }
+}
+
+/// An ordered static collection of single-plan rewrite rules.
+pub struct PlanRuleSet {
+    rules: &'static [&'static dyn DynPlanReduceRule],
+}
+
+impl PlanRuleSet {
+    /// Creates a rule set whose first successful rewrite wins.
+    pub const fn new(rules: &'static [&'static dyn DynPlanReduceRule]) -> Self {
+        Self { rules }
+    }
+
+    /// Evaluates rules registered for the concrete plan operator.
+    pub fn evaluate(&self, plan: &PlanRef) -> VortexResult<Option<PlanRef>> {
+        for rule in self.rules {
+            if !rule.matches(plan) {
+                continue;
+            }
+            let Some(reduced) = rule.reduce(plan)? else {
+                continue;
+            };
+
+            #[cfg(debug_assertions)]
+            {
+                vortex_error::vortex_ensure!(
+                    reduced.row_count() == plan.row_count(),
+                    "Plan rewrite from {rule:?} changed row count from {} to {}",
+                    plan.row_count(),
+                    reduced.row_count()
+                );
+                vortex_error::vortex_ensure!(
+                    reduced.dtype() == plan.dtype(),
+                    "Plan rewrite from {rule:?} changed dtype from {} to {}",
+                    plan.dtype(),
+                    reduced.dtype()
+                );
+            }
+
+            return Ok(Some(reduced));
+        }
+        Ok(None)
+    }
+}
+
+/// A metadata-only rewrite where a child plan rewrites its parent plan.
+///
+/// Rules return one rewrite without recursively optimizing the replacement. The plan optimizer
+/// owns traversal and drives further rewrites.
+pub trait PlanParentReduceRule<C: PlanVTable>: Debug + Send + Sync + 'static {
+    /// The concrete parent operator matched by this rule.
+    type Parent: PlanVTable;
+
+    /// Attempts to replace `parent` based on its child at `child_idx`.
+    fn reduce_parent(
+        &self,
+        child: &Plan<C>,
+        parent: &Plan<Self::Parent>,
+        child_idx: usize,
+    ) -> VortexResult<Option<PlanRef>>;
+}
+
+/// Type-erased interface used by [`PlanParentRuleSet`].
+pub trait DynPlanParentReduceRule: Debug + Send + Sync + 'static {
+    /// Returns whether this rule supports the concrete child and parent operators.
+    fn matches(&self, child: &PlanRef, parent: &PlanRef) -> bool;
+
+    /// Attempts to replace `parent` based on `child` at `child_idx`.
+    fn reduce_parent(
+        &self,
+        child: &PlanRef,
+        parent: &PlanRef,
+        child_idx: usize,
+    ) -> VortexResult<Option<PlanRef>>;
+}
+
+/// Bridges a typed [`PlanParentReduceRule`] to a type-erased static registry.
+pub struct PlanParentReduceRuleAdapter<C, R> {
+    rule: R,
+    _child: PhantomData<fn() -> C>,
+}
+
+impl<C, R> PlanParentReduceRuleAdapter<C, R> {
+    /// Creates an adapter for a typed parent-child rule.
+    pub const fn new(rule: R) -> Self {
+        Self {
+            rule,
+            _child: PhantomData,
+        }
+    }
+}
+
+impl<C, R> Debug for PlanParentReduceRuleAdapter<C, R>
+where
+    C: PlanVTable,
+    R: PlanParentReduceRule<C>,
+{
+    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        formatter
+            .debug_struct("PlanParentReduceRuleAdapter")
+            .field("parent", &type_name::<R::Parent>())
+            .field("child", &type_name::<C>())
+            .field("rule", &self.rule)
+            .finish()
+    }
+}
+
+impl<C, R> DynPlanParentReduceRule for PlanParentReduceRuleAdapter<C, R>
+where
+    C: PlanVTable,
+    R: PlanParentReduceRule<C>,
+{
+    fn matches(&self, child: &PlanRef, parent: &PlanRef) -> bool {
+        child.is::<C>() && parent.is::<R::Parent>()
+    }
+
+    fn reduce_parent(
+        &self,
+        child: &PlanRef,
+        parent: &PlanRef,
+        child_idx: usize,
+    ) -> VortexResult<Option<PlanRef>> {
+        let Some(child) = child.as_opt::<C>() else {
+            return Ok(None);
+        };
+        let Some(parent) = parent.as_opt::<R::Parent>() else {
+            return Ok(None);
+        };
+        self.rule.reduce_parent(child, parent, child_idx)
+    }
+}
+
+/// An ordered static collection of parent-child plan rewrite rules.
+pub struct PlanParentRuleSet {
+    rules: &'static [&'static dyn DynPlanParentReduceRule],
+}
+
+impl PlanParentRuleSet {
+    /// Creates a rule set whose first successful rewrite wins.
+    pub const fn new(rules: &'static [&'static dyn DynPlanParentReduceRule]) -> Self {
+        Self { rules }
+    }
+
+    /// Evaluates rules registered for the concrete `(parent, child)` pair.
+    pub fn evaluate(
+        &self,
+        child: &PlanRef,
+        parent: &PlanRef,
+        child_idx: usize,
+    ) -> VortexResult<Option<PlanRef>> {
+        for rule in self.rules {
+            if !rule.matches(child, parent) {
+                continue;
+            }
+            let Some(reduced) = rule.reduce_parent(child, parent, child_idx)? else {
+                continue;
+            };
+
+            #[cfg(debug_assertions)]
+            {
+                vortex_error::vortex_ensure!(
+                    reduced.row_count() == parent.row_count(),
+                    "Plan rewrite from {rule:?} changed row count from {} to {}",
+                    parent.row_count(),
+                    reduced.row_count()
+                );
+                vortex_error::vortex_ensure!(
+                    reduced.dtype() == parent.dtype(),
+                    "Plan rewrite from {rule:?} changed dtype from {} to {}",
+                    parent.dtype(),
+                    reduced.dtype()
+                );
+            }
+
+            return Ok(Some(reduced));
+        }
+        Ok(None)
+    }
+}