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::is_nan;
84pub use exprs::is_not_null;
85pub use exprs::is_null;
86pub use exprs::is_root;
87pub use exprs::like;
88pub use exprs::list_contains;
89pub use exprs::list_contains_opts;
90pub use exprs::list_length;
91pub use exprs::list_sum;
92pub use exprs::list_sum_opts;
93pub use exprs::lit;
94pub use exprs::lt;
95pub use exprs::lt_eq;
96pub use exprs::mask;
97pub use exprs::merge;
98pub use exprs::merge_opts;
99pub use exprs::nested_case_when;
100pub use exprs::not;
101pub use exprs::not_eq;
102pub use exprs::not_ilike;
103pub use exprs::not_like;
104pub use exprs::or;
105pub use exprs::or_collect;
106pub use exprs::pack;
107pub use exprs::root;
108pub use exprs::select;
109pub use exprs::select_exclude;
110pub use exprs::variant_get;
111pub use exprs::zip_expr;
112pub use scope::*;
113
114#[cfg(feature = "_test-harness")]
115pub mod test_harness {
116    use crate::dtype::DType;
117    use crate::dtype::Nullability;
118    use crate::dtype::PType;
119    use crate::dtype::StructFields;
120
121    pub fn struct_dtype() -> DType {
122        DType::Struct(
123            StructFields::new(
124                ["a", "col1", "col2", "bool1", "bool2"].into(),
125                vec![
126                    DType::Primitive(PType::I32, Nullability::NonNullable),
127                    DType::Primitive(PType::U16, Nullability::Nullable),
128                    DType::Primitive(PType::U16, Nullability::Nullable),
129                    DType::Bool(Nullability::NonNullable),
130                    DType::Bool(Nullability::NonNullable),
131                ],
132            ),
133            Nullability::NonNullable,
134        )
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use vortex_array::expr::eq;
141    use vortex_array::expr::lit;
142    use vortex_array::expr::root;
143
144    use super::*;
145    use crate::dtype::DType;
146    use crate::dtype::FieldName;
147    use crate::dtype::FieldNames;
148    use crate::dtype::Nullability;
149    use crate::dtype::PType;
150    use crate::dtype::StructFields;
151    use crate::expr::and;
152    use crate::expr::bound;
153    use crate::expr::case_when;
154    use crate::expr::col;
155    use crate::expr::gt;
156    use crate::expr::gt_eq;
157    use crate::expr::lt;
158    use crate::expr::lt_eq;
159    use crate::expr::not;
160    use crate::expr::not_eq;
161    use crate::expr::or;
162    use crate::expr::select;
163    use crate::expr::select_exclude;
164    use crate::scalar::Scalar;
165    use crate::scalar_fn::fns::literal::Literal;
166
167    #[test]
168    fn bound_constructors_preserve_order_and_types() -> vortex_error::VortexResult<()> {
169        let value_dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
170        let scope = DType::Struct(
171            StructFields::from_iter([("value", value_dtype.clone())]),
172            Nullability::NonNullable,
173        );
174
175        let root = bound::root(scope.clone());
176        let value = bound::get_item("value", root);
177        let literal = bound::lit(5i32);
178        let condition = bound::gt(value.clone(), literal.clone());
179        assert_eq!(condition.dtype(), &DType::Bool(Nullability::NonNullable));
180        assert_eq!(condition.children(), &[value.clone(), literal.clone()]);
181
182        let case = bound::case_when(condition.clone(), value.clone(), literal.clone());
183        assert_eq!(case.dtype(), &value_dtype);
184        assert_eq!(case.children(), &[condition.clone(), value, literal]);
185
186        let packed = bound::pack(
187            [("condition", condition.clone()), ("value", case.clone())],
188            Nullability::NonNullable,
189        );
190        assert_eq!(packed.children(), &[condition, case.clone()]);
191        assert_eq!(
192            packed.dtype(),
193            &DType::Struct(
194                StructFields::from_iter([
195                    ("condition", DType::Bool(Nullability::NonNullable)),
196                    ("value", value_dtype),
197                ]),
198                Nullability::NonNullable,
199            )
200        );
201
202        let unbound = case_when(gt(col("value"), lit(5i32)), col("value"), lit(5i32));
203        assert_eq!(unbound.bind(&scope)?, case);
204        Ok(())
205    }
206
207    #[test]
208    fn expr_display() {
209        assert_eq!(col("a").to_string(), "$.a");
210        assert_eq!(root().to_string(), "$");
211
212        let col1: Expression = col("col1");
213        let col2: Expression = col("col2");
214        assert_eq!(
215            and(col1.clone(), col2.clone()).to_string(),
216            "($.col1 and $.col2)"
217        );
218        assert_eq!(
219            or(col1.clone(), col2.clone()).to_string(),
220            "($.col1 or $.col2)"
221        );
222        assert_eq!(
223            eq(col1.clone(), col2.clone()).to_string(),
224            "($.col1 = $.col2)"
225        );
226        assert_eq!(
227            not_eq(col1.clone(), col2.clone()).to_string(),
228            "($.col1 != $.col2)"
229        );
230        assert_eq!(
231            gt(col1.clone(), col2.clone()).to_string(),
232            "($.col1 > $.col2)"
233        );
234        assert_eq!(
235            gt_eq(col1.clone(), col2.clone()).to_string(),
236            "($.col1 >= $.col2)"
237        );
238        assert_eq!(
239            lt(col1.clone(), col2.clone()).to_string(),
240            "($.col1 < $.col2)"
241        );
242        assert_eq!(
243            lt_eq(col1.clone(), col2.clone()).to_string(),
244            "($.col1 <= $.col2)"
245        );
246
247        assert_eq!(
248            or(lt(col1.clone(), col2.clone()), not_eq(col1.clone(), col2),).to_string(),
249            "(($.col1 < $.col2) or ($.col1 != $.col2))"
250        );
251
252        assert_eq!(not(col1).to_string(), "vortex.not($.col1)");
253
254        assert_eq!(
255            select(vec![FieldName::from("col1")], root()).to_string(),
256            "${col1}"
257        );
258        assert_eq!(
259            select(
260                vec![FieldName::from("col1"), FieldName::from("col2")],
261                root()
262            )
263            .to_string(),
264            "${col1, col2}"
265        );
266        assert_eq!(
267            select_exclude(
268                vec![FieldName::from("col1"), FieldName::from("col2")],
269                root()
270            )
271            .to_string(),
272            "${~ col1, col2}"
273        );
274
275        assert_eq!(lit(Scalar::from(0u8)).to_string(), "0u8");
276        assert_eq!(lit(Scalar::from(0.0f32)).to_string(), "0f32");
277        assert_eq!(
278            lit(Scalar::from(i64::MAX)).to_string(),
279            "9223372036854775807i64"
280        );
281        assert_eq!(lit(Scalar::from(true)).to_string(), "true");
282        assert_eq!(
283            lit(Scalar::null(DType::Bool(Nullability::Nullable))).to_string(),
284            "null"
285        );
286
287        assert_eq!(
288            lit(Scalar::struct_(
289                DType::Struct(
290                    StructFields::new(
291                        FieldNames::from(["dog", "cat"]),
292                        vec![
293                            DType::Primitive(PType::U32, Nullability::NonNullable),
294                            DType::Utf8(Nullability::NonNullable)
295                        ],
296                    ),
297                    Nullability::NonNullable
298                ),
299                vec![Scalar::from(32_u32), Scalar::from("rufus".to_string())]
300            ))
301            .to_string(),
302            "{dog: 32u32, cat: \"rufus\"}"
303        );
304    }
305
306    #[test]
307    fn expr_contains() {
308        let expression = &eq(root(), lit(3u64));
309        assert!(expression.contains::<Literal>().unwrap());
310        let expression = root();
311        assert!(!expression.contains::<Literal>().unwrap());
312    }
313}