Add plan optimizer rules and push expressions

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

Author: joseph-isaacs

vortex-array/src/expr/transform/bound_partition.rs

@@ -28,6 +28,7 @@ use crate::expr::traversal::NodeExt;
 use crate::expr::traversal::NodeRewriter;
 use crate::expr::traversal::Transformed;
 use crate::expr::traversal::TraversalOrder;
+use crate::scalar_fn::fns::get_item::GetItem;
 
 /// Partition an expression into sub-expressions that are uniquely associated with an annotation.
 /// A root expression is also returned that can be used to recombine the results of the partitions
@@ -300,6 +301,46 @@ fn replace_root_dtype(expr: BoundExpression, root_dtype: DType) -> VortexResult<
         .into_inner())
 }
 
+/// Rewrites partition accessors in `expression` to read from a partitioned root.
+pub fn rewrite_partition_root(
+    expression: BoundExpression,
+    root_dtype: DType,
+    collapsed: &[(FieldName, FieldName)],
+) -> VortexResult<BoundExpression> {
+    Ok(expression
+        .transform_down(|node| {
+            if let Some(value_name) = node.as_opt::<GetItem>() {
+                let partition_access = &node.children()[0];
+                if let Some(partition_name) = partition_access.as_opt::<GetItem>()
+                    && partition_access.children()[0].is_root()
+                    && collapsed.iter().any(|(partition, value)| {
+                        partition == partition_name && value == value_name
+                    })
+                {
+                    return Ok(Transformed {
+                        value: get_item(
+                            partition_name.clone(),
+                            BoundExpression::new_root(root_dtype.clone()),
+                        ),
+                        changed: true,
+                        order: TraversalOrder::Skip,
+                    });
+                }
+            }
+
+            if node.is_root() {
+                Ok(Transformed {
+                    value: BoundExpression::new_root(root_dtype.clone()),
+                    changed: true,
+                    order: TraversalOrder::Skip,
+                })
+            } else {
+                Ok(Transformed::no(node))
+            }
+        })?
+        .into_inner())
+}
+
 #[cfg(test)]
 mod tests {
     use rstest::fixture;

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)
+}