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_error::VortexResult;
13
14use crate::ArrayRef;
15use crate::IntoArray;
16use crate::arrays::ConstantArray;
17use crate::arrays::InterleaveArray;
18use crate::dtype::DType;
19use crate::dtype::FieldName;
20use crate::expr::Expression;
21use crate::optimizer::ArrayOptimizer;
22use crate::scalar::Scalar;
23use crate::scalar_fn::EmptyOptions;
24use crate::scalar_fn::ScalarFnVTableExt;
25use crate::scalar_fn::fns::between::Between;
26use crate::scalar_fn::fns::between::BetweenOptions;
27use crate::scalar_fn::fns::binary::Binary;
28use crate::scalar_fn::fns::cast::Cast;
29use crate::scalar_fn::fns::fill_null::FillNull;
30use crate::scalar_fn::fns::get_item::GetItem;
31use crate::scalar_fn::fns::is_not_null::IsNotNull;
32use crate::scalar_fn::fns::is_null::IsNull;
33use crate::scalar_fn::fns::list_contains::ListContains;
34use crate::scalar_fn::fns::mask::Mask;
35use crate::scalar_fn::fns::not::Not;
36use crate::scalar_fn::fns::operators::Operator;
37use crate::scalar_fn::fns::zip::Zip;
38
39/// A collection of built-in scalar functions that can be applied to expressions or arrays.
40pub trait ExprBuiltins: Sized {
41    /// Cast to the given data type.
42    fn cast(&self, dtype: DType) -> VortexResult<Expression>;
43
44    /// Replace null values with the given fill value.
45    fn fill_null(&self, fill_value: Expression) -> VortexResult<Expression>;
46
47    /// Get item by field name (for struct types).
48    fn get_item(&self, field_name: impl Into<FieldName>) -> VortexResult<Expression>;
49
50    /// Is null check.
51    fn is_null(&self) -> VortexResult<Expression>;
52
53    /// Is not null check.
54    fn is_not_null(&self) -> VortexResult<Expression>;
55
56    /// Mask the expression using the given boolean mask.
57    /// The resulting expression's validity is the intersection of the original expression's
58    /// validity.
59    fn mask(&self, mask: Expression) -> VortexResult<Expression>;
60
61    /// Boolean negation.
62    fn not(&self) -> VortexResult<Expression>;
63
64    /// Check if a list contains a value.
65    fn list_contains(&self, value: Expression) -> VortexResult<Expression>;
66
67    /// Conditional selection: `result[i] = if mask[i] then if_true[i] else if_false[i]`.
68    fn zip(&self, if_true: Expression, if_false: Expression) -> VortexResult<Expression>;
69
70    // TODO(joe): add an `interleave` expression builtin mirroring `ArrayBuiltins::interleave`.
71
72    /// Apply a binary operator to this expression and another.
73    fn binary(&self, rhs: Expression, op: Operator) -> VortexResult<Expression>;
74}
75
76impl ExprBuiltins for Expression {
77    fn cast(&self, dtype: DType) -> VortexResult<Expression> {
78        Cast.try_new_expr(dtype, [self.clone()])
79    }
80
81    fn fill_null(&self, fill_value: Expression) -> VortexResult<Expression> {
82        FillNull.try_new_expr(EmptyOptions, [self.clone(), fill_value])
83    }
84
85    fn get_item(&self, field_name: impl Into<FieldName>) -> VortexResult<Expression> {
86        GetItem.try_new_expr(field_name.into(), [self.clone()])
87    }
88
89    fn is_null(&self) -> VortexResult<Expression> {
90        IsNull.try_new_expr(EmptyOptions, [self.clone()])
91    }
92
93    fn is_not_null(&self) -> VortexResult<Expression> {
94        IsNotNull.try_new_expr(EmptyOptions, [self.clone()])
95    }
96
97    fn mask(&self, mask: Expression) -> VortexResult<Expression> {
98        Mask.try_new_expr(EmptyOptions, [self.clone(), mask])
99    }
100
101    fn not(&self) -> VortexResult<Expression> {
102        Not.try_new_expr(EmptyOptions, [self.clone()])
103    }
104
105    fn list_contains(&self, value: Expression) -> VortexResult<Expression> {
106        ListContains.try_new_expr(EmptyOptions, [self.clone(), value])
107    }
108
109    fn zip(&self, if_true: Expression, if_false: Expression) -> VortexResult<Expression> {
110        Zip.try_new_expr(EmptyOptions, [if_true, if_false, self.clone()])
111    }
112
113    fn binary(&self, rhs: Expression, op: Operator) -> VortexResult<Expression> {
114        Binary.try_new_expr(op, [self.clone(), rhs])
115    }
116}
117
118pub trait ArrayBuiltins: Sized {
119    /// Cast to the given data type.
120    fn cast(&self, dtype: DType) -> VortexResult<ArrayRef>;
121
122    /// Replace null values with the given fill value.
123    fn fill_null(&self, fill_value: impl Into<Scalar>) -> VortexResult<ArrayRef>;
124
125    /// Get item by field name (for struct types).
126    fn get_item(&self, field_name: impl Into<FieldName>) -> VortexResult<ArrayRef>;
127
128    /// Is null check.
129    fn is_null(&self) -> VortexResult<ArrayRef>;
130
131    /// Is not null check.
132    fn is_not_null(&self) -> VortexResult<ArrayRef>;
133
134    /// Mask the array using the given boolean mask.
135    /// The resulting array's validity is the intersection of the original array's validity
136    /// and the mask's validity.
137    fn mask(self, mask: ArrayRef) -> VortexResult<ArrayRef>;
138
139    /// Boolean negation.
140    fn not(&self) -> VortexResult<ArrayRef>;
141
142    /// Conditional selection: `result[i] = if mask[i] then if_true[i] else if_false[i]`.
143    fn zip(&self, if_true: ArrayRef, if_false: ArrayRef) -> VortexResult<ArrayRef>;
144
145    /// Random-access gather by `(array_index, row_index)`: output row `i` is taken from
146    /// `values[array_indices[i]][row_indices[i]]`, where `self` is the (non-nullable)
147    /// `array_indices` selector and `row_indices` names the position within the selected value.
148    /// See [`InterleaveArray`].
149    fn interleave(
150        &self,
151        values: impl IntoIterator<Item = ArrayRef>,
152        row_indices: ArrayRef,
153    ) -> VortexResult<ArrayRef>;
154
155    /// Check if a list contains a value.
156    fn list_contains(&self, value: ArrayRef) -> VortexResult<ArrayRef>;
157
158    /// Apply a binary operator to this array and another.
159    fn binary(&self, rhs: ArrayRef, op: Operator) -> VortexResult<ArrayRef>;
160
161    /// Compare a values between lower </<= value </<= upper
162    fn between(
163        self,
164        lower: ArrayRef,
165        upper: ArrayRef,
166        options: BetweenOptions,
167    ) -> VortexResult<ArrayRef>;
168}
169
170impl ArrayBuiltins for ArrayRef {
171    fn cast(&self, dtype: DType) -> VortexResult<ArrayRef> {
172        if self.dtype() == &dtype {
173            return Ok(self.clone());
174        }
175        Cast::new(self.clone(), dtype).into_array().optimize()
176    }
177
178    fn fill_null(&self, fill_value: impl Into<Scalar>) -> VortexResult<ArrayRef> {
179        let fill_value = fill_value.into();
180        if !self.dtype().is_nullable() {
181            return self.cast(fill_value.dtype().clone());
182        }
183        FillNull::try_new(
184            self.clone(),
185            ConstantArray::new(fill_value, self.len()).into_array(),
186        )?
187        .into_array()
188        .optimize()
189    }
190
191    fn get_item(&self, field_name: impl Into<FieldName>) -> VortexResult<ArrayRef> {
192        GetItem::try_new(self.clone(), field_name)?
193            .into_array()
194            .optimize()
195    }
196
197    fn is_null(&self) -> VortexResult<ArrayRef> {
198        IsNull::new(self.clone()).into_array().optimize()
199    }
200
201    fn is_not_null(&self) -> VortexResult<ArrayRef> {
202        IsNotNull::new(self.clone()).into_array().optimize()
203    }
204
205    fn mask(self, mask: ArrayRef) -> VortexResult<ArrayRef> {
206        Mask::try_new(self, mask)?.into_array().optimize()
207    }
208
209    fn not(&self) -> VortexResult<ArrayRef> {
210        Not::try_new(self.clone())?.into_array().optimize()
211    }
212
213    fn zip(&self, if_true: ArrayRef, if_false: ArrayRef) -> VortexResult<ArrayRef> {
214        Ok(Zip::try_new(if_true, if_false, self.clone())?.into_array())
215    }
216
217    fn interleave(
218        &self,
219        values: impl IntoIterator<Item = ArrayRef>,
220        row_indices: ArrayRef,
221    ) -> VortexResult<ArrayRef> {
222        Ok(
223            InterleaveArray::try_new(values.into_iter().collect(), self.clone(), row_indices)?
224                .into_array(),
225        )
226    }
227
228    fn list_contains(&self, value: ArrayRef) -> VortexResult<ArrayRef> {
229        ListContains::try_new(self.clone(), value)?
230            .into_array()
231            .optimize()
232    }
233
234    fn binary(&self, rhs: ArrayRef, op: Operator) -> VortexResult<ArrayRef> {
235        Binary::try_new(self.clone(), rhs, op)?
236            .into_array()
237            .optimize()
238    }
239
240    fn between(
241        self,
242        lower: ArrayRef,
243        upper: ArrayRef,
244        options: BetweenOptions,
245    ) -> VortexResult<ArrayRef> {
246        Between::try_new(self, lower, upper, options)?
247            .into_array()
248            .optimize()
249    }
250}
vortex_array/builtins.rs

vortex_array/
builtins.rs
