Push expressions through layout plans

vortex-layout/src/plan/children.rs

@@ -46,17 +46,24 @@ impl LazyPlanChildren {
         Ok(cell.get_or_try_init(|| (self.initializer)(index))?.clone())
     }
 
-    /// Lazily transforms each present child into a new child collection.
-    pub(crate) fn map(
+    /// Eagerly transforms each present child into a new child collection.
+    pub(crate) fn try_map(
         &self,
-        transform: impl Fn(usize, PlanRef) -> VortexResult<PlanRef> + 'static + Send + Sync,
-    ) -> Self {
-        let source = self.clone();
-        Self::new(self.len(), move |index| {
-            source
-                .get(index)?
-                .map(|child| transform(index, child))
-                .transpose()
-        })
+        transform: impl Fn(usize, PlanRef) -> VortexResult<PlanRef>,
+    ) -> VortexResult<Self> {
+        // TODO: Make recursive child optimization lazy again once the optimizer API can
+        // explicitly distinguish fully optimized plans from plans with deferred optimizer work.
+        let children = (0..self.len())
+            .map(|index| {
+                self.get(index)?
+                    .map(|child| transform(index, child))
+                    .transpose()
+            })
+            .collect::<VortexResult<Vec<_>>>()?;
+        let children: Arc<[Option<PlanRef>]> = children.into();
+        let len = children.len();
+        Ok(Self::new(len, move |index| {
+            Ok(children.get(index).cloned().flatten())
+        }))
     }
 }

vortex-layout/src/plan/mod.rs

@@ -24,7 +24,9 @@ pub use plans::DictPlan;
 pub use plans::ExpressionPlan;
 pub use plans::FlatPlan;
 pub use plans::ListPlan;
+pub use plans::RowIdxPartitionPlan;
 pub use plans::RowIdxPlan;
+pub use plans::RowIdxValuesPlan;
 pub use plans::StructPlan;
 use vortex_array::dtype::DType;
 use vortex_error::VortexResult;
@@ -58,9 +60,8 @@ pub trait Plan: 'static + Send + Sync {
         std::any::type_name::<Self>()
     }
 
-    /// Optimizes this plan while preserving its dtype and row domain.
-    ///
-    /// Implementations may defer child optimization until the child is accessed.
+    /// Recursively optimizes this plan and all of its children while preserving its dtype and row
+    /// domain.
     fn optimize(&self) -> VortexResult<PlanRef>;
 
     /// Returns the dtype produced by this plan.

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::ChunkedPlan;
+use super::DictPlan;
+use super::PlanRef;
+use super::RowIdxPlan;
+use super::StructPlan;
+use super::plans::ExpressionChunkedRule;
+use super::plans::ExpressionDictRule;
+use super::plans::ExpressionRowIdxRule;
+use super::plans::ExpressionStructRule;
+
+static EXPRESSION_CHUNKED_RULE: PlanParentReduceRuleAdapter<ChunkedPlan, ExpressionChunkedRule> =
+    PlanParentReduceRuleAdapter::new(ExpressionChunkedRule);
+static EXPRESSION_DICT_RULE: PlanParentReduceRuleAdapter<DictPlan, ExpressionDictRule> =
+    PlanParentReduceRuleAdapter::new(ExpressionDictRule);
+static EXPRESSION_ROW_IDX_RULE: PlanParentReduceRuleAdapter<RowIdxPlan, ExpressionRowIdxRule> =
+    PlanParentReduceRuleAdapter::new(ExpressionRowIdxRule);
+static EXPRESSION_STRUCT_RULE: PlanParentReduceRuleAdapter<StructPlan, ExpressionStructRule> =
+    PlanParentReduceRuleAdapter::new(ExpressionStructRule);
+
+static PARENT_RULES: PlanParentRuleSet = PlanParentRuleSet::new(&[
+    &EXPRESSION_CHUNKED_RULE,
+    &EXPRESSION_DICT_RULE,
+    &EXPRESSION_ROW_IDX_RULE,
+    &EXPRESSION_STRUCT_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/chunked.rs

@@ -5,13 +5,17 @@ use std::borrow::Cow;
 use std::sync::Arc;
 
 use vortex_array::dtype::DType;
+use vortex_array::expr::label_tree;
 use vortex_error::VortexResult;
 
 use crate::layouts::chunked::ChunkedLayout;
+use crate::layouts::row_idx::RowIdx;
+use crate::plan::ExpressionPlan;
 use crate::plan::LazyPlanChildren;
 use crate::plan::Plan;
 use crate::plan::PlanRef;
 use crate::plan::new_plan;
+use crate::plan::optimizer::PlanParentReduceRule;
 
 /// A physical plan with one child per row chunk.
 pub struct ChunkedPlan {
@@ -36,10 +40,10 @@ impl ChunkedPlan {
         }
     }
 
-    fn with_chunks(&self, chunks: LazyPlanChildren) -> Self {
+    fn with_chunks(&self, dtype: DType, chunks: LazyPlanChildren) -> Self {
         Self {
             layout: self.layout.clone(),
-            dtype: self.dtype.clone(),
+            dtype,
             chunks,
         }
     }
@@ -55,8 +59,8 @@ impl Plan for ChunkedPlan {
     }
 
     fn optimize(&self) -> VortexResult<PlanRef> {
-        let chunks = self.chunks.map(|_, chunk| chunk.optimize());
-        Ok(Arc::new(self.with_chunks(chunks)))
+        let chunks = self.chunks.try_map(|_, chunk| chunk.optimize())?;
+        Ok(Arc::new(self.with_chunks(self.dtype.clone(), chunks)))
     }
 
     fn dtype(&self) -> &DType {
@@ -82,3 +86,37 @@ impl Plan for ChunkedPlan {
         Cow::Owned(format!("chunks[{index}]"))
     }
 }
+
+/// Pushes an expression through every chunk of a chunked plan.
+#[derive(Debug)]
+pub(crate) struct ExpressionChunkedRule;
+
+impl PlanParentReduceRule<ChunkedPlan> for ExpressionChunkedRule {
+    type Parent = ExpressionPlan;
+
+    fn reduce_parent(
+        &self,
+        child: &ChunkedPlan,
+        parent: &ExpressionPlan,
+        _child_idx: usize,
+    ) -> VortexResult<Option<PlanRef>> {
+        let expression = parent.expression();
+        let references_row_idx = label_tree(
+            expression,
+            |node| node.is::<RowIdx>(),
+            |acc, &child| acc | child,
+        )
+        .get(expression)
+        .copied()
+        .unwrap_or(false);
+        if references_row_idx {
+            return Ok(None);
+        }
+
+        let dtype = expression.return_dtype(&child.dtype)?;
+        let chunks = child
+            .chunks
+            .try_map(|_, chunk| ExpressionPlan::try_new(expression.clone(), chunk)?.optimize())?;
+        Ok(Some(Arc::new(child.with_chunks(dtype, chunks))))
+    }
+}

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

@@ -4,13 +4,19 @@
 use std::borrow::Cow;
 use std::sync::Arc;
 
+use vortex_array::expr::is_root;
+use vortex_array::expr::label_is_fallible;
+use vortex_array::expr::label_strict;
+use vortex_array::expr::label_tree;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 
 use crate::layouts::dict::DictLayout;
+use crate::plan::ExpressionPlan;
 use crate::plan::Plan;
 use crate::plan::PlanRef;
 use crate::plan::new_plan;
+use crate::plan::optimizer::PlanParentReduceRule;
 
 /// A physical dictionary plan with children ordered as `[codes, values]`.
 pub struct DictPlan {
@@ -95,3 +101,44 @@ impl Plan for DictPlan {
         }
     }
 }
+
+/// Pushes a safe boolean expression into dictionary values.
+#[derive(Debug)]
+pub(crate) struct ExpressionDictRule;
+
+impl PlanParentReduceRule<DictPlan> for ExpressionDictRule {
+    type Parent = ExpressionPlan;
+
+    fn reduce_parent(
+        &self,
+        child: &DictPlan,
+        parent: &ExpressionPlan,
+        _child_idx: usize,
+    ) -> VortexResult<Option<PlanRef>> {
+        let expression = parent.expression();
+        if !expression.return_dtype(&child.dtype)?.is_boolean() {
+            return Ok(None);
+        }
+        let references_root = label_tree(expression, is_root, |acc, &child| acc | child)
+            .get(expression)
+            .copied()
+            .unwrap_or(false);
+        let is_strict = label_strict(expression)
+            .get(expression)
+            .copied()
+            .unwrap_or(false);
+        let is_fallible = label_is_fallible(expression)
+            .get(expression)
+            .copied()
+            .unwrap_or(true);
+        if !references_root || !is_strict || is_fallible {
+            return Ok(None);
+        }
+
+        let values =
+            ExpressionPlan::try_new(expression.clone(), Arc::clone(&child.values))?.optimize()?;
+        Ok(Some(Arc::new(
+            child.with_children(Arc::clone(&child.codes), values),
+        )))
+    }
+}

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

@@ -15,6 +15,7 @@ use vortex_error::vortex_bail;
 
 use crate::plan::Plan;
 use crate::plan::PlanRef;
+use crate::plan::optimizer::reduce_parent;
 
 /// A physical plan that applies an expression to the output of `child`.
 pub struct ExpressionPlan {
@@ -62,12 +63,13 @@ impl Plan for ExpressionPlan {
         }
         if let Some(inner) = child.as_any().downcast_ref::<Self>() {
             let expression = replace(expression, &root(), inner.expression.clone());
-            return Ok(Arc::new(Self::try_new(
-                expression,
-                Arc::clone(&inner.child),
-            )?));
+            return Self::try_new(expression, Arc::clone(&inner.child))?.optimize();
         }
-        Ok(Arc::new(Self::try_new(expression, child)?))
+        let parent: PlanRef = Arc::new(Self::try_new(expression, child)?);
+        if let Some(rewritten) = reduce_parent(&parent, 0)? {
+            return Ok(rewritten);
+        }
+        Ok(parent)
     }
 
     fn dtype(&self) -> &DType {

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

@@ -10,9 +10,15 @@ mod row_idx;
 mod struct_;
 
 pub use chunked::ChunkedPlan;
+pub(crate) use chunked::ExpressionChunkedRule;
 pub use dict::DictPlan;
+pub(crate) use dict::ExpressionDictRule;
 pub use expression::ExpressionPlan;
 pub use flat::FlatPlan;
 pub use list::ListPlan;
+pub(crate) use row_idx::ExpressionRowIdxRule;
+pub use row_idx::RowIdxPartitionPlan;
 pub use row_idx::RowIdxPlan;
+pub use row_idx::RowIdxValuesPlan;
+pub(crate) use struct_::ExpressionStructRule;
 pub use struct_::StructPlan;