Push expressions through chunked and row-index plans

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

Author: joseph-isaacs

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

@@ -23,7 +23,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_array::expr::Expression;
@@ -56,16 +58,14 @@ 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>;
 
     /// Attempts to rewrite `expression` through this plan.
     ///
-    /// Returns `None` when this plan has no applicable expression rewrite. Implementations may
-    /// request only the children needed by the rewrite and should preserve all other lazy child
-    /// slots.
+    /// Returns `None` when this plan has no applicable expression rewrite. Implementations should
+    /// preserve child slots that the rewrite does not change.
     fn optimize_expression(&self, expression: &Expression) -> VortexResult<Option<PlanRef>> {
         let _ = expression;
         Ok(None)

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

@@ -5,9 +5,13 @@ use std::borrow::Cow;
 use std::sync::Arc;
 
 use vortex_array::dtype::DType;
+use vortex_array::expr::Expression;
+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;
@@ -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,28 @@ 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 optimize_expression(&self, expression: &Expression) -> VortexResult<Option<PlanRef>> {
+        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(&self.dtype)?;
+        let chunks = self
+            .chunks
+            .try_map(|_, chunk| ExpressionPlan::try_new(expression.clone(), chunk)?.optimize())?;
+        Ok(Some(Arc::new(self.with_chunks(dtype, chunks))))
     }
 
     fn dtype(&self) -> &DType {

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

@@ -14,5 +14,7 @@ pub use dict::DictPlan;
 pub use expression::ExpressionPlan;
 pub use flat::FlatPlan;
 pub use list::ListPlan;
+pub use row_idx::RowIdxPartitionPlan;
 pub use row_idx::RowIdxPlan;
+pub use row_idx::RowIdxValuesPlan;
 pub use struct_::StructPlan;