runlength_arrow/
kernels.rs

1//! Element-wise kernels in the shape `zip_with` expects.
2//!
3//! Each takes two arrays of equal length and returns one value per position.
4//! They wrap `arrow-arith` and `arrow-ord`, so integer overflow and integer
5//! division by zero are errors.
6
7use std::sync::Arc;
8
9use arrow_arith::{boolean, numeric};
10use arrow_array::cast::AsArray;
11use arrow_array::{Array, ArrayRef, BooleanArray};
12use arrow_ord::cmp;
13use arrow_schema::ArrowError;
14
15/// A binary element-wise kernel.
16pub type Kernel = fn(&dyn Array, &dyn Array) -> Result<ArrayRef, ArrowError>;
17
18pub fn add(a: &dyn Array, b: &dyn Array) -> Result<ArrayRef, ArrowError> {
19    numeric::add(&a, &b)
20}
21
22pub fn sub(a: &dyn Array, b: &dyn Array) -> Result<ArrayRef, ArrowError> {
23    numeric::sub(&a, &b)
24}
25
26pub fn mul(a: &dyn Array, b: &dyn Array) -> Result<ArrayRef, ArrowError> {
27    numeric::mul(&a, &b)
28}
29
30pub fn div(a: &dyn Array, b: &dyn Array) -> Result<ArrayRef, ArrowError> {
31    numeric::div(&a, &b)
32}
33
34pub fn eq(a: &dyn Array, b: &dyn Array) -> Result<ArrayRef, ArrowError> {
35    Ok(Arc::new(cmp::eq(&a, &b)?))
36}
37
38pub fn ne(a: &dyn Array, b: &dyn Array) -> Result<ArrayRef, ArrowError> {
39    Ok(Arc::new(cmp::neq(&a, &b)?))
40}
41
42pub fn lt(a: &dyn Array, b: &dyn Array) -> Result<ArrayRef, ArrowError> {
43    Ok(Arc::new(cmp::lt(&a, &b)?))
44}
45
46pub fn le(a: &dyn Array, b: &dyn Array) -> Result<ArrayRef, ArrowError> {
47    Ok(Arc::new(cmp::lt_eq(&a, &b)?))
48}
49
50pub fn gt(a: &dyn Array, b: &dyn Array) -> Result<ArrayRef, ArrowError> {
51    Ok(Arc::new(cmp::gt(&a, &b)?))
52}
53
54pub fn ge(a: &dyn Array, b: &dyn Array) -> Result<ArrayRef, ArrowError> {
55    Ok(Arc::new(cmp::gt_eq(&a, &b)?))
56}
57
58/// Logical and, with SQL semantics for nulls.
59pub fn and(a: &dyn Array, b: &dyn Array) -> Result<ArrayRef, ArrowError> {
60    logical(a, b, boolean::and_kleene)
61}
62
63/// Logical or, with SQL semantics for nulls.
64pub fn or(a: &dyn Array, b: &dyn Array) -> Result<ArrayRef, ArrowError> {
65    logical(a, b, boolean::or_kleene)
66}
67
68fn logical(
69    a: &dyn Array,
70    b: &dyn Array,
71    f: fn(&BooleanArray, &BooleanArray) -> Result<BooleanArray, ArrowError>,
72) -> Result<ArrayRef, ArrowError> {
73    match (a.as_boolean_opt(), b.as_boolean_opt()) {
74        (Some(a), Some(b)) => Ok(Arc::new(f(a, b)?)),
75        _ => Err(ArrowError::InvalidArgumentError(format!(
76            "logical operations need boolean values, got {} and {}",
77            a.data_type(),
78            b.data_type()
79        ))),
80    }
81}
82
83/// Defines one method per kernel on a type with an
84/// `op(&self, &Self, Kernel) -> Result<Self>` method.
85macro_rules! kernel_methods {
86    ($Type:ident, $note:literal) => {
87        #[allow(clippy::should_implement_trait)]
88        impl $Type {
89            #[doc = concat!("Element-wise addition (integer overflow is an error).", $note)]
90            pub fn add(&self, other: &$Type) -> Result<$Type> {
91                self.op(other, $crate::kernels::add)
92            }
93            #[doc = concat!("Element-wise subtraction (integer overflow is an error).", $note)]
94            pub fn sub(&self, other: &$Type) -> Result<$Type> {
95                self.op(other, $crate::kernels::sub)
96            }
97            #[doc = concat!("Element-wise multiplication (integer overflow is an error).", $note)]
98            pub fn mul(&self, other: &$Type) -> Result<$Type> {
99                self.op(other, $crate::kernels::mul)
100            }
101            #[doc = concat!("Element-wise division (integer division by zero is an error).", $note)]
102            pub fn div(&self, other: &$Type) -> Result<$Type> {
103                self.op(other, $crate::kernels::div)
104            }
105            #[doc = concat!("Element-wise `==`.", $note)]
106            pub fn elem_eq(&self, other: &$Type) -> Result<$Type> {
107                self.op(other, $crate::kernels::eq)
108            }
109            #[doc = concat!("Element-wise `!=`.", $note)]
110            pub fn elem_ne(&self, other: &$Type) -> Result<$Type> {
111                self.op(other, $crate::kernels::ne)
112            }
113            #[doc = concat!("Element-wise `<`.", $note)]
114            pub fn elem_lt(&self, other: &$Type) -> Result<$Type> {
115                self.op(other, $crate::kernels::lt)
116            }
117            #[doc = concat!("Element-wise `<=`.", $note)]
118            pub fn elem_le(&self, other: &$Type) -> Result<$Type> {
119                self.op(other, $crate::kernels::le)
120            }
121            #[doc = concat!("Element-wise `>`.", $note)]
122            pub fn elem_gt(&self, other: &$Type) -> Result<$Type> {
123                self.op(other, $crate::kernels::gt)
124            }
125            #[doc = concat!("Element-wise `>=`.", $note)]
126            pub fn elem_ge(&self, other: &$Type) -> Result<$Type> {
127                self.op(other, $crate::kernels::ge)
128            }
129            #[doc = concat!("Element-wise logical and, with SQL semantics for nulls.", $note)]
130            pub fn and(&self, other: &$Type) -> Result<$Type> {
131                self.op(other, $crate::kernels::and)
132            }
133            #[doc = concat!("Element-wise logical or, with SQL semantics for nulls.", $note)]
134            pub fn or(&self, other: &$Type) -> Result<$Type> {
135                self.op(other, $crate::kernels::or)
136            }
137        }
138    };
139}
140
141pub(crate) use kernel_methods;
runlength_arrow/kernels.rs

runlength_arrow/
kernels.rs
