vortex_array/expr/
exprs.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Factory functions for creating [`Expression`]s from scalar function vtables.
5
6use std::sync::Arc;
7
8use vortex_error::VortexExpect;
9use vortex_utils::iter::ReduceBalancedIterExt;
10
11use crate::dtype::DType;
12use crate::dtype::FieldName;
13use crate::dtype::FieldNames;
14use crate::dtype::Nullability;
15use crate::expr::Expression;
16use crate::scalar::Scalar;
17use crate::scalar::ScalarValue;
18use crate::scalar_fn::EmptyOptions;
19use crate::scalar_fn::ScalarFnVTableExt;
20use crate::scalar_fn::fns::between::Between;
21use crate::scalar_fn::fns::between::BetweenOptions;
22use crate::scalar_fn::fns::binary::Binary;
23use crate::scalar_fn::fns::cast::Cast;
24use crate::scalar_fn::fns::dynamic::DynamicComparison;
25use crate::scalar_fn::fns::dynamic::DynamicComparisonExpr;
26use crate::scalar_fn::fns::dynamic::Rhs;
27use crate::scalar_fn::fns::fill_null::FillNull;
28use crate::scalar_fn::fns::get_item::GetItem;
29use crate::scalar_fn::fns::is_null::IsNull;
30use crate::scalar_fn::fns::like::Like;
31use crate::scalar_fn::fns::like::LikeOptions;
32use crate::scalar_fn::fns::list_contains::ListContains;
33use crate::scalar_fn::fns::literal::Literal;
34use crate::scalar_fn::fns::mask::Mask;
35use crate::scalar_fn::fns::merge::DuplicateHandling;
36use crate::scalar_fn::fns::merge::Merge;
37use crate::scalar_fn::fns::not::Not;
38use crate::scalar_fn::fns::operators::CompareOperator;
39use crate::scalar_fn::fns::operators::Operator;
40use crate::scalar_fn::fns::pack::Pack;
41use crate::scalar_fn::fns::pack::PackOptions;
42use crate::scalar_fn::fns::root::Root;
43use crate::scalar_fn::fns::select::FieldSelection;
44use crate::scalar_fn::fns::select::Select;
45use crate::scalar_fn::fns::zip::Zip;
46
47// ---- Root ----
48
49/// Creates an expression that references the root scope.
50///
51/// Returns the entire input array as passed to the expression evaluator.
52/// This is commonly used as the starting point for field access and other operations.
53pub fn root() -> Expression {
54    Root.try_new_expr(EmptyOptions, vec![])
55        .vortex_expect("Failed to create Root expression")
56}
57
58/// Return whether the expression is a root expression.
59pub fn is_root(expr: &Expression) -> bool {
60    expr.is::<Root>()
61}
62
63// ---- Literal ----
64
65/// Create a new `Literal` expression from a type that coerces to `Scalar`.
66///
67///
68/// ## Example usage
69///
70/// ```
71/// use vortex_array::arrays::PrimitiveArray;
72/// use vortex_array::dtype::Nullability;
73/// use vortex_array::expr::lit;
74/// use vortex_array::scalar_fn::fns::literal::Literal;
75/// use vortex_array::scalar::Scalar;
76///
77/// let number = lit(34i32);
78///
79/// let scalar = number.as_::<Literal>();
80/// assert_eq!(scalar, &Scalar::primitive(34i32, Nullability::NonNullable));
81/// ```
82pub fn lit(value: impl Into<Scalar>) -> Expression {
83    Literal.new_expr(value.into(), [])
84}
85
86// ---- GetItem / Col ----
87
88/// Creates an expression that accesses a field from the root array.
89///
90/// Equivalent to `get_item(field, root())` - extracts a named field from the input array.
91///
92/// ```rust
93/// # use vortex_array::expr::col;
94/// let expr = col("name");
95/// ```
96pub fn col(field: impl Into<FieldName>) -> Expression {
97    GetItem.new_expr(field.into(), vec![root()])
98}
99
100/// Creates an expression that extracts a named field from a struct expression.
101///
102/// Accesses the specified field from the result of the child expression.
103///
104/// ```rust
105/// # use vortex_array::expr::{get_item, root};
106/// let expr = get_item("user_id", root());
107/// ```
108pub fn get_item(field: impl Into<FieldName>, child: Expression) -> Expression {
109    GetItem.new_expr(field.into(), vec![child])
110}
111
112// ---- Binary operators ----
113
114/// Create a new [`Binary`] using the [`Eq`](Operator::Eq) operator.
115///
116/// ## Example usage
117///
118/// ```
119/// # use vortex_array::arrays::{BoolArray, PrimitiveArray};
120/// # use vortex_array::{Array, IntoArray, ToCanonical};
121/// # use vortex_array::validity::Validity;
122/// # use vortex_buffer::buffer;
123/// # use vortex_array::expr::{eq, root, lit};
124/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
125/// let result = xs.to_array().apply(&eq(root(), lit(3))).unwrap();
126///
127/// assert_eq!(
128///     result.to_bool().to_bit_buffer(),
129///     BoolArray::from_iter(vec![false, false, true]).to_bit_buffer(),
130/// );
131/// ```
132pub fn eq(lhs: Expression, rhs: Expression) -> Expression {
133    Binary
134        .try_new_expr(Operator::Eq, [lhs, rhs])
135        .vortex_expect("Failed to create Eq binary expression")
136}
137
138/// Create a new [`Binary`] using the [`NotEq`](Operator::NotEq) operator.
139///
140/// ## Example usage
141///
142/// ```
143/// # use vortex_array::arrays::{BoolArray, PrimitiveArray};
144/// # use vortex_array::{Array, IntoArray, ToCanonical};
145/// # use vortex_array::validity::Validity;
146/// # use vortex_buffer::buffer;
147/// # use vortex_array::expr::{root, lit, not_eq};
148/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
149/// let result = xs.to_array().apply(&not_eq(root(), lit(3))).unwrap();
150///
151/// assert_eq!(
152///     result.to_bool().to_bit_buffer(),
153///     BoolArray::from_iter(vec![true, true, false]).to_bit_buffer(),
154/// );
155/// ```
156pub fn not_eq(lhs: Expression, rhs: Expression) -> Expression {
157    Binary
158        .try_new_expr(Operator::NotEq, [lhs, rhs])
159        .vortex_expect("Failed to create NotEq binary expression")
160}
161
162/// Create a new [`Binary`] using the [`Gte`](Operator::Gte) operator.
163///
164/// ## Example usage
165///
166/// ```
167/// # use vortex_array::arrays::{BoolArray, PrimitiveArray };
168/// # use vortex_array::{Array, IntoArray, ToCanonical};
169/// # use vortex_array::validity::Validity;
170/// # use vortex_buffer::buffer;
171/// # use vortex_array::expr::{gt_eq, root, lit};
172/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
173/// let result = xs.to_array().apply(&gt_eq(root(), lit(3))).unwrap();
174///
175/// assert_eq!(
176///     result.to_bool().to_bit_buffer(),
177///     BoolArray::from_iter(vec![false, false, true]).to_bit_buffer(),
178/// );
179/// ```
180pub fn gt_eq(lhs: Expression, rhs: Expression) -> Expression {
181    Binary
182        .try_new_expr(Operator::Gte, [lhs, rhs])
183        .vortex_expect("Failed to create Gte binary expression")
184}
185
186/// Create a new [`Binary`] using the [`Gt`](Operator::Gt) operator.
187///
188/// ## Example usage
189///
190/// ```
191/// # use vortex_array::arrays::{BoolArray, PrimitiveArray };
192/// # use vortex_array::{Array, IntoArray, ToCanonical};
193/// # use vortex_array::validity::Validity;
194/// # use vortex_buffer::buffer;
195/// # use vortex_array::expr::{gt, root, lit};
196/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
197/// let result = xs.to_array().apply(&gt(root(), lit(2))).unwrap();
198///
199/// assert_eq!(
200///     result.to_bool().to_bit_buffer(),
201///     BoolArray::from_iter(vec![false, false, true]).to_bit_buffer(),
202/// );
203/// ```
204pub fn gt(lhs: Expression, rhs: Expression) -> Expression {
205    Binary
206        .try_new_expr(Operator::Gt, [lhs, rhs])
207        .vortex_expect("Failed to create Gt binary expression")
208}
209
210/// Create a new [`Binary`] using the [`Lte`](Operator::Lte) operator.
211///
212/// ## Example usage
213///
214/// ```
215/// # use vortex_array::arrays::{BoolArray, PrimitiveArray };
216/// # use vortex_array::{Array, IntoArray, ToCanonical};
217/// # use vortex_array::validity::Validity;
218/// # use vortex_buffer::buffer;
219/// # use vortex_array::expr::{root, lit, lt_eq};
220/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
221/// let result = xs.to_array().apply(&lt_eq(root(), lit(2))).unwrap();
222///
223/// assert_eq!(
224///     result.to_bool().to_bit_buffer(),
225///     BoolArray::from_iter(vec![true, true, false]).to_bit_buffer(),
226/// );
227/// ```
228pub fn lt_eq(lhs: Expression, rhs: Expression) -> Expression {
229    Binary
230        .try_new_expr(Operator::Lte, [lhs, rhs])
231        .vortex_expect("Failed to create Lte binary expression")
232}
233
234/// Create a new [`Binary`] using the [`Lt`](Operator::Lt) operator.
235///
236/// ## Example usage
237///
238/// ```
239/// # use vortex_array::arrays::{BoolArray, PrimitiveArray };
240/// # use vortex_array::{Array, IntoArray, ToCanonical};
241/// # use vortex_array::validity::Validity;
242/// # use vortex_buffer::buffer;
243/// # use vortex_array::expr::{root, lit, lt};
244/// let xs = PrimitiveArray::new(buffer![1i32, 2i32, 3i32], Validity::NonNullable);
245/// let result = xs.to_array().apply(&lt(root(), lit(3))).unwrap();
246///
247/// assert_eq!(
248///     result.to_bool().to_bit_buffer(),
249///     BoolArray::from_iter(vec![true, true, false]).to_bit_buffer(),
250/// );
251/// ```
252pub fn lt(lhs: Expression, rhs: Expression) -> Expression {
253    Binary
254        .try_new_expr(Operator::Lt, [lhs, rhs])
255        .vortex_expect("Failed to create Lt binary expression")
256}
257
258/// Create a new [`Binary`] using the [`Or`](Operator::Or) operator.
259///
260/// ## Example usage
261///
262/// ```
263/// # use vortex_array::arrays::BoolArray;
264/// # use vortex_array::{Array, IntoArray, ToCanonical};
265/// # use vortex_array::expr::{root, lit, or};
266/// let xs = BoolArray::from_iter(vec![true, false, true]);
267/// let result = xs.to_array().apply(&or(root(), lit(false))).unwrap();
268///
269/// assert_eq!(
270///     result.to_bool().to_bit_buffer(),
271///     BoolArray::from_iter(vec![true, false, true]).to_bit_buffer(),
272/// );
273/// ```
274pub fn or(lhs: Expression, rhs: Expression) -> Expression {
275    Binary
276        .try_new_expr(Operator::Or, [lhs, rhs])
277        .vortex_expect("Failed to create Or binary expression")
278}
279
280/// Collects a list of `or`ed values into a single expression using a balanced tree.
281///
282/// This creates a balanced binary tree to avoid deep nesting that could cause
283/// stack overflow during drop or evaluation.
284///
285/// [a, b, c, d] => or(or(a, b), or(c, d))
286pub fn or_collect<I>(iter: I) -> Option<Expression>
287where
288    I: IntoIterator<Item = Expression>,
289{
290    iter.into_iter().reduce_balanced(or)
291}
292
293/// Create a new [`Binary`] using the [`And`](Operator::And) operator.
294///
295/// ## Example usage
296///
297/// ```
298/// # use vortex_array::arrays::BoolArray;
299/// # use vortex_array::{Array, IntoArray, ToCanonical};
300/// # use vortex_array::expr::{and, root, lit};
301/// let xs = BoolArray::from_iter(vec![true, false, true]);
302/// let result = xs.to_array().apply(&and(root(), lit(true))).unwrap();
303///
304/// assert_eq!(
305///     result.to_bool().to_bit_buffer(),
306///     BoolArray::from_iter(vec![true, false, true]).to_bit_buffer(),
307/// );
308/// ```
309pub fn and(lhs: Expression, rhs: Expression) -> Expression {
310    Binary
311        .try_new_expr(Operator::And, [lhs, rhs])
312        .vortex_expect("Failed to create And binary expression")
313}
314
315/// Collects a list of `and`ed values into a single expression using a balanced tree.
316///
317/// This creates a balanced binary tree to avoid deep nesting that could cause
318/// stack overflow during drop or evaluation.
319///
320/// [a, b, c, d] => and(and(a, b), and(c, d))
321pub fn and_collect<I>(iter: I) -> Option<Expression>
322where
323    I: IntoIterator<Item = Expression>,
324{
325    iter.into_iter().reduce_balanced(and)
326}
327
328/// Create a new [`Binary`] using the [`Add`](Operator::Add) operator.
329///
330/// ## Example usage
331///
332/// ```
333/// # use vortex_array::{Array, IntoArray};
334/// # use vortex_array::arrow::IntoArrowArray as _;
335/// # use vortex_buffer::buffer;
336/// # use vortex_array::expr::{checked_add, lit, root};
337/// let xs = buffer![1, 2, 3].into_array();
338/// let result = xs.apply(&checked_add(root(), lit(5))).unwrap();
339///
340/// assert_eq!(
341///     &result.into_arrow_preferred().unwrap(),
342///     &buffer![6, 7, 8]
343///         .into_array()
344///         .into_arrow_preferred()
345///         .unwrap()
346/// );
347/// ```
348pub fn checked_add(lhs: Expression, rhs: Expression) -> Expression {
349    Binary
350        .try_new_expr(Operator::Add, [lhs, rhs])
351        .vortex_expect("Failed to create Add binary expression")
352}
353
354// ---- Not ----
355
356/// Creates an expression that logically inverts boolean values.
357///
358/// Returns the logical negation of the input boolean expression.
359///
360/// ```rust
361/// # use vortex_array::expr::{not, root};
362/// let expr = not(root());
363/// ```
364pub fn not(operand: Expression) -> Expression {
365    Not.new_expr(EmptyOptions, vec![operand])
366}
367
368// ---- Between ----
369
370/// Creates an expression that checks if values are between two bounds.
371///
372/// Returns a boolean array indicating which values fall within the specified range.
373/// The comparison strictness is controlled by the options parameter.
374///
375/// ```rust
376/// # use vortex_array::scalar_fn::fns::between::BetweenOptions;
377/// # use vortex_array::scalar_fn::fns::between::StrictComparison;
378/// # use vortex_array::expr::{between, lit, root};
379/// let opts = BetweenOptions {
380///     lower_strict: StrictComparison::NonStrict,
381///     upper_strict: StrictComparison::NonStrict,
382/// };
383/// let expr = between(root(), lit(10), lit(20), opts);
384/// ```
385pub fn between(
386    arr: Expression,
387    lower: Expression,
388    upper: Expression,
389    options: BetweenOptions,
390) -> Expression {
391    Between
392        .try_new_expr(options, [arr, lower, upper])
393        .vortex_expect("Failed to create Between expression")
394}
395
396// ---- Select ----
397
398/// Creates an expression that selects (includes) specific fields from an array.
399///
400/// Projects only the specified fields from the child expression, which must be of DType struct.
401/// ```rust
402/// # use vortex_array::expr::{select, root};
403/// let expr = select(["name", "age"], root());
404/// ```
405pub fn select(field_names: impl Into<FieldNames>, child: Expression) -> Expression {
406    Select
407        .try_new_expr(FieldSelection::Include(field_names.into()), [child])
408        .vortex_expect("Failed to create Select expression")
409}
410
411/// Creates an expression that excludes specific fields from an array.
412///
413/// Projects all fields except the specified ones from the input struct expression.
414///
415/// ```rust
416/// # use vortex_array::expr::{select_exclude, root};
417/// let expr = select_exclude(["internal_id", "metadata"], root());
418/// ```
419pub fn select_exclude(fields: impl Into<FieldNames>, child: Expression) -> Expression {
420    Select
421        .try_new_expr(FieldSelection::Exclude(fields.into()), [child])
422        .vortex_expect("Failed to create Select expression")
423}
424
425// ---- Pack ----
426
427/// Creates an expression that packs values into a struct with named fields.
428///
429/// ```rust
430/// # use vortex_array::dtype::Nullability;
431/// # use vortex_array::expr::{pack, col, lit};
432/// let expr = pack([("id", col("user_id")), ("constant", lit(42))], Nullability::NonNullable);
433/// ```
434pub fn pack(
435    elements: impl IntoIterator<Item = (impl Into<FieldName>, Expression)>,
436    nullability: Nullability,
437) -> Expression {
438    let (names, values): (Vec<_>, Vec<_>) = elements
439        .into_iter()
440        .map(|(name, value)| (name.into(), value))
441        .unzip();
442    Pack.new_expr(
443        PackOptions {
444            names: names.into(),
445            nullability,
446        },
447        values,
448    )
449}
450
451// ---- Cast ----
452
453/// Creates an expression that casts values to a target data type.
454///
455/// Converts the input expression's values to the specified target type.
456///
457/// ```rust
458/// # use vortex_array::dtype::{DType, Nullability, PType};
459/// # use vortex_array::expr::{cast, root};
460/// let expr = cast(root(), DType::Primitive(PType::I64, Nullability::NonNullable));
461/// ```
462pub fn cast(child: Expression, target: DType) -> Expression {
463    Cast.try_new_expr(target, [child])
464        .vortex_expect("Failed to create Cast expression")
465}
466
467// ---- FillNull ----
468
469/// Creates an expression that replaces null values with a fill value.
470///
471/// ```rust
472/// # use vortex_array::expr::{fill_null, root, lit};
473/// let expr = fill_null(root(), lit(0i32));
474/// ```
475pub fn fill_null(child: Expression, fill_value: Expression) -> Expression {
476    FillNull.new_expr(EmptyOptions, [child, fill_value])
477}
478
479// ---- IsNull ----
480
481/// Creates an expression that checks for null values.
482///
483/// Returns a boolean array indicating which positions contain null values.
484///
485/// ```rust
486/// # use vortex_array::expr::{is_null, root};
487/// let expr = is_null(root());
488/// ```
489pub fn is_null(child: Expression) -> Expression {
490    IsNull.new_expr(EmptyOptions, vec![child])
491}
492
493// ---- Like ----
494
495/// Creates a SQL LIKE expression.
496pub fn like(child: Expression, pattern: Expression) -> Expression {
497    Like.new_expr(
498        LikeOptions {
499            negated: false,
500            case_insensitive: false,
501        },
502        [child, pattern],
503    )
504}
505
506/// Creates a case-insensitive SQL ILIKE expression.
507pub fn ilike(child: Expression, pattern: Expression) -> Expression {
508    Like.new_expr(
509        LikeOptions {
510            negated: false,
511            case_insensitive: true,
512        },
513        [child, pattern],
514    )
515}
516
517/// Creates a negated SQL NOT LIKE expression.
518pub fn not_like(child: Expression, pattern: Expression) -> Expression {
519    Like.new_expr(
520        LikeOptions {
521            negated: true,
522            case_insensitive: false,
523        },
524        [child, pattern],
525    )
526}
527
528/// Creates a negated case-insensitive SQL NOT ILIKE expression.
529pub fn not_ilike(child: Expression, pattern: Expression) -> Expression {
530    Like.new_expr(
531        LikeOptions {
532            negated: true,
533            case_insensitive: true,
534        },
535        [child, pattern],
536    )
537}
538
539// ---- Mask ----
540
541/// Creates a mask expression that applies the given boolean mask to the input array.
542pub fn mask(array: Expression, mask: Expression) -> Expression {
543    Mask.new_expr(EmptyOptions, [array, mask])
544}
545
546// ---- Merge ----
547
548/// Creates an expression that merges struct expressions into a single struct.
549///
550/// Combines fields from all input expressions. If field names are duplicated,
551/// later expressions win. Fields are not recursively merged.
552///
553/// ```rust
554/// # use vortex_array::dtype::Nullability;
555/// # use vortex_array::expr::{merge, get_item, root};
556/// let expr = merge([get_item("a", root()), get_item("b", root())]);
557/// ```
558pub fn merge(elements: impl IntoIterator<Item = impl Into<Expression>>) -> Expression {
559    use itertools::Itertools as _;
560    let values = elements.into_iter().map(|value| value.into()).collect_vec();
561    Merge.new_expr(DuplicateHandling::default(), values)
562}
563
564/// Creates a merge expression with explicit duplicate handling.
565pub fn merge_opts(
566    elements: impl IntoIterator<Item = impl Into<Expression>>,
567    duplicate_handling: DuplicateHandling,
568) -> Expression {
569    use itertools::Itertools as _;
570    let values = elements.into_iter().map(|value| value.into()).collect_vec();
571    Merge.new_expr(duplicate_handling, values)
572}
573
574// ---- Zip ----
575
576/// Creates a zip expression that conditionally selects between two arrays.
577///
578/// ```rust
579/// # use vortex_array::expr::{zip_expr, root, lit};
580/// let expr = zip_expr(lit(true), root(), lit(0i32));
581/// ```
582pub fn zip_expr(mask: Expression, if_true: Expression, if_false: Expression) -> Expression {
583    Zip.new_expr(EmptyOptions, [if_true, if_false, mask])
584}
585
586// ---- Dynamic ----
587
588/// Creates a dynamic comparison expression.
589pub fn dynamic(
590    operator: CompareOperator,
591    rhs_value: impl Fn() -> Option<ScalarValue> + Send + Sync + 'static,
592    rhs_dtype: DType,
593    default: bool,
594    lhs: Expression,
595) -> Expression {
596    DynamicComparison.new_expr(
597        DynamicComparisonExpr {
598            operator,
599            rhs: Arc::new(Rhs {
600                value: Arc::new(rhs_value),
601                dtype: rhs_dtype,
602            }),
603            default,
604        },
605        [lhs],
606    )
607}
608
609// ---- ListContains ----
610
611/// Creates an expression that checks if a value is contained in a list.
612///
613/// Returns a boolean array indicating whether the value appears in each list.
614///
615/// ```rust
616/// # use vortex_array::expr::{list_contains, lit, root};
617/// let expr = list_contains(root(), lit(42));
618/// ```
619pub fn list_contains(list: Expression, value: Expression) -> Expression {
620    ListContains.new_expr(EmptyOptions, [list, value])
621}
vortex_array/expr/exprs.rs

vortex_array/expr/
exprs.rs
