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

vortex_layout/plan/physical/plans/
eval.rs
