arrow2/compute/arithmetics/basic/
add.rs

1//! Definition of basic add operations with primitive arrays
2use std::ops::Add;
3
4use num_traits::{ops::overflowing::OverflowingAdd, CheckedAdd, SaturatingAdd, WrappingAdd};
5
6use crate::{
7    array::PrimitiveArray,
8    bitmap::Bitmap,
9    compute::{
10        arithmetics::{
11            ArrayAdd, ArrayCheckedAdd, ArrayOverflowingAdd, ArraySaturatingAdd, ArrayWrappingAdd,
12        },
13        arity::{
14            binary, binary_checked, binary_with_bitmap, unary, unary_checked, unary_with_bitmap,
15        },
16    },
17};
18
19use super::NativeArithmetics;
20
21/// Adds two primitive arrays with the same type.
22/// Panics if the sum of one pair of values overflows.
23///
24/// # Examples
25/// ```
26/// use arrow2::compute::arithmetics::basic::add;
27/// use arrow2::array::PrimitiveArray;
28///
29/// let a = PrimitiveArray::from([None, Some(6), None, Some(6)]);
30/// let b = PrimitiveArray::from([Some(5), None, None, Some(6)]);
31/// let result = add(&a, &b);
32/// let expected = PrimitiveArray::from([None, None, None, Some(12)]);
33/// assert_eq!(result, expected)
34/// ```
35pub fn add<T>(lhs: &PrimitiveArray<T>, rhs: &PrimitiveArray<T>) -> PrimitiveArray<T>
36where
37    T: NativeArithmetics + Add<Output = T>,
38{
39    binary(lhs, rhs, lhs.data_type().clone(), |a, b| a + b)
40}
41
42/// Wrapping addition of two [`PrimitiveArray`]s.
43/// It wraps around at the boundary of the type if the result overflows.
44///
45/// # Examples
46/// ```
47/// use arrow2::compute::arithmetics::basic::wrapping_add;
48/// use arrow2::array::PrimitiveArray;
49///
50/// let a = PrimitiveArray::from([Some(-100i8), Some(100i8), Some(100i8)]);
51/// let b = PrimitiveArray::from([Some(0i8), Some(100i8), Some(0i8)]);
52/// let result = wrapping_add(&a, &b);
53/// let expected = PrimitiveArray::from([Some(-100i8), Some(-56i8), Some(100i8)]);
54/// assert_eq!(result, expected);
55/// ```
56pub fn wrapping_add<T>(lhs: &PrimitiveArray<T>, rhs: &PrimitiveArray<T>) -> PrimitiveArray<T>
57where
58    T: NativeArithmetics + WrappingAdd<Output = T>,
59{
60    let op = move |a: T, b: T| a.wrapping_add(&b);
61
62    binary(lhs, rhs, lhs.data_type().clone(), op)
63}
64
65/// Checked addition of two primitive arrays. If the result from the sum
66/// overflows, the validity for that index is changed to None
67///
68/// # Examples
69/// ```
70/// use arrow2::compute::arithmetics::basic::checked_add;
71/// use arrow2::array::PrimitiveArray;
72///
73/// let a = PrimitiveArray::from([Some(100i8), Some(100i8), Some(100i8)]);
74/// let b = PrimitiveArray::from([Some(0i8), Some(100i8), Some(0i8)]);
75/// let result = checked_add(&a, &b);
76/// let expected = PrimitiveArray::from([Some(100i8), None, Some(100i8)]);
77/// assert_eq!(result, expected);
78/// ```
79pub fn checked_add<T>(lhs: &PrimitiveArray<T>, rhs: &PrimitiveArray<T>) -> PrimitiveArray<T>
80where
81    T: NativeArithmetics + CheckedAdd<Output = T>,
82{
83    let op = move |a: T, b: T| a.checked_add(&b);
84
85    binary_checked(lhs, rhs, lhs.data_type().clone(), op)
86}
87
88/// Saturating addition of two primitive arrays. If the result from the sum is
89/// larger than the possible number for this type, the result for the operation
90/// will be the saturated value.
91///
92/// # Examples
93/// ```
94/// use arrow2::compute::arithmetics::basic::saturating_add;
95/// use arrow2::array::PrimitiveArray;
96///
97/// let a = PrimitiveArray::from([Some(100i8)]);
98/// let b = PrimitiveArray::from([Some(100i8)]);
99/// let result = saturating_add(&a, &b);
100/// let expected = PrimitiveArray::from([Some(127)]);
101/// assert_eq!(result, expected);
102/// ```
103pub fn saturating_add<T>(lhs: &PrimitiveArray<T>, rhs: &PrimitiveArray<T>) -> PrimitiveArray<T>
104where
105    T: NativeArithmetics + SaturatingAdd<Output = T>,
106{
107    let op = move |a: T, b: T| a.saturating_add(&b);
108
109    binary(lhs, rhs, lhs.data_type().clone(), op)
110}
111
112/// Overflowing addition of two primitive arrays. If the result from the sum is
113/// larger than the possible number for this type, the result for the operation
114/// will be an array with overflowed values and a  validity array indicating
115/// the overflowing elements from the array.
116///
117/// # Examples
118/// ```
119/// use arrow2::compute::arithmetics::basic::overflowing_add;
120/// use arrow2::array::PrimitiveArray;
121///
122/// let a = PrimitiveArray::from([Some(1i8), Some(100i8)]);
123/// let b = PrimitiveArray::from([Some(1i8), Some(100i8)]);
124/// let (result, overflow) = overflowing_add(&a, &b);
125/// let expected = PrimitiveArray::from([Some(2i8), Some(-56i8)]);
126/// assert_eq!(result, expected);
127/// ```
128pub fn overflowing_add<T>(
129    lhs: &PrimitiveArray<T>,
130    rhs: &PrimitiveArray<T>,
131) -> (PrimitiveArray<T>, Bitmap)
132where
133    T: NativeArithmetics + OverflowingAdd<Output = T>,
134{
135    let op = move |a: T, b: T| a.overflowing_add(&b);
136
137    binary_with_bitmap(lhs, rhs, lhs.data_type().clone(), op)
138}
139
140// Implementation of ArrayAdd trait for PrimitiveArrays
141impl<T> ArrayAdd<PrimitiveArray<T>> for PrimitiveArray<T>
142where
143    T: NativeArithmetics + Add<Output = T>,
144{
145    fn add(&self, rhs: &PrimitiveArray<T>) -> Self {
146        add(self, rhs)
147    }
148}
149
150impl<T> ArrayWrappingAdd<PrimitiveArray<T>> for PrimitiveArray<T>
151where
152    T: NativeArithmetics + WrappingAdd<Output = T>,
153{
154    fn wrapping_add(&self, rhs: &PrimitiveArray<T>) -> Self {
155        wrapping_add(self, rhs)
156    }
157}
158
159// Implementation of ArrayCheckedAdd trait for PrimitiveArrays
160impl<T> ArrayCheckedAdd<PrimitiveArray<T>> for PrimitiveArray<T>
161where
162    T: NativeArithmetics + CheckedAdd<Output = T>,
163{
164    fn checked_add(&self, rhs: &PrimitiveArray<T>) -> Self {
165        checked_add(self, rhs)
166    }
167}
168
169// Implementation of ArraySaturatingAdd trait for PrimitiveArrays
170impl<T> ArraySaturatingAdd<PrimitiveArray<T>> for PrimitiveArray<T>
171where
172    T: NativeArithmetics + SaturatingAdd<Output = T>,
173{
174    fn saturating_add(&self, rhs: &PrimitiveArray<T>) -> Self {
175        saturating_add(self, rhs)
176    }
177}
178
179// Implementation of ArraySaturatingAdd trait for PrimitiveArrays
180impl<T> ArrayOverflowingAdd<PrimitiveArray<T>> for PrimitiveArray<T>
181where
182    T: NativeArithmetics + OverflowingAdd<Output = T>,
183{
184    fn overflowing_add(&self, rhs: &PrimitiveArray<T>) -> (Self, Bitmap) {
185        overflowing_add(self, rhs)
186    }
187}
188
189/// Adds a scalar T to a primitive array of type T.
190/// Panics if the sum of the values overflows.
191///
192/// # Examples
193/// ```
194/// use arrow2::compute::arithmetics::basic::add_scalar;
195/// use arrow2::array::PrimitiveArray;
196///
197/// let a = PrimitiveArray::from([None, Some(6), None, Some(6)]);
198/// let result = add_scalar(&a, &1i32);
199/// let expected = PrimitiveArray::from([None, Some(7), None, Some(7)]);
200/// assert_eq!(result, expected)
201/// ```
202pub fn add_scalar<T>(lhs: &PrimitiveArray<T>, rhs: &T) -> PrimitiveArray<T>
203where
204    T: NativeArithmetics + Add<Output = T>,
205{
206    let rhs = *rhs;
207    unary(lhs, |a| a + rhs, lhs.data_type().clone())
208}
209
210/// Wrapping addition of a scalar T to a [`PrimitiveArray`] of type T.
211/// It do nothing if the result overflows.
212///
213/// # Examples
214/// ```
215/// use arrow2::compute::arithmetics::basic::wrapping_add_scalar;
216/// use arrow2::array::Int8Array;
217///
218/// let a = Int8Array::from(&[None, Some(100)]);
219/// let result = wrapping_add_scalar(&a, &100i8);
220/// let expected = Int8Array::from(&[None, Some(-56)]);
221/// assert_eq!(result, expected);
222/// ```
223pub fn wrapping_add_scalar<T>(lhs: &PrimitiveArray<T>, rhs: &T) -> PrimitiveArray<T>
224where
225    T: NativeArithmetics + WrappingAdd<Output = T>,
226{
227    unary(lhs, |a| a.wrapping_add(rhs), lhs.data_type().clone())
228}
229
230/// Checked addition of a scalar T to a primitive array of type T. If the
231/// result from the sum overflows then the validity index for that value is
232/// changed to None
233///
234/// # Examples
235/// ```
236/// use arrow2::compute::arithmetics::basic::checked_add_scalar;
237/// use arrow2::array::Int8Array;
238///
239/// let a = Int8Array::from(&[None, Some(100), None, Some(100)]);
240/// let result = checked_add_scalar(&a, &100i8);
241/// let expected = Int8Array::from(&[None, None, None, None]);
242/// assert_eq!(result, expected);
243/// ```
244pub fn checked_add_scalar<T>(lhs: &PrimitiveArray<T>, rhs: &T) -> PrimitiveArray<T>
245where
246    T: NativeArithmetics + CheckedAdd<Output = T>,
247{
248    let rhs = *rhs;
249    let op = move |a: T| a.checked_add(&rhs);
250
251    unary_checked(lhs, op, lhs.data_type().clone())
252}
253
254/// Saturated addition of a scalar T to a primitive array of type T. If the
255/// result from the sum is larger than the possible number for this type, then
256/// the result will be saturated
257///
258/// # Examples
259/// ```
260/// use arrow2::compute::arithmetics::basic::saturating_add_scalar;
261/// use arrow2::array::PrimitiveArray;
262///
263/// let a = PrimitiveArray::from([Some(100i8)]);
264/// let result = saturating_add_scalar(&a, &100i8);
265/// let expected = PrimitiveArray::from([Some(127)]);
266/// assert_eq!(result, expected);
267/// ```
268pub fn saturating_add_scalar<T>(lhs: &PrimitiveArray<T>, rhs: &T) -> PrimitiveArray<T>
269where
270    T: NativeArithmetics + SaturatingAdd<Output = T>,
271{
272    let rhs = *rhs;
273    let op = move |a: T| a.saturating_add(&rhs);
274
275    unary(lhs, op, lhs.data_type().clone())
276}
277
278/// Overflowing addition of a scalar T to a primitive array of type T. If the
279/// result from the sum is larger than the possible number for this type, then
280/// the result will be an array with overflowed values and a validity array
281/// indicating the overflowing elements from the array
282///
283/// # Examples
284/// ```
285/// use arrow2::compute::arithmetics::basic::overflowing_add_scalar;
286/// use arrow2::array::PrimitiveArray;
287///
288/// let a = PrimitiveArray::from([Some(1i8), Some(100i8)]);
289/// let (result, overflow) = overflowing_add_scalar(&a, &100i8);
290/// let expected = PrimitiveArray::from([Some(101i8), Some(-56i8)]);
291/// assert_eq!(result, expected);
292/// ```
293pub fn overflowing_add_scalar<T>(lhs: &PrimitiveArray<T>, rhs: &T) -> (PrimitiveArray<T>, Bitmap)
294where
295    T: NativeArithmetics + OverflowingAdd<Output = T>,
296{
297    let rhs = *rhs;
298    let op = move |a: T| a.overflowing_add(&rhs);
299
300    unary_with_bitmap(lhs, op, lhs.data_type().clone())
301}
302
303// Implementation of ArrayAdd trait for PrimitiveArrays with a scalar
304impl<T> ArrayAdd<T> for PrimitiveArray<T>
305where
306    T: NativeArithmetics + Add<Output = T>,
307{
308    fn add(&self, rhs: &T) -> Self {
309        add_scalar(self, rhs)
310    }
311}
312
313// Implementation of ArrayCheckedAdd trait for PrimitiveArrays with a scalar
314impl<T> ArrayCheckedAdd<T> for PrimitiveArray<T>
315where
316    T: NativeArithmetics + CheckedAdd<Output = T>,
317{
318    fn checked_add(&self, rhs: &T) -> Self {
319        checked_add_scalar(self, rhs)
320    }
321}
322
323// Implementation of ArraySaturatingAdd trait for PrimitiveArrays with a scalar
324impl<T> ArraySaturatingAdd<T> for PrimitiveArray<T>
325where
326    T: NativeArithmetics + SaturatingAdd<Output = T>,
327{
328    fn saturating_add(&self, rhs: &T) -> Self {
329        saturating_add_scalar(self, rhs)
330    }
331}
332
333// Implementation of ArraySaturatingAdd trait for PrimitiveArrays with a scalar
334impl<T> ArrayOverflowingAdd<T> for PrimitiveArray<T>
335where
336    T: NativeArithmetics + OverflowingAdd<Output = T>,
337{
338    fn overflowing_add(&self, rhs: &T) -> (Self, Bitmap) {
339        overflowing_add_scalar(self, rhs)
340    }
341}
arrow2/compute/arithmetics/basic/add.rs

arrow2/compute/arithmetics/basic/
add.rs
