vortex_array/expr/
bound_expression.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt;
5use std::fmt::Display;
6use std::fmt::Formatter;
7use std::hash::Hash;
8use std::hash::Hasher;
9use std::sync::Arc;
10
11use itertools::Itertools;
12use vortex_error::VortexExpect;
13use vortex_error::VortexResult;
14use vortex_error::vortex_ensure;
15use vortex_session::VortexSession;
16
17use crate::dtype::DType;
18use crate::expr::Expression;
19use crate::expr::display::DisplayTreeExpr;
20use crate::expr::scope::Scope;
21use crate::expr::traversal::TraversalOrder;
22use crate::expr::traversal::pre_order_visit_down;
23use crate::scalar_fn::ScalarFnRef;
24use crate::scalar_fn::ScalarFnVTable;
25use crate::stats::rewrite::StatsRewriteCtx;
26
27/// An [`Expression`] that has been type-checked against a [`Scope`].
28///
29/// Every node carries its own dtype, so reading one is a field access rather than a walk of the
30/// subtree. Holding a `BoundExpression` is proof that the whole tree type-checked.
31///
32/// Binding is purely logical: it deals only in [`DType`]s and never sees an array, a length, or an
33/// encoding.
34#[derive(Clone, Debug, PartialEq, Eq, Hash)]
35pub enum BoundExpression {
36    /// A scalar function applied to bound children.
37    Scalar {
38        /// The dtype this node evaluates to.
39        dtype: DType,
40        /// The scalar function for this node.
41        scalar_fn: ScalarFnRef,
42        /// The bound children, in argument order.
43        ///
44        /// Sharing keeps clones cheap even though the iterative [`Drop`] implementation prevents
45        /// consumers from destructuring a `BoundExpression` by value.
46        children: Arc<Vec<BoundExpression>>,
47    },
48    /// The scope itself. Its dtype is the scope's root dtype.
49    Root {
50        /// The dtype this node evaluates to.
51        dtype: DType,
52    },
53}
54
55/// A bound-expression wrapper that compares shared tree identity instead of structure.
56#[derive(Clone, Debug)]
57pub struct ExactBoundExpr(pub BoundExpression);
58
59impl PartialEq for ExactBoundExpr {
60    fn eq(&self, other: &Self) -> bool {
61        match (&self.0, &other.0) {
62            (
63                BoundExpression::Root { dtype: lhs_dtype },
64                BoundExpression::Root { dtype: rhs_dtype },
65            ) => lhs_dtype == rhs_dtype,
66            (
67                BoundExpression::Scalar {
68                    dtype: lhs_dtype,
69                    scalar_fn: lhs_fn,
70                    children: lhs_children,
71                },
72                BoundExpression::Scalar {
73                    dtype: rhs_dtype,
74                    scalar_fn: rhs_fn,
75                    children: rhs_children,
76                },
77            ) => {
78                lhs_fn == rhs_fn
79                    && Arc::ptr_eq(lhs_children, rhs_children)
80                    && lhs_dtype == rhs_dtype
81            }
82            _ => false,
83        }
84    }
85}
86
87impl Eq for ExactBoundExpr {}
88
89impl Hash for ExactBoundExpr {
90    fn hash<H: Hasher>(&self, state: &mut H) {
91        // DType differences are resolved by equality. Omitting the potentially lazy dtype keeps
92        // identity-keyed cache lookups from deserializing an entire schema just to compute a hash.
93        match &self.0 {
94            BoundExpression::Root { .. } => state.write_u8(0),
95            BoundExpression::Scalar {
96                scalar_fn,
97                children,
98                ..
99            } => {
100                state.write_u8(1);
101                scalar_fn.hash(state);
102                Arc::as_ptr(children).hash(state);
103            }
104        }
105    }
106}
107
108impl BoundExpression {
109    /// Create a bound root expression with the given dtype.
110    pub fn new_root(dtype: DType) -> Self {
111        Self::Root { dtype }
112    }
113
114    /// Create a bound scalar node from a scalar function and already-bound children.
115    pub fn try_new(
116        scalar_fn: ScalarFnRef,
117        children: impl IntoIterator<Item = BoundExpression>,
118    ) -> VortexResult<Self> {
119        Self::try_new_vec(scalar_fn, children.into_iter().collect())
120    }
121
122    fn try_new_vec(scalar_fn: ScalarFnRef, children: Vec<BoundExpression>) -> VortexResult<Self> {
123        vortex_ensure!(
124            scalar_fn.signature().arity().matches(children.len()),
125            "Expression arity mismatch: expected {} children but got {}",
126            scalar_fn.signature().arity(),
127            children.len()
128        );
129
130        let arg_dtypes = children
131            .iter()
132            .map(|child| child.dtype().clone())
133            .collect_vec();
134        let dtype = scalar_fn.return_dtype(&arg_dtypes)?;
135
136        Ok(Self::Scalar {
137            dtype,
138            scalar_fn,
139            children: children.into(),
140        })
141    }
142
143    /// Rebuild this node with new bound children, recomputing its dtype.
144    pub fn with_children(
145        self,
146        children: impl IntoIterator<Item = BoundExpression>,
147    ) -> VortexResult<Self> {
148        let children = Vec::from_iter(children);
149        let BoundExpression::Scalar { scalar_fn, .. } = &self else {
150            vortex_ensure!(
151                children.is_empty(),
152                "Root expression cannot have {} children",
153                children.len()
154            );
155            return Ok(self);
156        };
157
158        Self::try_new_vec(scalar_fn.clone(), children)
159    }
160
161    /// The dtype this expression evaluates to.
162    pub fn dtype(&self) -> &DType {
163        match self {
164            Self::Scalar { dtype, .. } | Self::Root { dtype } => dtype,
165        }
166    }
167
168    /// The bound children of this node, in argument order. Empty for [`BoundExpression::Root`].
169    pub fn children(&self) -> &[BoundExpression] {
170        match self {
171            Self::Scalar { children, .. } => children.as_slice(),
172            Self::Root { .. } => &[],
173        }
174    }
175
176    /// Return the child at `index`.
177    pub fn child(&self, index: usize) -> &BoundExpression {
178        &self.children()[index]
179    }
180
181    /// The scalar function for this node, or `None` if it is the scope root.
182    pub fn as_scalar(&self) -> Option<&ScalarFnRef> {
183        match self {
184            Self::Scalar { scalar_fn, .. } => Some(scalar_fn),
185            Self::Root { .. } => None,
186        }
187    }
188
189    /// Return whether this node uses the given scalar-function vtable.
190    pub fn is<V: ScalarFnVTable>(&self) -> bool {
191        self.as_scalar().is_some_and(ScalarFnRef::is::<V>)
192    }
193
194    /// Return whether this expression tree contains a node using the given scalar-function vtable.
195    pub fn contains<V: ScalarFnVTable>(&self) -> VortexResult<bool> {
196        let mut contains = false;
197        pre_order_visit_down(self, |node| {
198            if node.is::<V>() {
199                contains = true;
200                return Ok(TraversalOrder::Stop);
201            }
202            Ok(TraversalOrder::Continue)
203        })?;
204        Ok(contains)
205    }
206
207    /// Return the typed scalar-function options when this node uses the given vtable.
208    pub fn as_opt<V: ScalarFnVTable>(&self) -> Option<&V::Options> {
209        self.as_scalar().and_then(ScalarFnRef::as_opt::<V>)
210    }
211
212    /// Return the typed scalar-function options for this node.
213    ///
214    /// # Panics
215    ///
216    /// Panics when this node is the scope root or uses a different scalar-function vtable.
217    pub fn as_<V: ScalarFnVTable>(&self) -> &V::Options {
218        self.as_opt::<V>()
219            .vortex_expect("Bound expression options type mismatch")
220    }
221
222    /// Whether this node is the scope root.
223    pub fn is_root(&self) -> bool {
224        matches!(self, Self::Root { .. })
225    }
226
227    /// Return whether every scope root in this expression has `dtype`.
228    ///
229    /// Expressions without a scope root, such as literals, match every dtype.
230    pub fn is_root_bound_to(&self, dtype: &DType) -> bool {
231        let mut is_bound_to = true;
232        pre_order_visit_down(self, |node| {
233            if node.is_root() && node.dtype() != dtype {
234                is_bound_to = false;
235                return Ok(TraversalOrder::Stop);
236            }
237            Ok(TraversalOrder::Continue)
238        })
239        .vortex_expect("bound expression traversal cannot not fail");
240        is_bound_to
241    }
242
243    /// Return an expression that proves this predicate is definitely false from statistics.
244    pub fn falsify(&self, session: &VortexSession) -> VortexResult<Option<BoundExpression>> {
245        StatsRewriteCtx::new(session).falsify(self)
246    }
247
248    /// Return an expression that proves this predicate is definitely true from statistics.
249    pub fn satisfy(&self, session: &VortexSession) -> VortexResult<Option<BoundExpression>> {
250        StatsRewriteCtx::new(session).satisfy(self)
251    }
252
253    /// Display the bound expression as a formatted tree structure.
254    pub fn display_tree(&self) -> impl Display {
255        DisplayTreeExpr(self)
256    }
257}
258
259impl Display for BoundExpression {
260    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
261        match self {
262            Self::Scalar { scalar_fn, .. } => scalar_fn.fmt_sql(self, f),
263            Self::Root { .. } => f.write_str("$"),
264        }
265    }
266}
267
268impl Expression {
269    /// Bind this expression against a root dtype, type-checking every node in a single walk.
270    ///
271    /// The returned tree carries a dtype on each node, so callers needing types at more than one
272    /// node should bind once and read fields rather than calling
273    /// [`return_dtype`](Expression::return_dtype) repeatedly.
274    pub fn bind(&self, dtype: &DType) -> VortexResult<BoundExpression> {
275        self.bind_scope(&Scope::new(dtype.clone()))
276    }
277
278    /// Bind this expression against an explicit [`Scope`].
279    pub fn bind_scope(&self, scope: &Scope) -> VortexResult<BoundExpression> {
280        if self.is_root() {
281            return Ok(BoundExpression::new_root(scope.root().clone()));
282        }
283
284        let children: Vec<_> = self
285            .children()
286            .iter()
287            .map(|child| child.bind_scope(scope))
288            .try_collect()?;
289        let scalar_fn = self
290            .as_scalar()
291            .vortex_expect("root was handled above, so this is a scalar node");
292        BoundExpression::try_new(scalar_fn.clone(), children)
293    }
294}
295
296/// Iterative drop to avoid stack overflows on deep trees.
297impl Drop for BoundExpression {
298    fn drop(&mut self) {
299        let Self::Scalar { children, .. } = self else {
300            return;
301        };
302        let Some(children) = Arc::get_mut(children) else {
303            return;
304        };
305
306        let mut to_drop = std::mem::take(children);
307        while let Some(mut child) = to_drop.pop() {
308            if let BoundExpression::Scalar { children, .. } = &mut child
309                && let Some(grandchildren) = Arc::get_mut(children)
310            {
311                to_drop.append(grandchildren);
312            }
313        }
314    }
315}
316
317#[cfg(test)]
318mod tests {
319    use vortex_error::VortexResult;
320
321    use super::*;
322    use crate::dtype::Nullability;
323    use crate::dtype::PType;
324    use crate::expr::col;
325    use crate::expr::eq;
326    use crate::expr::lit;
327    use crate::expr::root;
328    use crate::expr::test_harness::struct_dtype;
329    use crate::scalar_fn::fns::literal::Literal;
330
331    fn scope() -> Scope {
332        Scope::new(struct_dtype())
333    }
334
335    #[test]
336    fn root_binds_to_the_scope() -> VortexResult<()> {
337        let bound = root().bind_scope(&scope())?;
338        assert!(bound.is_root());
339        assert_eq!(bound.dtype(), &struct_dtype());
340        assert_eq!(bound, BoundExpression::new_root(struct_dtype()));
341        Ok(())
342    }
343
344    #[test]
345    fn every_node_carries_its_dtype() -> VortexResult<()> {
346        let expr = eq(col("a"), lit(1_i32));
347        let bound = expr.bind_scope(&scope())?;
348
349        assert_eq!(bound.dtype(), &DType::Bool(Nullability::NonNullable));
350
351        let lhs = &bound.children()[0];
352        assert_eq!(
353            lhs.dtype(),
354            &DType::Primitive(PType::I32, Nullability::NonNullable)
355        );
356        assert_eq!(lhs.children()[0].dtype(), &struct_dtype());
357        Ok(())
358    }
359
360    #[test]
361    fn bind_agrees_with_return_dtype() -> VortexResult<()> {
362        for expr in [root(), col("a"), eq(col("a"), lit(1_i32)), lit(true)] {
363            assert_eq!(
364                expr.bind(&struct_dtype())?.dtype(),
365                &expr.return_dtype(&struct_dtype())?,
366                "disagreement for {expr}"
367            );
368        }
369        Ok(())
370    }
371
372    #[test]
373    fn contains_scalar_function() -> VortexResult<()> {
374        let bound = eq(col("a"), lit(1_i32)).bind_scope(&scope())?;
375        assert!(bound.contains::<Literal>()?);
376        assert!(!root().bind_scope(&scope())?.contains::<Literal>()?);
377        Ok(())
378    }
379
380    #[test]
381    fn bound_to_checks_every_root() -> VortexResult<()> {
382        let dtype = struct_dtype();
383        let bound = eq(col("a"), col("a")).bind(&dtype)?;
384        assert!(bound.is_root_bound_to(&dtype));
385        assert!(!bound.is_root_bound_to(&DType::Bool(Nullability::NonNullable)));
386        assert!(
387            lit(true)
388                .bind(&dtype)?
389                .is_root_bound_to(&DType::Bool(Nullability::NonNullable))
390        );
391        Ok(())
392    }
393
394    #[test]
395    fn bound_display_matches_unbound() -> VortexResult<()> {
396        for expr in [root(), col("a"), eq(col("a"), lit(1_i32)), lit(true)] {
397            let bound = expr.bind_scope(&scope())?;
398            assert_eq!(bound.to_string(), expr.to_string());
399            assert_eq!(
400                bound.display_tree().to_string(),
401                expr.display_tree().to_string()
402            );
403        }
404        Ok(())
405    }
406
407    #[test]
408    fn clone_shares_children() -> VortexResult<()> {
409        let bound = eq(col("a"), lit(1_i32)).bind_scope(&scope())?;
410        let cloned = bound.clone();
411
412        let (
413            BoundExpression::Scalar { children: a, .. },
414            BoundExpression::Scalar { children: b, .. },
415        ) = (&bound, &cloned)
416        else {
417            unreachable!("eq is a scalar node")
418        };
419        assert!(Arc::ptr_eq(a, b));
420        Ok(())
421    }
422
423    #[test]
424    fn repeated_subtree_is_bound_per_occurrence() -> VortexResult<()> {
425        let shared = col("a");
426        let bound = eq(shared.clone(), shared).bind_scope(&scope())?;
427        let children = bound.children();
428        assert_eq!(children[0].dtype(), children[1].dtype());
429        Ok(())
430    }
431
432    #[test]
433    fn structural_and_exact_equality_are_distinct() -> VortexResult<()> {
434        let expr = eq(col("a"), lit(1_i32));
435        let bound = expr.bind_scope(&scope())?;
436        let independently_bound = expr.bind_scope(&scope())?;
437
438        assert_eq!(bound, independently_bound);
439        assert_eq!(ExactBoundExpr(bound.clone()), ExactBoundExpr(bound.clone()));
440        assert_ne!(ExactBoundExpr(bound), ExactBoundExpr(independently_bound));
441        Ok(())
442    }
443
444    #[test]
445    fn binding_reports_a_type_error() {
446        let expr = eq(col("a"), lit("nope"));
447        assert!(expr.bind_scope(&scope()).is_err());
448    }
449}
vortex_array/expr/bound_expression.rs

vortex_array/expr/
bound_expression.rs
