Push expressions through plan operators

Add the four parent-reduction rules that push an `Eval` into the operator
below it, and drive them from a top-down optimizer:

  Eval x Concat   push into every chunk, unless the expression reads
                  row indices, which are relative to the whole domain
  Eval x Take     evaluate over dictionary values instead of codes, only
                  when the expression is strict, infallible, and boolean
  Eval x Pack     partition per field, push each partition into its field,
                  and prune fields the expression never reads
  Eval x RowIdx   partition between generated row indices and the data
                  child, evaluating each independently

Two operators support the row-index rule: `RowIdxValues` generates global
row indices over a domain, and `RowIdxPartition` pairs an independently
evaluated row-index branch with a data branch.

The optimizer becomes top-down, driven from `Eval` nodes. A rewrite can
leave a residual expression above the same operator kind, so the driver
tracks the child it just reduced and does not immediately re-fire the same
rule on that residual. This replaces the bottom-up rewriter, whose two
rules (identity elimination and expression fusion) are subsumed by it.

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

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_012obBhJ8oPZoBbKyeS79yMv
Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

vortex-layout/src/plan/mod.rs

@@ -38,13 +38,20 @@ pub use plans::PackData;
 pub use plans::PackPlan;
 pub use plans::RowIdx;
 pub use plans::RowIdxData;
+pub use plans::RowIdxPartition;
+pub use plans::RowIdxPartitionPlan;
 pub use plans::RowIdxPlan;
 pub use plans::RowIdxPlanMetadata;
+pub use plans::RowIdxValues;
+pub use plans::RowIdxValuesData;
+pub use plans::RowIdxValuesPlan;
+pub use plans::RowIdxValuesPlanMetadata;
 pub use plans::SegmentScan;
 pub use plans::SegmentScanData;
 pub use plans::SegmentScanPlan;
 pub use plans::Take;
 pub use plans::TakePlan;
+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,74 +1,27 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-//! Plan optimization rules.
+//! Plan optimization.
 //!
-//! Rules are written against the generic plan tree rather than as a method on each operator, so
-//! adding a rule does not require touching every operator.
+//! 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_array::expr::BoundExpression;
-use vortex_array::expr::traversal::NodeExt;
-use vortex_array::expr::traversal::NodeRewriter;
-use vortex_array::expr::traversal::Transformed;
-use vortex_array::expr::traversal::TraversalOrder;
 use vortex_error::VortexResult;
 
 use crate::plan::Eval;
-use crate::plan::EvalPlan;
 use crate::plan::PlanRef;
 
 /// Optimizes `plan`, preserving its dtype and row domain.
-///
-/// Rewrites run bottom-up in a single pass, so a rule always observes children that have already
-/// been rewritten.
 pub fn optimize(plan: PlanRef) -> VortexResult<PlanRef> {
-    Ok(plan.rewrite(&mut PlanRules)?.into_inner())
-}
-
-/// The built-in plan rewrite rules.
-struct PlanRules;
-
-impl NodeRewriter for PlanRules {
-    type NodeTy = PlanRef;
-
-    fn visit_up(&mut self, node: PlanRef) -> VortexResult<Transformed<PlanRef>> {
-        let Some(eval) = node.as_opt::<Eval>() else {
-            return Ok(Transformed::no(node));
-        };
-
-        // An identity expression contributes nothing over its child.
-        if eval.expression().is_root() {
-            return Ok(Transformed::yes(eval.child_plan().clone()));
-        }
-
-        // Fuse adjacent expressions so a chain evaluates in one pass.
-        if let Some(inner) = eval.child_plan().as_opt::<Eval>() {
-            let fused = replace_root(eval.expression().clone(), inner.expression().clone())?;
-            return Ok(Transformed::yes(
-                EvalPlan::new(fused, inner.child_plan().clone()).into_plan(),
-            ));
-        }
-
-        Ok(Transformed::no(node))
+    if let Some(eval) = plan.as_opt::<Eval>() {
+        return eval.optimize_top_down(None);
     }
-}
 
-/// Substitutes `replacement` for every root reference in `expression`.
-fn replace_root(
-    expression: BoundExpression,
-    replacement: BoundExpression,
-) -> VortexResult<BoundExpression> {
-    Ok(expression
-        .transform_down(|node| {
-            if node.is_root() {
-                Ok(Transformed {
-                    value: replacement.clone(),
-                    order: TraversalOrder::Skip,
-                    changed: true,
-                })
-            } else {
-                Ok(Transformed::no(node))
-            }
-        })?
-        .into_inner())
+    let children = plan
+        .children()
+        .iter()
+        .map(|child| optimize(child.clone()))
+        .collect::<VortexResult<Vec<_>>>()?;
+    plan.with_children(children)
 }

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

@@ -9,3 +9,38 @@ 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::RowIdx;
+use super::Take;
+use super::plans::ExpressionConcatRule;
+use super::plans::ExpressionPackRule;
+use super::plans::ExpressionRowIdxRule;
+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_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 `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,15 +6,21 @@ 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::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)]
@@ -93,3 +99,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| EvalPlan::new(expression.clone(), chunk.clone()).into_plan())
+            .collect::<Vec<_>>();
+        Ok(Some(
+            ConcatPlan::try_new(expression.dtype().clone(), chunks)?.into_plan(),
+        ))
+    }
+}

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

@@ -4,7 +4,14 @@
 use std::borrow::Cow;
 
 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;
 
@@ -14,6 +21,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)]
@@ -82,3 +91,122 @@ 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().clone());
+        }
+        if let Some(inner) = self.child_plan().as_opt::<Eval>() {
+            let expression = replace_root(self.expression().clone(), inner.expression().clone())?;
+            return EvalPlan::new(expression, inner.child_plan().clone()).optimize_top_down(None);
+        }
+
+        let child_type = self.child_plan().id();
+        let parent =
+            EvalPlan::new(self.expression().clone(), self.child_plan().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(self.child_plan().clone())?;
+        if let Some(inner) = child.as_opt::<Eval>() {
+            let expression = replace_root(self.expression().clone(), inner.expression().clone())?;
+            return EvalPlan::new(expression, inner.child_plan().clone()).optimize_top_down(None);
+        }
+
+        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)
+    }
+}
+
+/// Substitutes `replacement` for every root reference in `expression`.
+pub(crate) fn replace_root(
+    expression: BoundExpression,
+    replacement: BoundExpression,
+) -> VortexResult<BoundExpression> {
+    Ok(expression
+        .transform_down(|node| {
+            if node.is_root() {
+                Ok(Transformed {
+                    value: replacement.clone(),
+                    order: TraversalOrder::Skip,
+                    changed: true,
+                })
+            } else {
+                Ok(Transformed::no(node))
+            }
+        })?
+        .into_inner())
+}
+
+/// 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())
+}