Optimize expression plans top down

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

Author: joseph-isaacs

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

@@ -116,7 +116,7 @@ impl PlanParentReduceRule<ChunkedPlan> for ExpressionChunkedRule {
         let dtype = expression.dtype().clone();
         let chunks = child
             .chunks
-            .try_map(|_, chunk| ExpressionPlan::new(expression.clone(), chunk).optimize())?;
+            .try_map(|_, chunk| Ok(ExpressionPlan::new_ref(expression.clone(), chunk)))?;
         Ok(Some(Arc::new(child.with_chunks(dtype, chunks))))
     }
 }

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

@@ -133,8 +133,7 @@ impl PlanParentReduceRule<DictPlan> for ExpressionDictRule {
             return Ok(None);
         }
 
-        let values =
-            ExpressionPlan::new(expression.clone(), Arc::clone(&child.values)).optimize()?;
+        let values = ExpressionPlan::new_ref(expression.clone(), Arc::clone(&child.values));
         Ok(Some(Arc::new(
             child.with_children(Arc::clone(&child.codes), values),
         )))

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

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::any::TypeId;
 use std::borrow::Cow;
 use std::sync::Arc;
 
@@ -40,29 +41,65 @@ impl ExpressionPlan {
     pub fn child_plan(&self) -> &PlanRef {
         &self.child
     }
-}
 
-impl Plan for ExpressionPlan {
-    fn name(&self) -> &'static str {
-        "ExpressionPlan"
+    pub(crate) fn new_ref(expression: BoundExpression, child: PlanRef) -> PlanRef {
+        Arc::new(Self::new(expression, child))
     }
 
-    fn optimize(&self) -> VortexResult<PlanRef> {
-        let child = self.child.optimize()?;
+    fn optimize_top_down(&self, blocked_child_type: Option<TypeId>) -> VortexResult<PlanRef> {
         if self.expression.is_root() {
-            return Ok(child);
+            return self.child.optimize();
         }
+        if let Some(inner) = self.child.downcast_ref::<Self>() {
+            let expression = replace_root(self.expression.clone(), inner.expression.clone())?;
+            return Self::new(expression, Arc::clone(&inner.child)).optimize_top_down(None);
+        }
+
+        let child_type = self.child.as_ref().type_id();
+        let parent = Self::new_ref(self.expression.clone(), Arc::clone(&self.child));
+        if blocked_child_type != Some(child_type)
+            && let Some(rewritten) = reduce_parent(&parent, 0)?
+        {
+            return Self::optimize_rewrite(rewritten, child_type);
+        }
+
+        let child = self.child.optimize()?;
         if let Some(inner) = child.downcast_ref::<Self>() {
             let expression = replace_root(self.expression.clone(), inner.expression.clone())?;
-            return Self::new(expression, Arc::clone(&inner.child)).optimize();
+            return Self::new(expression, Arc::clone(&inner.child)).optimize_top_down(None);
         }
-        let parent: PlanRef = Arc::new(Self::new(self.expression.clone(), child));
-        if let Some(rewritten) = reduce_parent(&parent, 0)? {
-            return Ok(rewritten);
+
+        let child_type = child.as_ref().type_id();
+        let parent = Self::new_ref(self.expression.clone(), child);
+        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: TypeId) -> VortexResult<PlanRef> {
+        let Some(expression) = rewritten.downcast_ref::<Self>() else {
+            return rewritten.optimize();
+        };
+        let child_type = expression.child.as_ref().type_id();
+        // 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 blocked_child_type = (child_type == previous_child_type).then_some(previous_child_type);
+        expression.optimize_top_down(blocked_child_type)
+    }
+}
+
+impl Plan for ExpressionPlan {
+    fn name(&self) -> &'static str {
+        "ExpressionPlan"
+    }
+
+    fn optimize(&self) -> VortexResult<PlanRef> {
+        self.optimize_top_down(None)
+    }
+
     fn dtype(&self) -> &DType {
         self.expression.dtype()
     }

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

@@ -109,11 +109,12 @@ impl PlanParentReduceRule<RowIdxPlan> for ExpressionRowIdxRule {
                 RowIdxExpressionPartition::RowIdx => {
                     let expression = replace_row_idx(expression.clone())?;
                     let values = RowIdxValuesPlan::new_ref(child.row_offset, child.row_count());
-                    Ok(Some(ExpressionPlan::new(expression, values).optimize()?))
+                    Ok(Some(ExpressionPlan::new_ref(expression, values)))
                 }
-                RowIdxExpressionPartition::Child => Ok(Some(
-                    ExpressionPlan::new(expression.clone(), Arc::clone(&child.child)).optimize()?,
-                )),
+                RowIdxExpressionPartition::Child => Ok(Some(ExpressionPlan::new_ref(
+                    expression.clone(),
+                    Arc::clone(&child.child),
+                ))),
             };
         }
 
@@ -171,18 +172,16 @@ impl PlanParentReduceRule<RowIdxPlan> for ExpressionRowIdxRule {
         };
 
         let row_idx_expression = replace_row_idx(row_idx_expression)?;
-        let row_idx_plan = ExpressionPlan::new(
+        let row_idx_plan = ExpressionPlan::new_ref(
             row_idx_expression,
             RowIdxValuesPlan::new_ref(child.row_offset, child.row_count()),
-        )
-        .optimize()?;
-        let child_plan =
-            ExpressionPlan::new(child_expression, Arc::clone(&child.child)).optimize()?;
+        );
+        let child_plan = ExpressionPlan::new_ref(child_expression, Arc::clone(&child.child));
         let partitions = RowIdxPartitionPlan::try_new(row_idx_plan, child_plan)?;
         let residual =
             rewrite_partition_root(partitioned.root, partitions.dtype().clone(), &collapsed)?;
 
-        Ok(Some(ExpressionPlan::new(residual, partitions).optimize()?))
+        Ok(Some(ExpressionPlan::new_ref(residual, partitions)))
     }
 }
 

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

@@ -235,7 +235,7 @@ impl PlanParentReduceRule<StructPlan> for ExpressionStructRule {
                 .ok_or_else(|| vortex_err!("Struct field '{field_name}' has no plan"))?;
             let lowered = step_into_struct_field(expanded, field_name, field.dtype().clone())?;
 
-            return Ok(Some(ExpressionPlan::new(lowered, field).optimize()?));
+            return Ok(Some(ExpressionPlan::new_ref(lowered, field)));
         }
 
         let residual = partitioned.root;
@@ -283,13 +283,13 @@ impl PlanParentReduceRule<StructPlan> for ExpressionStructRule {
                 .children
                 .get(field_index)?
                 .ok_or_else(|| vortex_err!("Struct field '{field_name}' has no plan"))?;
-            let field = ExpressionPlan::new(expression, field).optimize()?;
+            let field = ExpressionPlan::new_ref(expression, field);
             pruned_fields.push((field_name, field));
         }
         let rewritten: PlanRef = Arc::new(child.with_pruned_fields(pruned_fields)?);
         let residual = rewrite_partition_root(residual, rewritten.dtype().clone(), &collapsed)?;
 
-        Ok(Some(Arc::new(ExpressionPlan::new(residual, rewritten))))
+        Ok(Some(ExpressionPlan::new_ref(residual, rewritten)))
     }
 }
 
@@ -401,3 +401,112 @@ fn bound_pack(names: FieldNames, children: Vec<BoundExpression>) -> VortexResult
         children,
     )
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+    use std::sync::atomic::AtomicUsize;
+    use std::sync::atomic::Ordering;
+
+    use vortex_array::dtype::DType;
+    use vortex_array::dtype::Nullability;
+    use vortex_array::dtype::PType;
+    use vortex_array::dtype::StructFields;
+    use vortex_array::expr::get_item;
+    use vortex_array::expr::root;
+    use vortex_error::VortexResult;
+    use vortex_session::registry::ReadContext;
+
+    use super::StructPlan;
+    use crate::LayoutRef;
+    use crate::layouts::flat::FlatLayout;
+    use crate::layouts::struct_::StructLayout;
+    use crate::plan::ExpressionPlan;
+    use crate::plan::LazyPlanChildren;
+    use crate::plan::Plan;
+    use crate::plan::PlanRef;
+    use crate::plan::RowIdxPlan;
+    use crate::segments::SegmentId;
+
+    struct CountingPlan {
+        dtype: DType,
+        optimizations: Arc<AtomicUsize>,
+    }
+
+    impl CountingPlan {
+        fn new_ref(dtype: DType, optimizations: Arc<AtomicUsize>) -> PlanRef {
+            Arc::new(Self {
+                dtype,
+                optimizations,
+            })
+        }
+    }
+
+    impl Plan for CountingPlan {
+        fn optimize(&self) -> VortexResult<PlanRef> {
+            self.optimizations.fetch_add(1, Ordering::Relaxed);
+            Ok(Self::new_ref(
+                self.dtype.clone(),
+                Arc::clone(&self.optimizations),
+            ))
+        }
+
+        fn dtype(&self) -> &DType {
+            &self.dtype
+        }
+
+        fn row_count(&self) -> u64 {
+            1
+        }
+    }
+
+    fn flat(dtype: DType, segment_id: u32) -> LayoutRef {
+        FlatLayout::new(1, dtype, SegmentId::from(segment_id), ReadContext::new([])).into_layout()
+    }
+
+    #[test]
+    fn expression_optimizes_only_referenced_struct_fields() -> VortexResult<()> {
+        let field_dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
+        let struct_dtype = DType::Struct(
+            StructFields::from_iter([("a", field_dtype.clone()), ("b", field_dtype.clone())]),
+            Nullability::NonNullable,
+        );
+        let layout = StructLayout::new(
+            1,
+            struct_dtype.clone(),
+            vec![flat(field_dtype.clone(), 0), flat(field_dtype.clone(), 1)],
+        );
+        let a_optimizations = Arc::new(AtomicUsize::new(0));
+        let b_optimizations = Arc::new(AtomicUsize::new(0));
+        let children: Arc<[Option<PlanRef>]> = [
+            Some(CountingPlan::new_ref(
+                field_dtype.clone(),
+                Arc::clone(&a_optimizations),
+            )),
+            Some(CountingPlan::new_ref(
+                field_dtype,
+                Arc::clone(&b_optimizations),
+            )),
+            None,
+        ]
+        .into();
+        let child_count = children.len();
+        let struct_plan: PlanRef = Arc::new(StructPlan {
+            layout,
+            dtype: struct_dtype,
+            children: LazyPlanChildren::new(child_count, move |index| {
+                Ok(children.get(index).cloned().flatten())
+            }),
+        });
+        let plan = RowIdxPlan::new_ref(0, struct_plan);
+
+        let expression = get_item("a", root())
+            .optimize_recursive(plan.dtype())?
+            .bind(plan.dtype())?;
+        ExpressionPlan::new(expression, plan).optimize()?;
+
+        assert_eq!(a_optimizations.load(Ordering::Relaxed), 1);
+        assert_eq!(b_optimizations.load(Ordering::Relaxed), 0);
+        Ok(())
+    }
+}