vortex_expr/
lib.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Vortex's expression language.
5//!
6//! All expressions are serializable, and own their own wire format.
7//!
8//! The implementation takes inspiration from [Postgres] and [Apache Datafusion].
9//!
10//! [Postgres]: https://www.postgresql.org/docs/current/sql-expressions.html
11//! [Apache Datafusion]: https://github.com/apache/datafusion/tree/5fac581efbaffd0e6a9edf931182517524526afd/datafusion/expr
12
13use std::any::Any;
14use std::fmt::{Debug, Display, Formatter};
15use std::hash::{Hash, Hasher};
16use std::sync::Arc;
17
18use dyn_hash::DynHash;
19pub use exprs::*;
20pub mod aliases;
21mod analysis;
22#[cfg(feature = "arbitrary")]
23pub mod arbitrary;
24pub mod dyn_traits;
25mod encoding;
26mod exprs;
27mod field;
28pub mod forms;
29mod operator;
30pub mod proto;
31pub mod pruning;
32mod registry;
33mod scope;
34mod scope_vars;
35pub mod transform;
36pub mod traversal;
37mod vtable;
38
39pub use analysis::*;
40pub use between::*;
41pub use binary::*;
42pub use cast::*;
43pub use encoding::*;
44pub use get_item::*;
45pub use is_null::*;
46pub use like::*;
47pub use list_contains::*;
48pub use literal::*;
49pub use merge::*;
50pub use not::*;
51pub use operator::*;
52pub use operators::*;
53pub use pack::*;
54pub use registry::*;
55pub use root::*;
56pub use scope::*;
57pub use scope_vars::*;
58pub use select::*;
59use vortex_array::{Array, ArrayRef, SerializeMetadata};
60use vortex_dtype::{DType, FieldName, FieldPath};
61use vortex_error::{VortexExpect, VortexResult, VortexUnwrap, vortex_bail};
62use vortex_utils::aliases::hash_set::HashSet;
63pub use vtable::*;
64
65pub mod display;
66
67use crate::display::{DisplayAs, DisplayFormat};
68use crate::dyn_traits::DynEq;
69use crate::traversal::{NodeExt, ReferenceCollector};
70
71pub trait IntoExpr {
72    /// Convert this type into an expression reference.
73    fn into_expr(self) -> ExprRef;
74}
75
76pub type ExprRef = Arc<dyn VortexExpr>;
77
78/// Represents logical operation on [`ArrayRef`]s
79pub trait VortexExpr:
80    'static + Send + Sync + Debug + DisplayAs + DynEq + DynHash + AnalysisExpr + private::Sealed
81{
82    /// Convert expression reference to reference of [`Any`] type
83    fn as_any(&self) -> &dyn Any;
84
85    /// Convert the expression to an [`ExprRef`].
86    fn to_expr(&self) -> ExprRef;
87
88    /// Return the encoding of the expression.
89    fn encoding(&self) -> ExprEncodingRef;
90
91    /// Serialize the metadata of this expression into a bytes vector.
92    ///
93    /// Returns `None` if the expression does not support serialization.
94    fn metadata(&self) -> Option<Vec<u8>> {
95        None
96    }
97
98    /// Compute result of expression on given batch producing a new batch
99    ///
100    /// "Unchecked" means that this function lacks a debug assertion that the returned array matches
101    /// the [VortexExpr::return_dtype] method. Use instead the
102    /// [`VortexExpr::evaluate`](./trait.VortexExpr.html#method.evaluate).
103    /// function which includes such an assertion.
104    fn unchecked_evaluate(&self, ctx: &Scope) -> VortexResult<ArrayRef>;
105
106    /// Returns the children of this expression.
107    fn children(&self) -> Vec<&ExprRef>;
108
109    /// Returns a new instance of this expression with the children replaced.
110    fn with_children(self: Arc<Self>, children: Vec<ExprRef>) -> VortexResult<ExprRef>;
111
112    /// Compute the type of the array returned by
113    /// [`VortexExpr::evaluate`](./trait.VortexExpr.html#method.evaluate).
114    fn return_dtype(&self, scope: &DType) -> VortexResult<DType>;
115}
116
117dyn_hash::hash_trait_object!(VortexExpr);
118
119impl PartialEq for dyn VortexExpr {
120    fn eq(&self, other: &Self) -> bool {
121        self.dyn_eq(other.as_any())
122    }
123}
124
125impl Eq for dyn VortexExpr {}
126
127impl dyn VortexExpr + '_ {
128    pub fn id(&self) -> ExprId {
129        self.encoding().id()
130    }
131
132    pub fn is<V: VTable>(&self) -> bool {
133        self.as_opt::<V>().is_some()
134    }
135
136    pub fn as_<V: VTable>(&self) -> &V::Expr {
137        self.as_opt::<V>()
138            .vortex_expect("Expr is not of the expected type")
139    }
140
141    pub fn as_opt<V: VTable>(&self) -> Option<&V::Expr> {
142        VortexExpr::as_any(self)
143            .downcast_ref::<ExprAdapter<V>>()
144            .map(|e| &e.0)
145    }
146
147    /// Compute result of expression on given batch producing a new batch
148    pub fn evaluate(&self, scope: &Scope) -> VortexResult<ArrayRef> {
149        let result = self.unchecked_evaluate(scope)?;
150        assert_eq!(
151            result.dtype(),
152            &self.return_dtype(scope.dtype())?,
153            "Expression {} returned dtype {} but declared return_dtype of {}",
154            &self,
155            result.dtype(),
156            self.return_dtype(scope.dtype())?,
157        );
158        Ok(result)
159    }
160
161    /// Display the expression as a formatted tree structure.
162    ///
163    /// This provides a hierarchical view of the expression that shows the relationships
164    /// between parent and child expressions, making complex nested expressions easier
165    /// to understand and debug.
166    ///
167    /// # Example
168    ///
169    /// ```rust
170    /// # use vortex_dtype::{DType, Nullability, PType};
171    /// # use vortex_expr::{and, cast, eq, get_item, gt, lit, not, root, select, IntoExpr, LikeExpr};
172    /// // Build a complex nested expression
173    /// let complex_expr = select(
174    ///     ["result"],
175    ///     and(
176    ///         not(eq(get_item("status", root()), lit("inactive"))),
177    ///         and(
178    ///             LikeExpr::new(get_item("name", root()), lit("%admin%"), false, false).into_expr(),
179    ///             gt(
180    ///                 cast(get_item("score", root()), DType::Primitive(PType::F64, Nullability::NonNullable)),
181    ///                 lit(75.0)
182    ///             )
183    ///         )
184    ///     )
185    /// );
186    ///
187    /// println!("{}", complex_expr.display_tree());
188    /// ```
189    ///
190    /// This produces output like:
191    ///
192    /// ```text
193    /// Select(include): {result}
194    /// └── Binary(and)
195    ///     ├── lhs: Not
196    ///     │   └── Binary(=)
197    ///     │       ├── lhs: GetItem(status)
198    ///     │       │   └── Root
199    ///     │       └── rhs: Literal(value: "inactive", dtype: utf8)
200    ///     └── rhs: Binary(and)
201    ///         ├── lhs: Like
202    ///         │   ├── child: GetItem(name)
203    ///         │   │   └── Root
204    ///         │   └── pattern: Literal(value: "%admin%", dtype: utf8)
205    ///         └── rhs: Binary(>)
206    ///             ├── lhs: Cast(target: f64)
207    ///             │   └── GetItem(score)
208    ///             │       └── Root
209    ///             └── rhs: Literal(value: 75f64, dtype: f64)
210    /// ```
211    pub fn display_tree(&self) -> impl Display {
212        display::DisplayTreeExpr(self)
213    }
214}
215
216impl Display for dyn VortexExpr + '_ {
217    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
218        DisplayAs::fmt_as(self, DisplayFormat::Compact, f)
219    }
220}
221
222pub trait VortexExprExt {
223    /// Accumulate all field references from this expression and its children in a set
224    fn field_references(&self) -> HashSet<FieldName>;
225}
226
227impl VortexExprExt for ExprRef {
228    fn field_references(&self) -> HashSet<FieldName> {
229        let mut collector = ReferenceCollector::new();
230        // The collector is infallible, so we can unwrap the result
231        self.accept(&mut collector).vortex_unwrap();
232        collector.into_fields()
233    }
234}
235
236#[derive(Clone)]
237#[repr(transparent)]
238pub struct ExprAdapter<V: VTable>(V::Expr);
239
240impl<V: VTable> VortexExpr for ExprAdapter<V> {
241    fn as_any(&self) -> &dyn Any {
242        self
243    }
244
245    fn to_expr(&self) -> ExprRef {
246        Arc::new(ExprAdapter::<V>(self.0.clone()))
247    }
248
249    fn encoding(&self) -> ExprEncodingRef {
250        V::encoding(&self.0)
251    }
252
253    fn metadata(&self) -> Option<Vec<u8>> {
254        V::metadata(&self.0).map(|m| m.serialize())
255    }
256
257    fn unchecked_evaluate(&self, ctx: &Scope) -> VortexResult<ArrayRef> {
258        V::evaluate(&self.0, ctx)
259    }
260
261    fn children(&self) -> Vec<&ExprRef> {
262        V::children(&self.0)
263    }
264
265    fn with_children(self: Arc<Self>, children: Vec<ExprRef>) -> VortexResult<ExprRef> {
266        if self.children().len() != children.len() {
267            vortex_bail!(
268                "Expected {} children, got {}",
269                self.children().len(),
270                children.len()
271            );
272        }
273        Ok(V::with_children(&self.0, children)?.to_expr())
274    }
275
276    fn return_dtype(&self, scope: &DType) -> VortexResult<DType> {
277        V::return_dtype(&self.0, scope)
278    }
279}
280
281impl<V: VTable> Debug for ExprAdapter<V> {
282    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
283        Debug::fmt(&self.0, f)
284    }
285}
286
287impl<V: VTable> Display for ExprAdapter<V> {
288    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
289        DisplayAs::fmt_as(&self.0, DisplayFormat::Compact, f)
290    }
291}
292
293impl<V: VTable> DisplayAs for ExprAdapter<V> {
294    fn fmt_as(&self, df: DisplayFormat, f: &mut Formatter) -> std::fmt::Result {
295        DisplayAs::fmt_as(&self.0, df, f)
296    }
297
298    fn child_names(&self) -> Option<Vec<String>> {
299        DisplayAs::child_names(&self.0)
300    }
301}
302
303impl<V: VTable> PartialEq for ExprAdapter<V> {
304    fn eq(&self, other: &Self) -> bool {
305        PartialEq::eq(&self.0, &other.0)
306    }
307}
308
309impl<V: VTable> Eq for ExprAdapter<V> {}
310
311impl<V: VTable> Hash for ExprAdapter<V> {
312    fn hash<H: Hasher>(&self, state: &mut H) {
313        Hash::hash(&self.0, state);
314    }
315}
316
317impl<V: VTable> AnalysisExpr for ExprAdapter<V> {
318    fn stat_falsification(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
319        <V::Expr as AnalysisExpr>::stat_falsification(&self.0, catalog)
320    }
321
322    fn max(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
323        <V::Expr as AnalysisExpr>::max(&self.0, catalog)
324    }
325
326    fn min(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
327        <V::Expr as AnalysisExpr>::min(&self.0, catalog)
328    }
329
330    fn nan_count(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
331        <V::Expr as AnalysisExpr>::nan_count(&self.0, catalog)
332    }
333
334    fn field_path(&self) -> Option<FieldPath> {
335        <V::Expr as AnalysisExpr>::field_path(&self.0)
336    }
337}
338
339mod private {
340    use super::*;
341
342    pub trait Sealed {}
343
344    impl<V: VTable> Sealed for ExprAdapter<V> {}
345}
346
347/// Splits top level and operations into separate expressions.
348pub fn split_conjunction(expr: &ExprRef) -> Vec<ExprRef> {
349    let mut conjunctions = vec![];
350    split_inner(expr, &mut conjunctions);
351    conjunctions
352}
353
354fn split_inner(expr: &ExprRef, exprs: &mut Vec<ExprRef>) {
355    match expr.as_opt::<BinaryVTable>() {
356        Some(bexp) if bexp.op() == Operator::And => {
357            split_inner(bexp.lhs(), exprs);
358            split_inner(bexp.rhs(), exprs);
359        }
360        Some(_) | None => {
361            exprs.push(expr.clone());
362        }
363    }
364}
365
366/// An expression wrapper that performs pointer equality.
367#[derive(Clone)]
368pub struct ExactExpr(pub ExprRef);
369
370impl PartialEq for ExactExpr {
371    fn eq(&self, other: &Self) -> bool {
372        Arc::ptr_eq(&self.0, &other.0)
373    }
374}
375
376impl Eq for ExactExpr {}
377
378impl Hash for ExactExpr {
379    fn hash<H: Hasher>(&self, state: &mut H) {
380        Arc::as_ptr(&self.0).hash(state)
381    }
382}
383
384#[cfg(feature = "test-harness")]
385pub mod test_harness {
386
387    use vortex_dtype::{DType, Nullability, PType, StructFields};
388
389    pub fn struct_dtype() -> DType {
390        DType::Struct(
391            StructFields::new(
392                ["a", "col1", "col2", "bool1", "bool2"].into(),
393                vec![
394                    DType::Primitive(PType::I32, Nullability::NonNullable),
395                    DType::Primitive(PType::U16, Nullability::Nullable),
396                    DType::Primitive(PType::U16, Nullability::Nullable),
397                    DType::Bool(Nullability::NonNullable),
398                    DType::Bool(Nullability::NonNullable),
399                ],
400            ),
401            Nullability::NonNullable,
402        )
403    }
404}
405
406#[cfg(test)]
407mod tests {
408    use vortex_dtype::{DType, FieldNames, Nullability, PType, StructFields};
409    use vortex_scalar::Scalar;
410
411    use super::*;
412
413    #[test]
414    fn basic_expr_split_test() {
415        let lhs = get_item("col1", root());
416        let rhs = lit(1);
417        let expr = eq(lhs, rhs);
418        let conjunction = split_conjunction(&expr);
419        assert_eq!(conjunction.len(), 1);
420    }
421
422    #[test]
423    fn basic_conjunction_split_test() {
424        let lhs = get_item("col1", root());
425        let rhs = lit(1);
426        let expr = and(lhs, rhs);
427        let conjunction = split_conjunction(&expr);
428        assert_eq!(conjunction.len(), 2, "Conjunction is {conjunction:?}");
429    }
430
431    #[test]
432    fn expr_display() {
433        assert_eq!(col("a").to_string(), "$.a");
434        assert_eq!(root().to_string(), "$");
435
436        let col1: Arc<dyn VortexExpr> = col("col1");
437        let col2: Arc<dyn VortexExpr> = col("col2");
438        assert_eq!(
439            and(col1.clone(), col2.clone()).to_string(),
440            "($.col1 and $.col2)"
441        );
442        assert_eq!(
443            or(col1.clone(), col2.clone()).to_string(),
444            "($.col1 or $.col2)"
445        );
446        assert_eq!(
447            eq(col1.clone(), col2.clone()).to_string(),
448            "($.col1 = $.col2)"
449        );
450        assert_eq!(
451            not_eq(col1.clone(), col2.clone()).to_string(),
452            "($.col1 != $.col2)"
453        );
454        assert_eq!(
455            gt(col1.clone(), col2.clone()).to_string(),
456            "($.col1 > $.col2)"
457        );
458        assert_eq!(
459            gt_eq(col1.clone(), col2.clone()).to_string(),
460            "($.col1 >= $.col2)"
461        );
462        assert_eq!(
463            lt(col1.clone(), col2.clone()).to_string(),
464            "($.col1 < $.col2)"
465        );
466        assert_eq!(
467            lt_eq(col1.clone(), col2.clone()).to_string(),
468            "($.col1 <= $.col2)"
469        );
470
471        assert_eq!(
472            or(
473                lt(col1.clone(), col2.clone()),
474                not_eq(col1.clone(), col2.clone()),
475            )
476            .to_string(),
477            "(($.col1 < $.col2) or ($.col1 != $.col2))"
478        );
479
480        assert_eq!(not(col1.clone()).to_string(), "(!$.col1)");
481
482        assert_eq!(
483            select(vec![FieldName::from("col1")], root()).to_string(),
484            "${col1}"
485        );
486        assert_eq!(
487            select(
488                vec![FieldName::from("col1"), FieldName::from("col2")],
489                root()
490            )
491            .to_string(),
492            "${col1, col2}"
493        );
494        assert_eq!(
495            select_exclude(
496                vec![FieldName::from("col1"), FieldName::from("col2")],
497                root()
498            )
499            .to_string(),
500            "$~{col1, col2}"
501        );
502
503        assert_eq!(lit(Scalar::from(0u8)).to_string(), "0u8");
504        assert_eq!(lit(Scalar::from(0.0f32)).to_string(), "0f32");
505        assert_eq!(
506            lit(Scalar::from(i64::MAX)).to_string(),
507            "9223372036854775807i64"
508        );
509        assert_eq!(lit(Scalar::from(true)).to_string(), "true");
510        assert_eq!(
511            lit(Scalar::null(DType::Bool(Nullability::Nullable))).to_string(),
512            "null"
513        );
514
515        assert_eq!(
516            lit(Scalar::struct_(
517                DType::Struct(
518                    StructFields::new(
519                        FieldNames::from(["dog", "cat"]),
520                        vec![
521                            DType::Primitive(PType::U32, Nullability::NonNullable),
522                            DType::Utf8(Nullability::NonNullable)
523                        ],
524                    ),
525                    Nullability::NonNullable
526                ),
527                vec![Scalar::from(32_u32), Scalar::from("rufus".to_string())]
528            ))
529            .to_string(),
530            "{dog: 32u32, cat: \"rufus\"}"
531        );
532    }
533}
vortex_expr/lib.rs

vortex_expr/
lib.rs
