1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors
use std::fmt;
use std::fmt::Debug;
use std::fmt::Display;
use std::fmt::Formatter;
use std::hash::Hash;
use std::ops::Deref;
use std::sync::Arc;
use itertools::Itertools;
use vortex_dtype::DType;
use vortex_error::VortexResult;
use vortex_error::vortex_ensure;
use crate::expr::Root;
use crate::expr::ScalarFn;
use crate::expr::StatsCatalog;
use crate::expr::display::DisplayTreeExpr;
use crate::expr::stats::Stat;
/// A node in a Vortex expression tree.
///
/// Expressions represent scalar computations that can be performed on data. Each
/// expression consists of an encoding (vtable), heap-allocated metadata, and child expressions.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct Expression {
/// The scalar fn for this node.
scalar_fn: ScalarFn,
/// Any children of this expression.
children: Arc<Vec<Expression>>,
}
impl Deref for Expression {
type Target = ScalarFn;
fn deref(&self) -> &Self::Target {
&self.scalar_fn
}
}
impl Expression {
/// Create a new expression node from a scalar_fn expression and its children.
pub fn try_new(
scalar_fn: ScalarFn,
children: impl IntoIterator<Item = Expression>,
) -> VortexResult<Self> {
let children = Vec::from_iter(children);
vortex_ensure!(
scalar_fn.signature().arity().matches(children.len()),
"Expression arity mismatch: expected {} children but got {}",
scalar_fn.signature().arity(),
children.len()
);
Ok(Self {
scalar_fn,
children: children.into(),
})
}
/// Returns the scalar fn vtable for this expression.
pub fn scalar_fn(&self) -> &ScalarFn {
&self.scalar_fn
}
/// Returns the children of this expression.
pub fn children(&self) -> &Arc<Vec<Expression>> {
&self.children
}
/// Returns the n'th child of this expression.
pub fn child(&self, n: usize) -> &Expression {
&self.children[n]
}
/// Replace the children of this expression with the provided new children.
pub fn with_children(
mut self,
children: impl IntoIterator<Item = Expression>,
) -> VortexResult<Self> {
let children = Vec::from_iter(children);
vortex_ensure!(
self.signature().arity().matches(children.len()),
"Expression arity mismatch: expected {} children but got {}",
self.signature().arity(),
children.len()
);
self.children = Arc::new(children);
Ok(self)
}
/// Computes the return dtype of this expression given the input dtype.
pub fn return_dtype(&self, scope: &DType) -> VortexResult<DType> {
if self.is::<Root>() {
return Ok(scope.clone());
}
let dtypes: Vec<_> = self
.children
.iter()
.map(|c| c.return_dtype(scope))
.try_collect()?;
self.scalar_fn.return_dtype(&dtypes)
}
/// Returns a new expression representing the validity mask output of this expression.
///
/// The returned expression evaluates to a non-nullable boolean array.
pub fn validity(&self) -> VortexResult<Expression> {
self.scalar_fn.validity(self)
}
/// An expression over zone-statistics which implies all records in the zone evaluate to false.
///
/// Given an expression, `e`, if `e.stat_falsification(..)` evaluates to true, it is guaranteed
/// that `e` evaluates to false on all records in the zone. However, the inverse is not
/// necessarily true: even if the falsification evaluates to false, `e` need not evaluate to
/// true on all records.
///
/// The [`StatsCatalog`] can be used to constrain or rename stats used in the final expr.
///
/// # Examples
///
/// - An expression over one variable: `x > 0` is false for all records in a zone if the maximum
/// value of the column `x` in that zone is less than or equal to zero: `max(x) <= 0`.
/// - An expression over two variables: `x > y` becomes `max(x) <= min(y)`.
/// - A conjunctive expression: `x > y AND z < x` becomes `max(x) <= min(y) OR min(z) >= max(x).
///
/// Some expressions, in theory, have falsifications but this function does not support them
/// such as `x < (y < z)` or `x LIKE "needle%"`.
pub fn stat_falsification(&self, catalog: &dyn StatsCatalog) -> Option<Expression> {
self.vtable().as_dyn().stat_falsification(self, catalog)
}
/// Returns an expression representing the zoned statistic for the given stat, if available.
///
/// The [`StatsCatalog`] returns expressions that can be evaluated using the zone map as a
/// scope. Expressions can implement this function to propagate such statistics through the
/// expression tree. For example, the `a + 10` expression could propagate `min: min(a) + 10`.
///
/// NOTE(gatesn): we currently cannot represent statistics over nested fields. Please file an
/// issue to discuss a solution to this.
pub fn stat_expression(&self, stat: Stat, catalog: &dyn StatsCatalog) -> Option<Expression> {
self.vtable().as_dyn().stat_expression(self, stat, catalog)
}
/// Returns an expression representing the zoned maximum statistic, if available.
pub fn stat_min(&self, catalog: &dyn StatsCatalog) -> Option<Expression> {
self.stat_expression(Stat::Min, catalog)
}
/// Returns an expression representing the zoned maximum statistic, if available.
pub fn stat_max(&self, catalog: &dyn StatsCatalog) -> Option<Expression> {
self.stat_expression(Stat::Max, catalog)
}
/// Format the expression as a compact string.
///
/// Since this is a recursive formatter, it is exposed on the public Expression type.
/// See fmt_data that is only implemented on the vtable trait.
pub fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result {
self.vtable().as_dyn().fmt_sql(self, f)
}
/// Display the expression as a formatted tree structure.
///
/// This provides a hierarchical view of the expression that shows the relationships
/// between parent and child expressions, making complex nested expressions easier
/// to understand and debug.
///
/// # Example
///
/// ```rust
/// # use vortex_array::expr::LikeOptions;
/// # use vortex_array::expr::VTableExt;
/// # use vortex_dtype::{DType, Nullability, PType};
/// # use vortex_array::expr::{and, cast, eq, get_item, gt, lit, not, root, select, Like};
/// // Build a complex nested expression
/// let complex_expr = select(
/// ["result"],
/// and(
/// not(eq(get_item("status", root()), lit("inactive"))),
/// and(
/// Like.new_expr(LikeOptions::default(), [get_item("name", root()), lit("%admin%")]),
/// gt(
/// cast(get_item("score", root()), DType::Primitive(PType::F64, Nullability::NonNullable)),
/// lit(75.0)
/// )
/// )
/// )
/// );
///
/// println!("{}", complex_expr.display_tree());
/// ```
///
/// This produces output like:
///
/// ```text
/// Select(include): {result}
/// └── Binary(and)
/// ├── lhs: Not
/// │ └── Binary(=)
/// │ ├── lhs: GetItem(status)
/// │ │ └── Root
/// │ └── rhs: Literal(value: "inactive", dtype: utf8)
/// └── rhs: Binary(and)
/// ├── lhs: Like
/// │ ├── child: GetItem(name)
/// │ │ └── Root
/// │ └── pattern: Literal(value: "%admin%", dtype: utf8)
/// └── rhs: Binary(>)
/// ├── lhs: Cast(target: f64)
/// │ └── GetItem(score)
/// │ └── Root
/// └── rhs: Literal(value: 75f64, dtype: f64)
/// ```
pub fn display_tree(&self) -> impl Display {
DisplayTreeExpr(self)
}
}
/// The default display implementation for expressions uses the 'SQL'-style format.
impl Display for Expression {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
self.fmt_sql(f)
}
}
/// Iterative drop for expression to avoid stack overflows.
impl Drop for Expression {
fn drop(&mut self) {
if let Some(children) = Arc::get_mut(&mut self.children) {
let mut children_to_drop = std::mem::take(children);
while let Some(mut child) = children_to_drop.pop() {
if let Some(expr_children) = Arc::get_mut(&mut child.children) {
children_to_drop.append(expr_children);
}
}
}
}
}