vortex_array/scalar_fns/
mod.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::Expression;
20use crate::expr::ScalarFnExprExt;
21use crate::expr::functions::EmptyOptions;
22
23pub mod binary;
24pub mod cast;
25pub mod get_item;
26pub mod is_null;
27pub mod mask;
28pub mod not;
29
30/// A collection of built-in scalar functions that can be applied to expressions or arrays.
31pub trait ExprBuiltins: Sized {
32    /// Cast to the given data type.
33    fn cast(&self, dtype: DType) -> VortexResult<Expression>;
34
35    /// Get item by field name (for struct types).
36    fn get_item(&self, field_name: impl Into<FieldName>) -> VortexResult<Expression>;
37
38    /// Is null check.
39    fn is_null(&self) -> VortexResult<Expression>;
40
41    /// Mask the expression using the given boolean mask.
42    /// The resulting expression's validity is the intersection of the original expression's
43    /// validity.
44    fn mask(&self, mask: Expression) -> VortexResult<Expression>;
45
46    /// Boolean negation.
47    fn not(&self) -> VortexResult<Expression>;
48}
49
50impl ExprBuiltins for Expression {
51    fn cast(&self, dtype: DType) -> VortexResult<Expression> {
52        cast::CastFn.try_new_expr(dtype, [self.clone()])
53    }
54
55    fn get_item(&self, field_name: impl Into<FieldName>) -> VortexResult<Expression> {
56        get_item::GetItemFn.try_new_expr(field_name.into(), [self.clone()])
57    }
58
59    fn is_null(&self) -> VortexResult<Expression> {
60        is_null::IsNullFn.try_new_expr(EmptyOptions, [self.clone()])
61    }
62
63    fn mask(&self, mask: Expression) -> VortexResult<Expression> {
64        mask::MaskFn.try_new_expr(EmptyOptions, [self.clone(), mask])
65    }
66
67    fn not(&self) -> VortexResult<Expression> {
68        not::NotFn.try_new_expr(EmptyOptions, [self.clone()])
69    }
70}
71
72pub trait ArrayBuiltins: Sized {
73    /// Cast to the given data type.
74    fn cast(&self, dtype: DType) -> VortexResult<ArrayRef>;
75
76    /// Get item by field name (for struct types).
77    fn get_item(&self, field_name: impl Into<FieldName>) -> VortexResult<ArrayRef>;
78
79    /// Is null check.
80    fn is_null(&self) -> VortexResult<ArrayRef>;
81
82    /// Mask the array using the given boolean mask.
83    /// The resulting array's validity is the intersection of the original array's validity
84    /// and the mask's validity.
85    fn mask(&self, mask: &ArrayRef) -> VortexResult<ArrayRef>;
86
87    /// Boolean negation.
88    fn not(&self) -> VortexResult<ArrayRef>;
89}
90
91impl ArrayBuiltins for ArrayRef {
92    fn cast(&self, dtype: DType) -> VortexResult<ArrayRef> {
93        cast::CastFn.try_new_array(self.len(), dtype, [self.clone()])
94    }
95
96    fn get_item(&self, field_name: impl Into<FieldName>) -> VortexResult<ArrayRef> {
97        get_item::GetItemFn.try_new_array(self.len(), field_name.into(), [self.clone()])
98    }
99
100    fn is_null(&self) -> VortexResult<ArrayRef> {
101        is_null::IsNullFn.try_new_array(self.len(), EmptyOptions, [self.clone()])
102    }
103
104    fn mask(&self, mask: &ArrayRef) -> VortexResult<ArrayRef> {
105        mask::MaskFn.try_new_array(self.len(), EmptyOptions, [self.clone(), mask.clone()])
106    }
107
108    fn not(&self) -> VortexResult<ArrayRef> {
109        not::NotFn.try_new_array(self.len(), EmptyOptions, [self.clone()])
110    }
111}