Skip to main content

vortex_array/expr/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Vortex's expression language: scalar operations over [arrays](crate::ArrayRef).
5//!
6//! An [`Expression`] is a tree of scalar operations rooted at a scope (see [`root`]). Expressions
7//! are the common currency of scans: a scan takes a *filter* expression that resolves to a boolean
8//! and a *projection* expression that shapes the output. All expressions are serializable and own
9//! their own wire format, so they can be pushed down to remote sources and reconstructed on workers.
10//!
11//! # Scalar functions
12//!
13//! Each node references a scalar function defined by a
14//! [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable). The vtable declares the function signature,
15//! properties such as strictness, and the logic that executes it over input arrays. Built-in
16//! functions live in [`crate::scalar_fn`]; integration and plugin crates supply additional,
17//! use-case-specific functions.
18//!
19//! # Deferred execution
20//!
21//! Applying an expression to an array does not compute the result eagerly. Instead it builds a
22//! [`ScalarFnArray`](crate::arrays::ScalarFnArray) representing the deferred application, letting
23//! downstream encodings push the computation into compressed data, or fuse several expressions
24//! together, before any data is materialized. The deferred tree is executed toward canonical form
25//! only when a result is actually required.
26//!
27//! # Type checking
28//!
29//! Expressions are strictly typed: an input array's dtype must match the function signature exactly,
30//! so callers perform any required casts themselves before building the expression. The one
31//! relaxation is nullability—for example, equality may compare a `u32` against a `u32?`, but never
32//! a `u32` against an `i32`.
33//!
34//! Scans decompose bound filter expressions into independent conjuncts so that they can evaluate
35//! and reorder the most selective predicates first.
36//!
37//! The implementation takes inspiration from [Postgres] and [Apache Datafusion].
38//!
39//! [Postgres]: https://www.postgresql.org/docs/current/sql-expressions.html
40//! [Apache Datafusion]: https://github.com/apache/datafusion/tree/5fac581efbaffd0e6a9edf931182517524526afd/datafusion/expr
41
42pub mod aliases;
43pub mod analysis;
44#[cfg(feature = "arbitrary")]
45pub mod arbitrary;
46pub mod bound_expression;
47pub mod display;
48pub(crate) mod expression;
49mod exprs;
50pub(crate) mod field;
51pub mod forms;
52mod optimize;
53pub mod proto;
54pub mod scope;
55pub mod stats;
56pub mod transform;
57pub mod traversal;
58
59pub use analysis::*;
60pub use bound_expression::*;
61pub use expression::*;
62pub use exprs::and;
63pub use exprs::and_collect;
64pub use exprs::between;
65pub use exprs::binary;
66pub use exprs::bound;
67pub use exprs::byte_length;
68pub use exprs::case_when;
69pub use exprs::case_when_no_else;
70pub use exprs::cast;
71pub use exprs::checked_add;
72pub use exprs::col;
73pub use exprs::dynamic;
74pub use exprs::dynamic_with_options;
75pub use exprs::eq;
76pub use exprs::ext_storage;
77pub use exprs::fill_null;
78pub use exprs::get_item;
79pub use exprs::gt;
80pub use exprs::gt_eq;
81pub use exprs::ilike;
82pub use exprs::in_list;
83pub use exprs::in_list_of;
84pub use exprs::is_nan;
85pub use exprs::is_not_null;
86pub use exprs::is_null;
87pub use exprs::is_root;
88pub use exprs::like;
89pub use exprs::list_contains;
90pub use exprs::list_contains_opts;
91pub use exprs::list_length;
92pub use exprs::list_sum;
93pub use exprs::list_sum_opts;
94pub use exprs::lit;
95pub use exprs::lt;
96pub use exprs::lt_eq;
97pub use exprs::mask;
98pub use exprs::merge;
99pub use exprs::merge_opts;
100pub use exprs::nested_case_when;
101pub use exprs::not;
102pub use exprs::not_eq;
103pub use exprs::not_ilike;
104pub use exprs::not_like;
105pub use exprs::or;
106pub use exprs::or_collect;
107pub use exprs::pack;
108pub use exprs::replace_time_zone;
109pub use exprs::root;
110pub use exprs::select;
111pub use exprs::select_exclude;
112pub use exprs::variant_get;
113pub use exprs::zip_expr;
114pub use scope::*;
115
116#[cfg(feature = "_test-harness")]
117pub mod test_harness {
118    use crate::dtype::DType;
119    use crate::dtype::Nullability;
120    use crate::dtype::PType;
121    use crate::dtype::StructFields;
122
123    pub fn struct_dtype() -> DType {
124        DType::Struct(
125            StructFields::new(
126                ["a", "col1", "col2", "bool1", "bool2"].into(),
127                vec![
128                    DType::Primitive(PType::I32, Nullability::NonNullable),
129                    DType::Primitive(PType::U16, Nullability::Nullable),
130                    DType::Primitive(PType::U16, Nullability::Nullable),
131                    DType::Bool(Nullability::NonNullable),
132                    DType::Bool(Nullability::NonNullable),
133                ],
134            ),
135            Nullability::NonNullable,
136        )
137    }
138}
139
140#[cfg(test)]
141mod tests {
142    use vortex_array::expr::eq;
143    use vortex_array::expr::lit;
144    use vortex_array::expr::root;
145
146    use super::*;
147    use crate::dtype::DType;
148    use crate::dtype::FieldName;
149    use crate::dtype::FieldNames;
150    use crate::dtype::Nullability;
151    use crate::dtype::PType;
152    use crate::dtype::StructFields;
153    use crate::expr::and;
154    use crate::expr::bound;
155    use crate::expr::case_when;
156    use crate::expr::col;
157    use crate::expr::gt;
158    use crate::expr::gt_eq;
159    use crate::expr::lt;
160    use crate::expr::lt_eq;
161    use crate::expr::not;
162    use crate::expr::not_eq;
163    use crate::expr::or;
164    use crate::expr::select;
165    use crate::expr::select_exclude;
166    use crate::scalar::Scalar;
167    use crate::scalar_fn::fns::literal::Literal;
168
169    #[test]
170    fn bound_constructors_preserve_order_and_types() -> vortex_error::VortexResult<()> {
171        let value_dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
172        let scope = DType::Struct(
173            StructFields::from_iter([("value", value_dtype.clone())]),
174            Nullability::NonNullable,
175        );
176
177        let root = bound::root(scope.clone());
178        let value = bound::get_item("value", root);
179        let literal = bound::lit(5i32);
180        let condition = bound::gt(value.clone(), literal.clone());
181        assert_eq!(condition.dtype(), &DType::Bool(Nullability::NonNullable));
182        assert_eq!(condition.children(), &[value.clone(), literal.clone()]);
183
184        let case = bound::case_when(condition.clone(), value.clone(), literal.clone());
185        assert_eq!(case.dtype(), &value_dtype);
186        assert_eq!(case.children(), &[condition.clone(), value, literal]);
187
188        let packed = bound::pack(
189            [("condition", condition.clone()), ("value", case.clone())],
190            Nullability::NonNullable,
191        );
192        assert_eq!(packed.children(), &[condition, case.clone()]);
193        assert_eq!(
194            packed.dtype(),
195            &DType::Struct(
196                StructFields::from_iter([
197                    ("condition", DType::Bool(Nullability::NonNullable)),
198                    ("value", value_dtype),
199                ]),
200                Nullability::NonNullable,
201            )
202        );
203
204        let unbound = case_when(gt(col("value"), lit(5i32)), col("value"), lit(5i32));
205        assert_eq!(unbound.bind(&scope)?, case);
206        Ok(())
207    }
208
209    #[test]
210    fn expr_display() {
211        assert_eq!(col("a").to_string(), "$.a");
212        assert_eq!(root().to_string(), "$");
213
214        let col1: Expression = col("col1");
215        let col2: Expression = col("col2");
216        assert_eq!(
217            and(col1.clone(), col2.clone()).to_string(),
218            "($.col1 and $.col2)"
219        );
220        assert_eq!(
221            or(col1.clone(), col2.clone()).to_string(),
222            "($.col1 or $.col2)"
223        );
224        assert_eq!(
225            eq(col1.clone(), col2.clone()).to_string(),
226            "($.col1 = $.col2)"
227        );
228        assert_eq!(
229            not_eq(col1.clone(), col2.clone()).to_string(),
230            "($.col1 != $.col2)"
231        );
232        assert_eq!(
233            gt(col1.clone(), col2.clone()).to_string(),
234            "($.col1 > $.col2)"
235        );
236        assert_eq!(
237            gt_eq(col1.clone(), col2.clone()).to_string(),
238            "($.col1 >= $.col2)"
239        );
240        assert_eq!(
241            lt(col1.clone(), col2.clone()).to_string(),
242            "($.col1 < $.col2)"
243        );
244        assert_eq!(
245            lt_eq(col1.clone(), col2.clone()).to_string(),
246            "($.col1 <= $.col2)"
247        );
248
249        assert_eq!(
250            or(lt(col1.clone(), col2.clone()), not_eq(col1.clone(), col2),).to_string(),
251            "(($.col1 < $.col2) or ($.col1 != $.col2))"
252        );
253
254        assert_eq!(not(col1).to_string(), "vortex.not($.col1)");
255
256        assert_eq!(
257            select(vec![FieldName::from("col1")], root()).to_string(),
258            "${col1}"
259        );
260        assert_eq!(
261            select(
262                vec![FieldName::from("col1"), FieldName::from("col2")],
263                root()
264            )
265            .to_string(),
266            "${col1, col2}"
267        );
268        assert_eq!(
269            select_exclude(
270                vec![FieldName::from("col1"), FieldName::from("col2")],
271                root()
272            )
273            .to_string(),
274            "${~ col1, col2}"
275        );
276
277        assert_eq!(lit(Scalar::from(0u8)).to_string(), "0u8");
278        assert_eq!(lit(Scalar::from(0.0f32)).to_string(), "0f32");
279        assert_eq!(
280            lit(Scalar::from(i64::MAX)).to_string(),
281            "9223372036854775807i64"
282        );
283        assert_eq!(lit(Scalar::from(true)).to_string(), "true");
284        assert_eq!(
285            lit(Scalar::null(DType::Bool(Nullability::Nullable))).to_string(),
286            "null"
287        );
288
289        assert_eq!(
290            lit(Scalar::struct_(
291                DType::Struct(
292                    StructFields::new(
293                        FieldNames::from(["dog", "cat"]),
294                        vec![
295                            DType::Primitive(PType::U32, Nullability::NonNullable),
296                            DType::Utf8(Nullability::NonNullable)
297                        ],
298                    ),
299                    Nullability::NonNullable
300                ),
301                vec![Scalar::from(32_u32), Scalar::from("rufus".to_string())]
302            ))
303            .to_string(),
304            "{dog: 32u32, cat: \"rufus\"}"
305        );
306    }
307
308    #[test]
309    fn expr_contains() {
310        let expression = &eq(root(), lit(3u64));
311        assert!(expression.contains::<Literal>().unwrap());
312        let expression = root();
313        assert!(!expression.contains::<Literal>().unwrap());
314    }
315}