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//! # Typing and coercion
28//!
29//! Expressions are strictly typed: an input array's dtype must match the function signature exactly,
30//! so callers perform any required type coercion themselves before building the expression (see the
31//! [`transform`] passes). The one relaxation is null-coercion — for example, equality may compare a
32//! `u32` against a `u32?`, but never 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        self.0.scalar_fn() == other.0.scalar_fn()
165            && Arc::ptr_eq(self.0.children(), other.0.children())
166    }
167}
168impl Eq for ExactExpr {}
169
170impl Hash for ExactExpr {
171    fn hash<H: Hasher>(&self, state: &mut H) {
172        self.0.scalar_fn().hash(state);
173        Arc::as_ptr(self.0.children()).hash(state);
174    }
175}
176
177#[cfg(feature = "_test-harness")]
178pub mod test_harness {
179    use crate::dtype::DType;
180    use crate::dtype::Nullability;
181    use crate::dtype::PType;
182    use crate::dtype::StructFields;
183
184    pub fn struct_dtype() -> DType {
185        DType::Struct(
186            StructFields::new(
187                ["a", "col1", "col2", "bool1", "bool2"].into(),
188                vec![
189                    DType::Primitive(PType::I32, Nullability::NonNullable),
190                    DType::Primitive(PType::U16, Nullability::Nullable),
191                    DType::Primitive(PType::U16, Nullability::Nullable),
192                    DType::Bool(Nullability::NonNullable),
193                    DType::Bool(Nullability::NonNullable),
194                ],
195            ),
196            Nullability::NonNullable,
197        )
198    }
199}
200
201#[cfg(test)]
202mod tests {
203    use std::collections::hash_map::RandomState;
204    use std::hash::BuildHasher;
205
206    use vortex_array::expr::eq;
207    use vortex_array::expr::lit;
208    use vortex_array::expr::root;
209
210    use super::*;
211    use crate::dtype::DType;
212    use crate::dtype::FieldNames;
213    use crate::dtype::Nullability;
214    use crate::dtype::PType;
215    use crate::dtype::StructFields;
216    use crate::expr::and;
217    use crate::expr::bound;
218    use crate::expr::case_when;
219    use crate::expr::col;
220    use crate::expr::get_item;
221    use crate::expr::gt;
222    use crate::expr::gt_eq;
223    use crate::expr::lt;
224    use crate::expr::lt_eq;
225    use crate::expr::not;
226    use crate::expr::not_eq;
227    use crate::expr::or;
228    use crate::expr::select;
229    use crate::expr::select_exclude;
230    use crate::scalar::Scalar;
231    use crate::scalar_fn::fns::literal::Literal;
232
233    #[test]
234    fn basic_expr_split_test() {
235        let lhs = get_item("col1", root());
236        let rhs = lit(1);
237        let expr = eq(lhs, rhs);
238        let conjunction = split_conjunction(&expr);
239        assert_eq!(conjunction.len(), 1);
240    }
241
242    #[test]
243    fn basic_conjunction_split_test() {
244        let lhs = get_item("col1", root());
245        let rhs = lit(1);
246        let expr = and(lhs, rhs);
247        let conjunction = split_conjunction(&expr);
248        assert_eq!(conjunction.len(), 2, "Conjunction is {conjunction:?}");
249    }
250
251    #[test]
252    fn exact_expr_hash_consistent_with_eq() {
253        let state = RandomState::new();
254        let expr = eq(get_item("col1", root()), lit(1));
255
256        // Clones share the children Arc, so they are equal and must hash equally.
257        let a = ExactExpr(expr.clone());
258        let b = ExactExpr(expr);
259        assert_eq!(a, b);
260        assert_eq!(state.hash_one(&a), state.hash_one(&b));
261
262        // Structurally identical expressions built separately are distinct keys.
263        let rebuilt = ExactExpr(eq(get_item("col1", root()), lit(1)));
264        assert_ne!(a, rebuilt);
265    }
266
267    #[test]
268    fn bound_constructors_preserve_order_and_types() -> vortex_error::VortexResult<()> {
269        let value_dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
270        let scope = DType::Struct(
271            StructFields::from_iter([("value", value_dtype.clone())]),
272            Nullability::NonNullable,
273        );
274
275        let root = bound::root(scope.clone());
276        let value = bound::get_item("value", root);
277        let literal = bound::lit(5i32);
278        let condition = bound::gt(value.clone(), literal.clone());
279        assert_eq!(condition.dtype(), &DType::Bool(Nullability::NonNullable));
280        assert_eq!(condition.children(), &[value.clone(), literal.clone()]);
281
282        let case = bound::case_when(condition.clone(), value.clone(), literal.clone());
283        assert_eq!(case.dtype(), &value_dtype);
284        assert_eq!(case.children(), &[condition.clone(), value, literal]);
285
286        let packed = bound::pack(
287            [("condition", condition.clone()), ("value", case.clone())],
288            Nullability::NonNullable,
289        );
290        assert_eq!(packed.children(), &[condition, case.clone()]);
291        assert_eq!(
292            packed.dtype(),
293            &DType::Struct(
294                StructFields::from_iter([
295                    ("condition", DType::Bool(Nullability::NonNullable)),
296                    ("value", value_dtype),
297                ]),
298                Nullability::NonNullable,
299            )
300        );
301
302        let unbound = case_when(gt(col("value"), lit(5i32)), col("value"), lit(5i32));
303        assert_eq!(unbound.bind(&scope)?, case);
304        Ok(())
305    }
306
307    #[test]
308    fn expr_display() {
309        assert_eq!(col("a").to_string(), "$.a");
310        assert_eq!(root().to_string(), "$");
311
312        let col1: Expression = col("col1");
313        let col2: Expression = col("col2");
314        assert_eq!(
315            and(col1.clone(), col2.clone()).to_string(),
316            "($.col1 and $.col2)"
317        );
318        assert_eq!(
319            or(col1.clone(), col2.clone()).to_string(),
320            "($.col1 or $.col2)"
321        );
322        assert_eq!(
323            eq(col1.clone(), col2.clone()).to_string(),
324            "($.col1 = $.col2)"
325        );
326        assert_eq!(
327            not_eq(col1.clone(), col2.clone()).to_string(),
328            "($.col1 != $.col2)"
329        );
330        assert_eq!(
331            gt(col1.clone(), col2.clone()).to_string(),
332            "($.col1 > $.col2)"
333        );
334        assert_eq!(
335            gt_eq(col1.clone(), col2.clone()).to_string(),
336            "($.col1 >= $.col2)"
337        );
338        assert_eq!(
339            lt(col1.clone(), col2.clone()).to_string(),
340            "($.col1 < $.col2)"
341        );
342        assert_eq!(
343            lt_eq(col1.clone(), col2.clone()).to_string(),
344            "($.col1 <= $.col2)"
345        );
346
347        assert_eq!(
348            or(lt(col1.clone(), col2.clone()), not_eq(col1.clone(), col2),).to_string(),
349            "(($.col1 < $.col2) or ($.col1 != $.col2))"
350        );
351
352        assert_eq!(not(col1).to_string(), "vortex.not($.col1)");
353
354        assert_eq!(
355            select(vec![FieldName::from("col1")], root()).to_string(),
356            "${col1}"
357        );
358        assert_eq!(
359            select(
360                vec![FieldName::from("col1"), FieldName::from("col2")],
361                root()
362            )
363            .to_string(),
364            "${col1, col2}"
365        );
366        assert_eq!(
367            select_exclude(
368                vec![FieldName::from("col1"), FieldName::from("col2")],
369                root()
370            )
371            .to_string(),
372            "${~ col1, col2}"
373        );
374
375        assert_eq!(lit(Scalar::from(0u8)).to_string(), "0u8");
376        assert_eq!(lit(Scalar::from(0.0f32)).to_string(), "0f32");
377        assert_eq!(
378            lit(Scalar::from(i64::MAX)).to_string(),
379            "9223372036854775807i64"
380        );
381        assert_eq!(lit(Scalar::from(true)).to_string(), "true");
382        assert_eq!(
383            lit(Scalar::null(DType::Bool(Nullability::Nullable))).to_string(),
384            "null"
385        );
386
387        assert_eq!(
388            lit(Scalar::struct_(
389                DType::Struct(
390                    StructFields::new(
391                        FieldNames::from(["dog", "cat"]),
392                        vec![
393                            DType::Primitive(PType::U32, Nullability::NonNullable),
394                            DType::Utf8(Nullability::NonNullable)
395                        ],
396                    ),
397                    Nullability::NonNullable
398                ),
399                vec![Scalar::from(32_u32), Scalar::from("rufus".to_string())]
400            ))
401            .to_string(),
402            "{dog: 32u32, cat: \"rufus\"}"
403        );
404    }
405
406    #[test]
407    fn expr_contains() {
408        let expression = &eq(root(), lit(3u64));
409        assert!(expression.contains::<Literal>().unwrap());
410        let expression = root();
411        assert!(!expression.contains::<Literal>().unwrap());
412    }
413}
vortex_array/expr/mod.rs

vortex_array/expr/
mod.rs
