vortex_array/
builtins.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! A collection of built-in common scalar functions.
5//!
6//! It is expected that each Vortex integration may provide its own set of scalar functions with
7//! semantics that exactly match the underlying system (e.g. SQL engine, DataFrame library, etc).
8//!
9//! This set of functions should cover the basics, and in general leans towards the semantics of
10//! the equivalent Arrow compute function.
11
12use vortex_dtype::DType;
13use vortex_dtype::FieldName;
14use vortex_error::VortexResult;
15
16use crate::Array;
17use crate::ArrayRef;
18use crate::arrays::ScalarFnArrayExt;
19use crate::expr::Cast;
20use crate::expr::EmptyOptions;
21use crate::expr::Expression;
22use crate::expr::GetItem;
23use crate::expr::IsNull;
24use crate::expr::Mask;
25use crate::expr::Not;
26use crate::expr::VTableExt;
27use crate::optimizer::ArrayOptimizer;
28
29/// A collection of built-in scalar functions that can be applied to expressions or arrays.
30pub trait ExprBuiltins: Sized {
31    /// Cast to the given data type.
32    fn cast(&self, dtype: DType) -> VortexResult<Expression>;
33
34    /// Get item by field name (for struct types).
35    fn get_item(&self, field_name: impl Into<FieldName>) -> VortexResult<Expression>;
36
37    /// Is null check.
38    fn is_null(&self) -> VortexResult<Expression>;
39
40    /// Mask the expression using the given boolean mask.
41    /// The resulting expression's validity is the intersection of the original expression's
42    /// validity.
43    fn mask(&self, mask: Expression) -> VortexResult<Expression>;
44
45    /// Boolean negation.
46    fn not(&self) -> VortexResult<Expression>;
47}
48
49impl ExprBuiltins for Expression {
50    fn cast(&self, dtype: DType) -> VortexResult<Expression> {
51        Cast.try_new_expr(dtype, [self.clone()])
52    }
53
54    fn get_item(&self, field_name: impl Into<FieldName>) -> VortexResult<Expression> {
55        GetItem.try_new_expr(field_name.into(), [self.clone()])
56    }
57
58    fn is_null(&self) -> VortexResult<Expression> {
59        IsNull.try_new_expr(EmptyOptions, [self.clone()])
60    }
61
62    fn mask(&self, mask: Expression) -> VortexResult<Expression> {
63        Mask.try_new_expr(EmptyOptions, [self.clone(), mask])
64    }
65
66    fn not(&self) -> VortexResult<Expression> {
67        Not.try_new_expr(EmptyOptions, [self.clone()])
68    }
69}
70
71pub trait ArrayBuiltins: Sized {
72    /// Cast to the given data type.
73    fn cast(&self, dtype: DType) -> VortexResult<ArrayRef>;
74
75    /// Get item by field name (for struct types).
76    fn get_item(&self, field_name: impl Into<FieldName>) -> VortexResult<ArrayRef>;
77
78    /// Is null check.
79    fn is_null(&self) -> VortexResult<ArrayRef>;
80
81    /// Mask the array using the given boolean mask.
82    /// The resulting array's validity is the intersection of the original array's validity
83    /// and the mask's validity.
84    fn mask(&self, mask: &ArrayRef) -> VortexResult<ArrayRef>;
85
86    /// Boolean negation.
87    fn not(&self) -> VortexResult<ArrayRef>;
88}
89
90impl ArrayBuiltins for ArrayRef {
91    fn cast(&self, dtype: DType) -> VortexResult<ArrayRef> {
92        Cast.try_new_array(self.len(), dtype, [self.clone()])?
93            .optimize()
94    }
95
96    fn get_item(&self, field_name: impl Into<FieldName>) -> VortexResult<ArrayRef> {
97        GetItem
98            .try_new_array(self.len(), field_name.into(), [self.clone()])?
99            .optimize()
100    }
101
102    fn is_null(&self) -> VortexResult<ArrayRef> {
103        IsNull
104            .try_new_array(self.len(), EmptyOptions, [self.clone()])?
105            .optimize()
106    }
107
108    fn mask(&self, mask: &ArrayRef) -> VortexResult<ArrayRef> {
109        Mask.try_new_array(self.len(), EmptyOptions, [self.clone(), mask.clone()])?
110            .optimize()
111    }
112
113    fn not(&self) -> VortexResult<ArrayRef> {
114        Not.try_new_array(self.len(), EmptyOptions, [self.clone()])?
115            .optimize()
116    }
117}
vortex_array/builtins.rs

vortex_array/
builtins.rs
