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 null-sensitivity, 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 display;
60pub(crate) mod expression;
61mod exprs;
62pub(crate) mod field;
63pub mod forms;
64mod optimize;
65pub mod proto;
66pub mod stats;
67pub mod transform;
68pub mod traversal;
69
70pub use analysis::*;
71pub use expression::*;
72pub use exprs::*;
73
74pub trait VortexExprExt {
75    /// Accumulate all field references from this expression and its children in a set
76    fn field_references(&self) -> HashSet<FieldName>;
77}
78
79impl VortexExprExt for Expression {
80    fn field_references(&self) -> HashSet<FieldName> {
81        let mut collector = ReferenceCollector::new();
82        // The collector is infallible, so we can unwrap the result
83        self.accept(&mut collector)
84            .vortex_expect("reference collector should never fail");
85        collector.into_fields()
86    }
87}
88
89/// Splits top level and operations into separate expressions.
90pub fn split_conjunction(expr: &Expression) -> Vec<Expression> {
91    let mut conjunctions = vec![];
92    split_inner(expr, &mut conjunctions);
93    conjunctions
94}
95
96fn split_inner(expr: &Expression, exprs: &mut Vec<Expression>) {
97    match expr.as_opt::<Binary>() {
98        Some(operator) if *operator == Operator::And => {
99            split_inner(expr.child(0), exprs);
100            split_inner(expr.child(1), exprs);
101        }
102        Some(_) | None => {
103            exprs.push(expr.clone());
104        }
105    }
106}
107
108/// An expression wrapper that performs pointer equality on child expressions.
109#[derive(Clone, Debug)]
110pub struct ExactExpr(pub Expression);
111impl PartialEq for ExactExpr {
112    fn eq(&self, other: &Self) -> bool {
113        self.0.scalar_fn() == other.0.scalar_fn()
114            && Arc::ptr_eq(self.0.children(), other.0.children())
115    }
116}
117impl Eq for ExactExpr {}
118
119impl Hash for ExactExpr {
120    fn hash<H: Hasher>(&self, state: &mut H) {
121        self.0.scalar_fn().hash(state);
122        Arc::as_ptr(self.0.children()).hash(state);
123    }
124}
125
126#[cfg(feature = "_test-harness")]
127pub mod test_harness {
128    use crate::dtype::DType;
129    use crate::dtype::Nullability;
130    use crate::dtype::PType;
131    use crate::dtype::StructFields;
132
133    pub fn struct_dtype() -> DType {
134        DType::Struct(
135            StructFields::new(
136                ["a", "col1", "col2", "bool1", "bool2"].into(),
137                vec![
138                    DType::Primitive(PType::I32, Nullability::NonNullable),
139                    DType::Primitive(PType::U16, Nullability::Nullable),
140                    DType::Primitive(PType::U16, Nullability::Nullable),
141                    DType::Bool(Nullability::NonNullable),
142                    DType::Bool(Nullability::NonNullable),
143                ],
144            ),
145            Nullability::NonNullable,
146        )
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use std::collections::hash_map::RandomState;
153    use std::hash::BuildHasher;
154
155    use super::*;
156    use crate::dtype::DType;
157    use crate::dtype::FieldNames;
158    use crate::dtype::Nullability;
159    use crate::dtype::PType;
160    use crate::dtype::StructFields;
161    use crate::expr::and;
162    use crate::expr::col;
163    use crate::expr::eq;
164    use crate::expr::get_item;
165    use crate::expr::gt;
166    use crate::expr::gt_eq;
167    use crate::expr::lit;
168    use crate::expr::lt;
169    use crate::expr::lt_eq;
170    use crate::expr::not;
171    use crate::expr::not_eq;
172    use crate::expr::or;
173    use crate::expr::root;
174    use crate::expr::select;
175    use crate::expr::select_exclude;
176    use crate::scalar::Scalar;
177
178    #[test]
179    fn basic_expr_split_test() {
180        let lhs = get_item("col1", root());
181        let rhs = lit(1);
182        let expr = eq(lhs, rhs);
183        let conjunction = split_conjunction(&expr);
184        assert_eq!(conjunction.len(), 1);
185    }
186
187    #[test]
188    fn basic_conjunction_split_test() {
189        let lhs = get_item("col1", root());
190        let rhs = lit(1);
191        let expr = and(lhs, rhs);
192        let conjunction = split_conjunction(&expr);
193        assert_eq!(conjunction.len(), 2, "Conjunction is {conjunction:?}");
194    }
195
196    #[test]
197    fn exact_expr_hash_consistent_with_eq() {
198        let state = RandomState::new();
199        let expr = eq(get_item("col1", root()), lit(1));
200
201        // Clones share the children Arc, so they are equal and must hash equally.
202        let a = ExactExpr(expr.clone());
203        let b = ExactExpr(expr);
204        assert_eq!(a, b);
205        assert_eq!(state.hash_one(&a), state.hash_one(&b));
206
207        // Structurally identical expressions built separately are distinct keys.
208        let rebuilt = ExactExpr(eq(get_item("col1", root()), lit(1)));
209        assert_ne!(a, rebuilt);
210    }
211
212    #[test]
213    fn expr_display() {
214        assert_eq!(col("a").to_string(), "$.a");
215        assert_eq!(root().to_string(), "$");
216
217        let col1: Expression = col("col1");
218        let col2: Expression = col("col2");
219        assert_eq!(
220            and(col1.clone(), col2.clone()).to_string(),
221            "($.col1 and $.col2)"
222        );
223        assert_eq!(
224            or(col1.clone(), col2.clone()).to_string(),
225            "($.col1 or $.col2)"
226        );
227        assert_eq!(
228            eq(col1.clone(), col2.clone()).to_string(),
229            "($.col1 = $.col2)"
230        );
231        assert_eq!(
232            not_eq(col1.clone(), col2.clone()).to_string(),
233            "($.col1 != $.col2)"
234        );
235        assert_eq!(
236            gt(col1.clone(), col2.clone()).to_string(),
237            "($.col1 > $.col2)"
238        );
239        assert_eq!(
240            gt_eq(col1.clone(), col2.clone()).to_string(),
241            "($.col1 >= $.col2)"
242        );
243        assert_eq!(
244            lt(col1.clone(), col2.clone()).to_string(),
245            "($.col1 < $.col2)"
246        );
247        assert_eq!(
248            lt_eq(col1.clone(), col2.clone()).to_string(),
249            "($.col1 <= $.col2)"
250        );
251
252        assert_eq!(
253            or(lt(col1.clone(), col2.clone()), not_eq(col1.clone(), col2),).to_string(),
254            "(($.col1 < $.col2) or ($.col1 != $.col2))"
255        );
256
257        assert_eq!(not(col1).to_string(), "vortex.not($.col1)");
258
259        assert_eq!(
260            select(vec![FieldName::from("col1")], root()).to_string(),
261            "${col1}"
262        );
263        assert_eq!(
264            select(
265                vec![FieldName::from("col1"), FieldName::from("col2")],
266                root()
267            )
268            .to_string(),
269            "${col1, col2}"
270        );
271        assert_eq!(
272            select_exclude(
273                vec![FieldName::from("col1"), FieldName::from("col2")],
274                root()
275            )
276            .to_string(),
277            "${~ col1, col2}"
278        );
279
280        assert_eq!(lit(Scalar::from(0u8)).to_string(), "0u8");
281        assert_eq!(lit(Scalar::from(0.0f32)).to_string(), "0f32");
282        assert_eq!(
283            lit(Scalar::from(i64::MAX)).to_string(),
284            "9223372036854775807i64"
285        );
286        assert_eq!(lit(Scalar::from(true)).to_string(), "true");
287        assert_eq!(
288            lit(Scalar::null(DType::Bool(Nullability::Nullable))).to_string(),
289            "null"
290        );
291
292        assert_eq!(
293            lit(Scalar::struct_(
294                DType::Struct(
295                    StructFields::new(
296                        FieldNames::from(["dog", "cat"]),
297                        vec![
298                            DType::Primitive(PType::U32, Nullability::NonNullable),
299                            DType::Utf8(Nullability::NonNullable)
300                        ],
301                    ),
302                    Nullability::NonNullable
303                ),
304                vec![Scalar::from(32_u32), Scalar::from("rufus".to_string())]
305            ))
306            .to_string(),
307            "{dog: 32u32, cat: \"rufus\"}"
308        );
309    }
310}
vortex_array/expr/mod.rs

vortex_array/expr/
mod.rs
