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//! Filter expressions are decomposed into independent conjuncts with [`split_conjunction`] so that
35//! scans can evaluate 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
42use std::hash::Hash;
43use std::hash::Hasher;
44use std::sync::Arc;
45
46use vortex_error::VortexExpect;
47use vortex_utils::aliases::hash_set::HashSet;
48
49use crate::dtype::FieldName;
50use crate::expr::traversal::NodeExt;
51use crate::expr::traversal::ReferenceCollector;
52use crate::scalar_fn::fns::binary::Binary;
53use crate::scalar_fn::fns::operators::Operator;
54
55pub mod aliases;
56pub mod analysis;
57#[cfg(feature = "arbitrary")]
58pub mod arbitrary;
59pub mod bound_expression;
60pub mod display;
61pub(crate) mod expression;
62mod exprs;
63pub(crate) mod field;
64pub mod forms;
65mod optimize;
66pub mod proto;
67pub mod scope;
68pub mod stats;
69pub mod transform;
70pub mod traversal;
71
72pub use analysis::*;
73pub use bound_expression::*;
74pub use expression::*;
75pub use exprs::and;
76pub use exprs::and_collect;
77pub use exprs::between;
78pub use exprs::binary;
79pub use exprs::bound;
80pub use exprs::byte_length;
81pub use exprs::case_when;
82pub use exprs::case_when_no_else;
83pub use exprs::cast;
84pub use exprs::checked_add;
85pub use exprs::col;
86pub use exprs::dynamic;
87pub use exprs::dynamic_with_options;
88pub use exprs::eq;
89pub use exprs::ext_storage;
90pub use exprs::fill_null;
91pub use exprs::get_item;
92pub use exprs::gt;
93pub use exprs::gt_eq;
94pub use exprs::ilike;
95pub use exprs::is_not_null;
96pub use exprs::is_null;
97pub use exprs::is_root;
98pub use exprs::like;
99pub use exprs::list_contains;
100pub use exprs::list_length;
101pub use exprs::list_sum;
102pub use exprs::list_sum_opts;
103pub use exprs::lit;
104pub use exprs::lt;
105pub use exprs::lt_eq;
106pub use exprs::mask;
107pub use exprs::merge;
108pub use exprs::merge_opts;
109pub use exprs::nested_case_when;
110pub use exprs::not;
111pub use exprs::not_eq;
112pub use exprs::not_ilike;
113pub use exprs::not_like;
114pub use exprs::or;
115pub use exprs::or_collect;
116pub use exprs::pack;
117pub use exprs::root;
118pub use exprs::select;
119pub use exprs::select_exclude;
120pub use exprs::union_child_validities;
121pub use exprs::variant_get;
122pub use exprs::zip_expr;
123pub use scope::*;
124
125pub trait VortexExprExt {
126    /// Accumulate all field references from this expression and its children in a set
127    fn field_references(&self) -> HashSet<FieldName>;
128}
129
130impl VortexExprExt for Expression {
131    fn field_references(&self) -> HashSet<FieldName> {
132        let mut collector = ReferenceCollector::new();
133        // The collector is infallible, so we can unwrap the result
134        self.accept(&mut collector)
135            .vortex_expect("reference collector should never fail");
136        collector.into_fields()
137    }
138}
139
140/// Splits top level and operations into separate expressions.
141pub fn split_conjunction(expr: &Expression) -> Vec<Expression> {
142    let mut conjunctions = vec![];
143    split_inner(expr, &mut conjunctions);
144    conjunctions
145}
146
147fn split_inner(expr: &Expression, exprs: &mut Vec<Expression>) {
148    match expr.as_opt::<Binary>() {
149        Some(operator) if *operator == Operator::And => {
150            split_inner(expr.child(0), exprs);
151            split_inner(expr.child(1), exprs);
152        }
153        Some(_) | None => {
154            exprs.push(expr.clone());
155        }
156    }
157}
158
159/// An expression wrapper that performs pointer equality on child expressions.
160#[derive(Clone, Debug)]
161pub struct ExactExpr(pub Expression);
162impl PartialEq for ExactExpr {
163    fn eq(&self, other: &Self) -> bool {
164        match (&self.0, &other.0) {
165            (Expression::Root, Expression::Root) => true,
166            (
167                Expression::Scalar {
168                    scalar_fn: lhs_fn,
169                    children: lhs_children,
170                },
171                Expression::Scalar {
172                    scalar_fn: rhs_fn,
173                    children: rhs_children,
174                },
175            ) => lhs_fn == rhs_fn && Arc::ptr_eq(lhs_children, rhs_children),
176            _ => false,
177        }
178    }
179}
180impl Eq for ExactExpr {}
181
182impl Hash for ExactExpr {
183    fn hash<H: Hasher>(&self, state: &mut H) {
184        match &self.0 {
185            Expression::Root => state.write_u8(0),
186            Expression::Scalar {
187                scalar_fn,
188                children,
189            } => {
190                state.write_u8(1);
191                scalar_fn.hash(state);
192                Arc::as_ptr(children).hash(state);
193            }
194        }
195    }
196}
197
198#[cfg(feature = "_test-harness")]
199pub mod test_harness {
200    use crate::dtype::DType;
201    use crate::dtype::Nullability;
202    use crate::dtype::PType;
203    use crate::dtype::StructFields;
204
205    pub fn struct_dtype() -> DType {
206        DType::Struct(
207            StructFields::new(
208                ["a", "col1", "col2", "bool1", "bool2"].into(),
209                vec![
210                    DType::Primitive(PType::I32, Nullability::NonNullable),
211                    DType::Primitive(PType::U16, Nullability::Nullable),
212                    DType::Primitive(PType::U16, Nullability::Nullable),
213                    DType::Bool(Nullability::NonNullable),
214                    DType::Bool(Nullability::NonNullable),
215                ],
216            ),
217            Nullability::NonNullable,
218        )
219    }
220}
221
222#[cfg(test)]
223mod tests {
224    use std::collections::hash_map::RandomState;
225    use std::hash::BuildHasher;
226
227    use vortex_array::expr::eq;
228    use vortex_array::expr::lit;
229    use vortex_array::expr::root;
230
231    use super::*;
232    use crate::dtype::DType;
233    use crate::dtype::FieldNames;
234    use crate::dtype::Nullability;
235    use crate::dtype::PType;
236    use crate::dtype::StructFields;
237    use crate::expr::and;
238    use crate::expr::bound;
239    use crate::expr::case_when;
240    use crate::expr::col;
241    use crate::expr::get_item;
242    use crate::expr::gt;
243    use crate::expr::gt_eq;
244    use crate::expr::lt;
245    use crate::expr::lt_eq;
246    use crate::expr::not;
247    use crate::expr::not_eq;
248    use crate::expr::or;
249    use crate::expr::select;
250    use crate::expr::select_exclude;
251    use crate::scalar::Scalar;
252    use crate::scalar_fn::fns::literal::Literal;
253
254    #[test]
255    fn basic_expr_split_test() {
256        let lhs = get_item("col1", root());
257        let rhs = lit(1);
258        let expr = eq(lhs, rhs);
259        let conjunction = split_conjunction(&expr);
260        assert_eq!(conjunction.len(), 1);
261    }
262
263    #[test]
264    fn basic_conjunction_split_test() {
265        let lhs = get_item("col1", root());
266        let rhs = lit(1);
267        let expr = and(lhs, rhs);
268        let conjunction = split_conjunction(&expr);
269        assert_eq!(conjunction.len(), 2, "Conjunction is {conjunction:?}");
270    }
271
272    #[test]
273    fn exact_expr_hash_consistent_with_eq() {
274        let state = RandomState::new();
275        let expr = eq(get_item("col1", root()), lit(1));
276
277        // Clones share the children Arc, so they are equal and must hash equally.
278        let a = ExactExpr(expr.clone());
279        let b = ExactExpr(expr);
280        assert_eq!(a, b);
281        assert_eq!(state.hash_one(&a), state.hash_one(&b));
282
283        // Structurally identical expressions built separately are distinct keys.
284        let rebuilt = ExactExpr(eq(get_item("col1", root()), lit(1)));
285        assert_ne!(a, rebuilt);
286    }
287
288    #[test]
289    fn bound_constructors_preserve_order_and_types() -> vortex_error::VortexResult<()> {
290        let value_dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
291        let scope = DType::Struct(
292            StructFields::from_iter([("value", value_dtype.clone())]),
293            Nullability::NonNullable,
294        );
295
296        let root = bound::root(scope.clone());
297        let value = bound::get_item("value", root);
298        let literal = bound::lit(5i32);
299        let condition = bound::gt(value.clone(), literal.clone());
300        assert_eq!(condition.dtype(), &DType::Bool(Nullability::NonNullable));
301        assert_eq!(condition.children(), &[value.clone(), literal.clone()]);
302
303        let case = bound::case_when(condition.clone(), value.clone(), literal.clone());
304        assert_eq!(case.dtype(), &value_dtype);
305        assert_eq!(case.children(), &[condition.clone(), value, literal]);
306
307        let packed = bound::pack(
308            [("condition", condition.clone()), ("value", case.clone())],
309            Nullability::NonNullable,
310        );
311        assert_eq!(packed.children(), &[condition, case.clone()]);
312        assert_eq!(
313            packed.dtype(),
314            &DType::Struct(
315                StructFields::from_iter([
316                    ("condition", DType::Bool(Nullability::NonNullable)),
317                    ("value", value_dtype),
318                ]),
319                Nullability::NonNullable,
320            )
321        );
322
323        let unbound = case_when(gt(col("value"), lit(5i32)), col("value"), lit(5i32));
324        assert_eq!(unbound.bind(&scope)?, case);
325        Ok(())
326    }
327
328    #[test]
329    fn expr_display() {
330        assert_eq!(col("a").to_string(), "$.a");
331        assert_eq!(root().to_string(), "$");
332
333        let col1: Expression = col("col1");
334        let col2: Expression = col("col2");
335        assert_eq!(
336            and(col1.clone(), col2.clone()).to_string(),
337            "($.col1 and $.col2)"
338        );
339        assert_eq!(
340            or(col1.clone(), col2.clone()).to_string(),
341            "($.col1 or $.col2)"
342        );
343        assert_eq!(
344            eq(col1.clone(), col2.clone()).to_string(),
345            "($.col1 = $.col2)"
346        );
347        assert_eq!(
348            not_eq(col1.clone(), col2.clone()).to_string(),
349            "($.col1 != $.col2)"
350        );
351        assert_eq!(
352            gt(col1.clone(), col2.clone()).to_string(),
353            "($.col1 > $.col2)"
354        );
355        assert_eq!(
356            gt_eq(col1.clone(), col2.clone()).to_string(),
357            "($.col1 >= $.col2)"
358        );
359        assert_eq!(
360            lt(col1.clone(), col2.clone()).to_string(),
361            "($.col1 < $.col2)"
362        );
363        assert_eq!(
364            lt_eq(col1.clone(), col2.clone()).to_string(),
365            "($.col1 <= $.col2)"
366        );
367
368        assert_eq!(
369            or(lt(col1.clone(), col2.clone()), not_eq(col1.clone(), col2),).to_string(),
370            "(($.col1 < $.col2) or ($.col1 != $.col2))"
371        );
372
373        assert_eq!(not(col1).to_string(), "vortex.not($.col1)");
374
375        assert_eq!(
376            select(vec![FieldName::from("col1")], root()).to_string(),
377            "${col1}"
378        );
379        assert_eq!(
380            select(
381                vec![FieldName::from("col1"), FieldName::from("col2")],
382                root()
383            )
384            .to_string(),
385            "${col1, col2}"
386        );
387        assert_eq!(
388            select_exclude(
389                vec![FieldName::from("col1"), FieldName::from("col2")],
390                root()
391            )
392            .to_string(),
393            "${~ col1, col2}"
394        );
395
396        assert_eq!(lit(Scalar::from(0u8)).to_string(), "0u8");
397        assert_eq!(lit(Scalar::from(0.0f32)).to_string(), "0f32");
398        assert_eq!(
399            lit(Scalar::from(i64::MAX)).to_string(),
400            "9223372036854775807i64"
401        );
402        assert_eq!(lit(Scalar::from(true)).to_string(), "true");
403        assert_eq!(
404            lit(Scalar::null(DType::Bool(Nullability::Nullable))).to_string(),
405            "null"
406        );
407
408        assert_eq!(
409            lit(Scalar::struct_(
410                DType::Struct(
411                    StructFields::new(
412                        FieldNames::from(["dog", "cat"]),
413                        vec![
414                            DType::Primitive(PType::U32, Nullability::NonNullable),
415                            DType::Utf8(Nullability::NonNullable)
416                        ],
417                    ),
418                    Nullability::NonNullable
419                ),
420                vec![Scalar::from(32_u32), Scalar::from("rufus".to_string())]
421            ))
422            .to_string(),
423            "{dog: 32u32, cat: \"rufus\"}"
424        );
425    }
426
427    #[test]
428    fn expr_contains() {
429        let expression = &eq(root(), lit(3u64));
430        assert!(expression.contains::<Literal>().unwrap());
431        let expression = root();
432        assert!(!expression.contains::<Literal>().unwrap());
433    }
434}
vortex_array/expr/mod.rs

vortex_array/expr/
mod.rs
