Use bound expressions in physical plans

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

Author: joseph-isaacs

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

@@ -5,10 +5,10 @@ use std::borrow::Cow;
 use std::sync::Arc;
 
 use vortex_array::dtype::DType;
-use vortex_array::expr::Expression;
-use vortex_array::expr::is_root;
-use vortex_array::expr::root;
-use vortex_array::expr::transform::replace;
+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_error::VortexResult;
 use vortex_error::vortex_bail;
 
@@ -17,24 +17,18 @@ use crate::plan::PlanRef;
 
 /// A physical plan that applies an expression to the output of `child`.
 pub struct ExpressionPlan {
-    expression: Expression,
+    expression: BoundExpression,
     child: PlanRef,
-    dtype: DType,
 }
 
 impl ExpressionPlan {
-    /// Creates an expression plan and validates its output dtype.
-    pub fn try_new(expression: Expression, child: PlanRef) -> VortexResult<Self> {
-        let dtype = expression.return_dtype(child.dtype())?;
-        Ok(Self {
-            expression,
-            child,
-            dtype,
-        })
+    /// Creates an expression plan from an expression bound to the child's dtype.
+    pub fn new(expression: BoundExpression, child: PlanRef) -> Self {
+        Self { expression, child }
     }
 
     /// Returns the expression evaluated by this plan.
-    pub fn expression(&self) -> &Expression {
+    pub fn expression(&self) -> &BoundExpression {
         &self.expression
     }
 
@@ -51,22 +45,18 @@ impl Plan for ExpressionPlan {
 
     fn optimize(&self) -> VortexResult<PlanRef> {
         let child = self.child.optimize()?;
-        let expression = self.expression.optimize_recursive(child.dtype())?;
-        if is_root(&expression) {
+        if self.expression.is_root() {
             return Ok(child);
         }
         if let Some(inner) = child.downcast_ref::<Self>() {
-            let expression = replace(expression, &root(), inner.expression.clone());
-            return Ok(Arc::new(Self::try_new(
-                expression,
-                Arc::clone(&inner.child),
-            )?));
+            let expression = replace_root(self.expression.clone(), inner.expression.clone())?;
+            return Ok(Arc::new(Self::new(expression, Arc::clone(&inner.child))));
         }
-        Ok(Arc::new(Self::try_new(expression, child)?))
+        Ok(Arc::new(Self::new(self.expression.clone(), child)))
     }
 
     fn dtype(&self) -> &DType {
-        &self.dtype
+        self.expression.dtype()
     }
 
     fn row_count(&self) -> u64 {
@@ -92,3 +82,22 @@ impl Plan for ExpressionPlan {
         }
     }
 }
+
+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())
+}

vortex-layout/src/plan/tests.rs

@@ -308,10 +308,9 @@ fn plan_display_matches_array_tree_display_shape() -> VortexResult<()> {
         vec![flat(3, field_dtype.clone(), 0), flat(3, field_dtype, 1)],
     )
     .into_layout();
-    let plan: PlanRef = Arc::new(ExpressionPlan::try_new(
-        get_item("a", root()),
-        make_plan(layout)?,
-    )?);
+    let child = make_plan(layout)?;
+    let expression = get_item("a", root()).bind(child.dtype())?;
+    let plan: PlanRef = Arc::new(ExpressionPlan::new(expression, child));
 
     assert_eq!(plan.to_string(), "ExpressionPlan(i32, rows=3)");
     insta::assert_snapshot!(plan.tree_display(), @r"
@@ -442,12 +441,13 @@ fn list_plan_display_handles_optional_validity() -> VortexResult<()> {
 fn row_idx_plan_preserves_row_index_expressions() -> VortexResult<()> {
     let layout = flat(3, primitive(PType::I32, Nullability::NonNullable), 0);
     let plan = RowIdxPlan::new_ref(10, make_plan(layout)?);
-    let plan = ExpressionPlan::try_new(row_idx(), plan)?.optimize()?;
+    let bound_expression = row_idx().bind(plan.dtype())?;
+    let plan = ExpressionPlan::new(bound_expression.clone(), plan).optimize()?;
     let expression = plan
         .downcast_ref::<ExpressionPlan>()
         .ok_or_else(|| vortex_err!("optimized plan is not an expression plan"))?;
 
-    assert_eq!(expression.expression(), &row_idx());
+    assert_eq!(expression.expression(), &bound_expression);
     assert!(expression.child_plan().is::<RowIdxPlan>());
     assert_eq!(expression.row_count(), 3);
     Ok(())