Push expressions through plan operators

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

Author: joseph-isaacs

vortex-layout/src/plan/optimize.rs

@@ -1,7 +1,11 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-//! Generic bottom-up optimization over physical plans.
+//! Plan optimization.
+//!
+//! Optimization is driven top-down from [`Eval`] nodes, which apply the static parent-reduction
+//! rules in [`crate::plan::optimizer`] as they become applicable. Operators without a rule simply
+//! optimize their children.
 
 use vortex_error::VortexResult;
 
@@ -10,26 +14,14 @@ use crate::plan::PlanRef;
 
 /// Optimizes `plan`, preserving its dtype and row domain.
 pub fn optimize(plan: PlanRef) -> VortexResult<PlanRef> {
-    let mut children = Vec::with_capacity(plan.child_count());
-    let mut changed = false;
-    for child in plan.children().iter() {
-        let child = child?;
-        let optimized = optimize(child.clone())?;
-        changed |= !PlanRef::ptr_eq(&child, &optimized);
-        children.push(optimized);
+    if let Some(eval) = plan.as_opt::<Eval>() {
+        return eval.optimize_top_down(None);
     }
 
-    let plan = if changed {
-        plan.with_children(children)?
-    } else {
-        plan
-    };
-
-    let Some(eval) = plan.as_opt::<Eval>() else {
-        return Ok(plan);
-    };
-    if eval.expression().is_root() {
-        return eval.child_plan();
-    }
-    Ok(plan)
+    let children = plan
+        .children()
+        .iter()
+        .map(|child| optimize(child?))
+        .collect::<VortexResult<Vec<_>>>()?;
+    plan.with_children(children)
 }

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

@@ -9,3 +9,33 @@ pub use rules::DynPlanParentReduceRule;
 pub use rules::PlanParentReduceRule;
 pub use rules::PlanParentReduceRuleAdapter;
 pub use rules::PlanParentRuleSet;
+use vortex_error::VortexResult;
+
+use super::Concat;
+use super::Pack;
+use super::PlanRef;
+use super::Take;
+use super::plans::ExpressionConcatRule;
+use super::plans::ExpressionPackRule;
+use super::plans::ExpressionTakeRule;
+
+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 `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/plans/concat.rs

@@ -6,16 +6,22 @@ use std::sync::Arc;
 
 use vortex_array::EmptyMetadata;
 use vortex_array::dtype::DType;
+use vortex_array::expr::ExactBoundExpr;
+use vortex_array::expr::label_bound_tree;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_session::registry::CachedId;
 
+use crate::layouts::row_idx::RowIdx as RowIdxFn;
+use crate::plan::Eval;
+use crate::plan::EvalPlan;
 use crate::plan::Plan;
 use crate::plan::PlanChildren;
 use crate::plan::PlanId;
 use crate::plan::PlanParts;
 use crate::plan::PlanRef;
 use crate::plan::PlanVTable;
+use crate::plan::optimizer::PlanParentReduceRule;
 
 /// Concatenates its children row-wise.
 #[derive(Clone, Debug)]
@@ -134,3 +140,45 @@ impl PlanVTable for Concat {
         Cow::Owned(format!("chunks[{index}]"))
     }
 }
+
+/// Pushes an expression into every chunk of a [`Concat`].
+#[derive(Debug)]
+pub(crate) struct ExpressionConcatRule;
+
+impl PlanParentReduceRule<Concat> for ExpressionConcatRule {
+    type Parent = Eval;
+
+    fn reduce_parent(
+        &self,
+        child: &Plan<Concat>,
+        parent: &Plan<Eval>,
+        _child_idx: usize,
+    ) -> VortexResult<Option<PlanRef>> {
+        let expression = parent.expression();
+        // Row-index expressions are relative to the whole row domain, so they cannot be evaluated
+        // chunk by chunk.
+        let references_row_idx = label_bound_tree(
+            expression,
+            |node| {
+                node.as_scalar()
+                    .is_some_and(|scalar_fn| scalar_fn.is::<RowIdxFn>())
+            },
+            |acc, &child| acc | child,
+        )
+        .get(&ExactBoundExpr(expression.clone()))
+        .copied()
+        .unwrap_or(false);
+        if references_row_idx {
+            return Ok(None);
+        }
+
+        let chunks = child
+            .children()
+            .iter()
+            .map(|chunk| Ok(EvalPlan::new(expression.clone(), chunk?).into_plan()))
+            .collect::<VortexResult<Vec<_>>>()?;
+        Ok(Some(
+            ConcatPlan::try_new(expression.dtype().clone(), chunks)?.into_plan(),
+        ))
+    }
+}

vortex-layout/src/plan/plans/eval.rs

@@ -5,7 +5,14 @@ use std::borrow::Cow;
 use std::fmt;
 
 use vortex_array::EmptyMetadata;
+use vortex_array::dtype::DType;
+use vortex_array::dtype::FieldName;
 use vortex_array::expr::BoundExpression;
+use vortex_array::expr::traversal::NodeExt;
+use vortex_array::expr::traversal::Transformed;
+use vortex_array::expr::traversal::TraversalOrder;
+use vortex_array::scalar_fn::ScalarFnVTableExt;
+use vortex_array::scalar_fn::fns::get_item::GetItem;
 use vortex_error::VortexResult;
 use vortex_session::registry::CachedId;
 
@@ -16,6 +23,8 @@ use crate::plan::PlanParts;
 use crate::plan::PlanRef;
 use crate::plan::PlanVTable;
 use crate::plan::check_child_count;
+use crate::plan::optimize;
+use crate::plan::optimizer::reduce_parent;
 
 /// Applies an expression to the output of its child.
 #[derive(Clone, Debug)]
@@ -99,3 +108,94 @@ impl PlanVTable for Eval {
         }
     }
 }
+
+impl EvalPlan {
+    /// Optimizes this plan top-down, applying parent-reduction rules as they become applicable.
+    ///
+    /// `blocked_child_type` suppresses one rule re-firing on its own residual output, which would
+    /// otherwise loop when a rewrite leaves an expression above the same child kind.
+    pub(crate) fn optimize_top_down(
+        &self,
+        blocked_child_type: Option<PlanId>,
+    ) -> VortexResult<PlanRef> {
+        if self.expression().is_root() {
+            return optimize(self.child_plan()?);
+        }
+
+        let child = self.child_plan()?;
+        let child_type = child.id();
+        let parent = EvalPlan::new(self.expression().clone(), child.clone()).into_plan();
+        if blocked_child_type != Some(child_type)
+            && let Some(rewritten) = reduce_parent(&parent, 0)?
+        {
+            return Self::optimize_rewrite(rewritten, child_type);
+        }
+
+        let child = optimize(child)?;
+
+        let child_type = child.id();
+        let parent = EvalPlan::new(self.expression().clone(), child).into_plan();
+        if blocked_child_type != Some(child_type)
+            && let Some(rewritten) = reduce_parent(&parent, 0)?
+        {
+            return Self::optimize_rewrite(rewritten, child_type);
+        }
+        Ok(parent)
+    }
+
+    fn optimize_rewrite(rewritten: PlanRef, previous_child_type: PlanId) -> VortexResult<PlanRef> {
+        let Some(eval) = rewritten.as_opt::<Eval>() else {
+            return optimize(rewritten);
+        };
+        // A residual expression may remain above the same child kind after a successful rewrite.
+        // Do not immediately apply that rule again; recursively optimize only the retained child.
+        let child_type = eval.child_plan()?.id();
+        let blocked = (child_type == previous_child_type).then_some(previous_child_type);
+        eval.optimize_top_down(blocked)
+    }
+}
+
+/// Rewrites partition accessors in `expression` to read from a partitioned root.
+pub(crate) 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_scalar()
+                .and_then(|scalar_fn| scalar_fn.as_opt::<GetItem>())
+            {
+                let partition_access = &node.children()[0];
+                if let Some(partition_name) = partition_access
+                    .as_scalar()
+                    .and_then(|scalar_fn| scalar_fn.as_opt::<GetItem>())
+                    && partition_access.children()[0].is_root()
+                    && collapsed.iter().any(|(partition, value)| {
+                        partition == partition_name && value == value_name
+                    })
+                {
+                    return Ok(Transformed {
+                        value: BoundExpression::try_new(
+                            GetItem.bind(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())
+}

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

@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 mod concat;
-mod eval;
+pub(crate) mod eval;
 mod list_pack;
 mod pack;
 mod row_idx;
@@ -12,12 +12,14 @@ mod take;
 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 use eval::EvalPlan;
 pub use list_pack::ListPack;
 pub use list_pack::ListPackData;
 pub use list_pack::ListPackPlan;
+pub(crate) use pack::ExpressionPackRule;
 pub use pack::Pack;
 pub use pack::PackData;
 pub use pack::PackPlan;
@@ -28,5 +30,6 @@ pub use row_idx::RowIdxPlanMetadata;
 pub use segment_scan::SegmentScan;
 pub use segment_scan::SegmentScanData;
 pub use segment_scan::SegmentScanPlan;
+pub(crate) use take::ExpressionTakeRule;
 pub use take::Take;
 pub use take::TakePlan;