vortex_layout/plan/plans/
eval.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::borrow::Cow;
5use std::fmt;
6
7use vortex_array::EmptyMetadata;
8use vortex_array::expr::BoundExpression;
9use vortex_error::VortexResult;
10use vortex_error::vortex_bail;
11use vortex_session::registry::CachedId;
12
13use crate::plan::Plan;
14use crate::plan::PlanChildren;
15use crate::plan::PlanId;
16use crate::plan::PlanParts;
17use crate::plan::PlanRef;
18use crate::plan::PlanVTable;
19use crate::plan::check_child_count;
20use crate::plan::optimizer::PlanReduceRule;
21
22/// Applies an expression to the output of its child.
23#[derive(Clone, Debug)]
24pub struct Eval;
25
26/// The expression evaluated by an [`Eval`].
27#[derive(Clone, Debug)]
28pub struct EvalData {
29    expression: BoundExpression,
30}
31
32/// A plan that applies an expression to its child.
33pub type EvalPlan = Plan<Eval>;
34
35impl EvalPlan {
36    /// Creates an evaluation of `expression`, which must be bound to the child's dtype.
37    pub fn try_new(expression: BoundExpression, child: PlanRef) -> VortexResult<Self> {
38        validate_expression_child(&expression, &child)?;
39
40        // SAFETY: The expression root dtype was validated against the child dtype above.
41        Ok(unsafe { Self::new_unchecked(expression, child) })
42    }
43
44    /// Creates an evaluation without validating the expression's root dtype.
45    ///
46    /// # Safety
47    ///
48    /// Every scope root in `expression` must have the same dtype as `child`.
49    pub unsafe fn new_unchecked(expression: BoundExpression, child: PlanRef) -> Self {
50        PlanParts {
51            vtable: Eval,
52            dtype: expression.dtype().clone(),
53            row_count: child.row_count(),
54            children: vec![child].into(),
55            data: EvalData { expression },
56        }
57        .into_typed()
58    }
59
60    /// Returns the expression evaluated by this plan.
61    pub fn expression(&self) -> &BoundExpression {
62        &self.data().expression
63    }
64
65    /// Returns the child plan supplying the expression root.
66    pub fn child_plan(&self) -> VortexResult<PlanRef> {
67        self.child_required(0)
68    }
69}
70
71impl PlanVTable for Eval {
72    type PlanData = EvalData;
73    type Metadata = EmptyMetadata;
74
75    fn id(&self) -> PlanId {
76        static ID: CachedId = CachedId::new("vortex.plan.eval");
77        *ID
78    }
79
80    fn fmt(plan: &Plan<Self>, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
81        write!(formatter, " expr={}", plan.expression())
82    }
83
84    fn metadata(_plan: &Plan<Self>) -> Option<Self::Metadata> {
85        // Expressions serialize through `vortex.expr` protobuf, which is not wired up here yet.
86        None
87    }
88
89    fn with_children(
90        plan: &Plan<Self>,
91        children: &PlanChildren,
92        _data: &mut Self::PlanData,
93    ) -> VortexResult<()> {
94        check_child_count("Eval", children, 1)?;
95        let child = children
96            .get(0)?
97            .ok_or_else(|| vortex_error::vortex_err!("Eval child is absent"))?;
98        validate_expression_child(plan.expression(), &child)?;
99        if child.row_count() != plan.row_count() {
100            vortex_error::vortex_bail!(
101                "Eval child has {} rows but the plan has {}",
102                child.row_count(),
103                plan.row_count()
104            );
105        }
106        Ok(())
107    }
108
109    fn child_name(_plan: &Plan<Self>, index: usize) -> Cow<'_, str> {
110        if index == 0 {
111            Cow::Borrowed("child")
112        } else {
113            Cow::Owned(format!("child[{index}]"))
114        }
115    }
116}
117
118fn validate_expression_child(expression: &BoundExpression, child: &PlanRef) -> VortexResult<()> {
119    if !expression.is_root_bound_to(child.dtype()) {
120        vortex_bail!(
121            "Eval expression is not bound to child dtype {}",
122            child.dtype()
123        );
124    }
125    Ok(())
126}
127
128/// Removes an [`Eval`] whose expression is the identity expression.
129#[derive(Debug)]
130pub(crate) struct EvalIdentityRule;
131
132impl PlanReduceRule<Eval> for EvalIdentityRule {
133    fn reduce(&self, plan: &Plan<Eval>) -> VortexResult<Option<PlanRef>> {
134        if plan.expression().is_root() {
135            Ok(Some(plan.child_plan()?))
136        } else {
137            Ok(None)
138        }
139    }
140}
vortex_layout/plan/plans/eval.rs

vortex_layout/plan/plans/
eval.rs
