polars_core/chunked_array/arithmetic/
numeric.rs

Help
1use polars_compute::arithmetic::ArithmeticKernel;
2
3use super::*;
4use crate::chunked_array::arity::{
5    apply_binary_kernel_broadcast, apply_binary_kernel_broadcast_owned, unary_kernel,
6    unary_kernel_owned,
7};
8
9macro_rules! impl_op_overload {
10    ($op: ident, $trait_method: ident, $ca_method: ident, $ca_method_scalar: ident) => {
11        impl<T: PolarsNumericType> $op for ChunkedArray<T> {
12            type Output = ChunkedArray<T>;
13
14            fn $trait_method(self, rhs: Self) -> Self::Output {
15                ArithmeticChunked::$ca_method(self, rhs)
16            }
17        }
18
19        impl<T: PolarsNumericType> $op for &ChunkedArray<T> {
20            type Output = ChunkedArray<T>;
21
22            fn $trait_method(self, rhs: Self) -> Self::Output {
23                ArithmeticChunked::$ca_method(self, rhs)
24            }
25        }
26
27        // TODO: make this more strict instead of casting.
28        impl<T: PolarsNumericType, N: Num + ToPrimitive> $op<N> for ChunkedArray<T> {
29            type Output = ChunkedArray<T>;
30
31            fn $trait_method(self, rhs: N) -> Self::Output {
32                let rhs: T::Native = NumCast::from(rhs).unwrap();
33                ArithmeticChunked::$ca_method_scalar(self, rhs)
34            }
35        }
36
37        impl<T: PolarsNumericType, N: Num + ToPrimitive> $op<N> for &ChunkedArray<T> {
38            type Output = ChunkedArray<T>;
39
40            fn $trait_method(self, rhs: N) -> Self::Output {
41                let rhs: T::Native = NumCast::from(rhs).unwrap();
42                ArithmeticChunked::$ca_method_scalar(self, rhs)
43            }
44        }
45    };
46}
47
48impl_op_overload!(Add, add, wrapping_add, wrapping_add_scalar);
49impl_op_overload!(Sub, sub, wrapping_sub, wrapping_sub_scalar);
50impl_op_overload!(Mul, mul, wrapping_mul, wrapping_mul_scalar);
51impl_op_overload!(Div, div, legacy_div, legacy_div_scalar); // FIXME: replace this with true division.
52impl_op_overload!(Rem, rem, wrapping_mod, wrapping_mod_scalar);
53
54pub trait ArithmeticChunked {
55    type Scalar;
56    type Out;
57    type TrueDivOut;
58
59    fn wrapping_abs(self) -> Self::Out;
60    fn wrapping_neg(self) -> Self::Out;
61    fn wrapping_add(self, rhs: Self) -> Self::Out;
62    fn wrapping_sub(self, rhs: Self) -> Self::Out;
63    fn wrapping_mul(self, rhs: Self) -> Self::Out;
64    fn wrapping_floor_div(self, rhs: Self) -> Self::Out;
65    fn wrapping_trunc_div(self, rhs: Self) -> Self::Out;
66    fn wrapping_mod(self, rhs: Self) -> Self::Out;
67
68    fn wrapping_add_scalar(self, rhs: Self::Scalar) -> Self::Out;
69    fn wrapping_sub_scalar(self, rhs: Self::Scalar) -> Self::Out;
70    fn wrapping_sub_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out;
71    fn wrapping_mul_scalar(self, rhs: Self::Scalar) -> Self::Out;
72    fn wrapping_floor_div_scalar(self, rhs: Self::Scalar) -> Self::Out;
73    fn wrapping_floor_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out;
74    fn wrapping_trunc_div_scalar(self, rhs: Self::Scalar) -> Self::Out;
75    fn wrapping_trunc_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out;
76    fn wrapping_mod_scalar(self, rhs: Self::Scalar) -> Self::Out;
77    fn wrapping_mod_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out;
78
79    fn true_div(self, rhs: Self) -> Self::TrueDivOut;
80    fn true_div_scalar(self, rhs: Self::Scalar) -> Self::TrueDivOut;
81    fn true_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::TrueDivOut;
82
83    // TODO: remove these.
84    // These are flooring division for integer types, true division for floating point types.
85    fn legacy_div(self, rhs: Self) -> Self::Out;
86    fn legacy_div_scalar(self, rhs: Self::Scalar) -> Self::Out;
87    fn legacy_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out;
88}
89
90impl<T: PolarsNumericType> ArithmeticChunked for ChunkedArray<T> {
91    type Scalar = T::Native;
92    type Out = ChunkedArray<T>;
93    type TrueDivOut = ChunkedArray<<T::Native as NumericNative>::TrueDivPolarsType>;
94
95    fn wrapping_abs(self) -> Self::Out {
96        unary_kernel_owned(self, ArithmeticKernel::wrapping_abs)
97    }
98
99    fn wrapping_neg(self) -> Self::Out {
100        unary_kernel_owned(self, ArithmeticKernel::wrapping_neg)
101    }
102
103    fn wrapping_add(self, rhs: Self) -> Self::Out {
104        apply_binary_kernel_broadcast_owned(
105            self,
106            rhs,
107            ArithmeticKernel::wrapping_add,
108            |l, r| ArithmeticKernel::wrapping_add_scalar(r, l),
109            ArithmeticKernel::wrapping_add_scalar,
110        )
111    }
112
113    fn wrapping_sub(self, rhs: Self) -> Self::Out {
114        apply_binary_kernel_broadcast_owned(
115            self,
116            rhs,
117            ArithmeticKernel::wrapping_sub,
118            ArithmeticKernel::wrapping_sub_scalar_lhs,
119            ArithmeticKernel::wrapping_sub_scalar,
120        )
121    }
122
123    fn wrapping_mul(self, rhs: Self) -> Self::Out {
124        apply_binary_kernel_broadcast_owned(
125            self,
126            rhs,
127            ArithmeticKernel::wrapping_mul,
128            |l, r| ArithmeticKernel::wrapping_mul_scalar(r, l),
129            ArithmeticKernel::wrapping_mul_scalar,
130        )
131    }
132
133    fn wrapping_floor_div(self, rhs: Self) -> Self::Out {
134        apply_binary_kernel_broadcast_owned(
135            self,
136            rhs,
137            ArithmeticKernel::wrapping_floor_div,
138            ArithmeticKernel::wrapping_floor_div_scalar_lhs,
139            ArithmeticKernel::wrapping_floor_div_scalar,
140        )
141    }
142
143    fn wrapping_trunc_div(self, rhs: Self) -> Self::Out {
144        apply_binary_kernel_broadcast_owned(
145            self,
146            rhs,
147            ArithmeticKernel::wrapping_trunc_div,
148            ArithmeticKernel::wrapping_trunc_div_scalar_lhs,
149            ArithmeticKernel::wrapping_trunc_div_scalar,
150        )
151    }
152
153    fn wrapping_mod(self, rhs: Self) -> Self::Out {
154        apply_binary_kernel_broadcast_owned(
155            self,
156            rhs,
157            ArithmeticKernel::wrapping_mod,
158            ArithmeticKernel::wrapping_mod_scalar_lhs,
159            ArithmeticKernel::wrapping_mod_scalar,
160        )
161    }
162
163    fn wrapping_add_scalar(self, rhs: Self::Scalar) -> Self::Out {
164        unary_kernel_owned(self, |a| ArithmeticKernel::wrapping_add_scalar(a, rhs))
165    }
166
167    fn wrapping_sub_scalar(self, rhs: Self::Scalar) -> Self::Out {
168        unary_kernel_owned(self, |a| ArithmeticKernel::wrapping_sub_scalar(a, rhs))
169    }
170
171    fn wrapping_sub_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out {
172        unary_kernel_owned(rhs, |a| ArithmeticKernel::wrapping_sub_scalar_lhs(lhs, a))
173    }
174
175    fn wrapping_mul_scalar(self, rhs: Self::Scalar) -> Self::Out {
176        unary_kernel_owned(self, |a| ArithmeticKernel::wrapping_mul_scalar(a, rhs))
177    }
178
179    fn wrapping_floor_div_scalar(self, rhs: Self::Scalar) -> Self::Out {
180        unary_kernel_owned(self, |a| {
181            ArithmeticKernel::wrapping_floor_div_scalar(a, rhs)
182        })
183    }
184
185    fn wrapping_floor_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out {
186        unary_kernel_owned(rhs, |a| {
187            ArithmeticKernel::wrapping_floor_div_scalar_lhs(lhs, a)
188        })
189    }
190
191    fn wrapping_trunc_div_scalar(self, rhs: Self::Scalar) -> Self::Out {
192        unary_kernel_owned(self, |a| {
193            ArithmeticKernel::wrapping_trunc_div_scalar(a, rhs)
194        })
195    }
196
197    fn wrapping_trunc_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out {
198        unary_kernel_owned(rhs, |a| {
199            ArithmeticKernel::wrapping_trunc_div_scalar_lhs(lhs, a)
200        })
201    }
202
203    fn wrapping_mod_scalar(self, rhs: Self::Scalar) -> Self::Out {
204        unary_kernel_owned(self, |a| ArithmeticKernel::wrapping_mod_scalar(a, rhs))
205    }
206
207    fn wrapping_mod_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out {
208        unary_kernel_owned(rhs, |a| ArithmeticKernel::wrapping_mod_scalar_lhs(lhs, a))
209    }
210
211    fn true_div(self, rhs: Self) -> Self::TrueDivOut {
212        apply_binary_kernel_broadcast_owned(
213            self,
214            rhs,
215            ArithmeticKernel::true_div,
216            ArithmeticKernel::true_div_scalar_lhs,
217            ArithmeticKernel::true_div_scalar,
218        )
219    }
220
221    fn true_div_scalar(self, rhs: Self::Scalar) -> Self::TrueDivOut {
222        unary_kernel_owned(self, |a| ArithmeticKernel::true_div_scalar(a, rhs))
223    }
224
225    fn true_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::TrueDivOut {
226        unary_kernel_owned(rhs, |a| ArithmeticKernel::true_div_scalar_lhs(lhs, a))
227    }
228
229    fn legacy_div(self, rhs: Self) -> Self::Out {
230        apply_binary_kernel_broadcast_owned(
231            self,
232            rhs,
233            ArithmeticKernel::legacy_div,
234            ArithmeticKernel::legacy_div_scalar_lhs,
235            ArithmeticKernel::legacy_div_scalar,
236        )
237    }
238
239    fn legacy_div_scalar(self, rhs: Self::Scalar) -> Self::Out {
240        unary_kernel_owned(self, |a| ArithmeticKernel::legacy_div_scalar(a, rhs))
241    }
242
243    fn legacy_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out {
244        unary_kernel_owned(rhs, |a| ArithmeticKernel::legacy_div_scalar_lhs(lhs, a))
245    }
246}
247
248impl<T: PolarsNumericType> ArithmeticChunked for &ChunkedArray<T> {
249    type Scalar = T::Native;
250    type Out = ChunkedArray<T>;
251    type TrueDivOut = ChunkedArray<<T::Native as NumericNative>::TrueDivPolarsType>;
252
253    fn wrapping_abs(self) -> Self::Out {
254        unary_kernel(self, |a| ArithmeticKernel::wrapping_abs(a.clone()))
255    }
256
257    fn wrapping_neg(self) -> Self::Out {
258        unary_kernel(self, |a| ArithmeticKernel::wrapping_neg(a.clone()))
259    }
260
261    fn wrapping_add(self, rhs: Self) -> Self::Out {
262        apply_binary_kernel_broadcast(
263            self,
264            rhs,
265            |l, r| ArithmeticKernel::wrapping_add(l.clone(), r.clone()),
266            |l, r| ArithmeticKernel::wrapping_add_scalar(r.clone(), l),
267            |l, r| ArithmeticKernel::wrapping_add_scalar(l.clone(), r),
268        )
269    }
270
271    fn wrapping_sub(self, rhs: Self) -> Self::Out {
272        apply_binary_kernel_broadcast(
273            self,
274            rhs,
275            |l, r| ArithmeticKernel::wrapping_sub(l.clone(), r.clone()),
276            |l, r| ArithmeticKernel::wrapping_sub_scalar_lhs(l, r.clone()),
277            |l, r| ArithmeticKernel::wrapping_sub_scalar(l.clone(), r),
278        )
279    }
280
281    fn wrapping_mul(self, rhs: Self) -> Self::Out {
282        apply_binary_kernel_broadcast(
283            self,
284            rhs,
285            |l, r| ArithmeticKernel::wrapping_mul(l.clone(), r.clone()),
286            |l, r| ArithmeticKernel::wrapping_mul_scalar(r.clone(), l),
287            |l, r| ArithmeticKernel::wrapping_mul_scalar(l.clone(), r),
288        )
289    }
290
291    fn wrapping_floor_div(self, rhs: Self) -> Self::Out {
292        apply_binary_kernel_broadcast(
293            self,
294            rhs,
295            |l, r| ArithmeticKernel::wrapping_floor_div(l.clone(), r.clone()),
296            |l, r| ArithmeticKernel::wrapping_floor_div_scalar_lhs(l, r.clone()),
297            |l, r| ArithmeticKernel::wrapping_floor_div_scalar(l.clone(), r),
298        )
299    }
300
301    fn wrapping_trunc_div(self, rhs: Self) -> Self::Out {
302        apply_binary_kernel_broadcast(
303            self,
304            rhs,
305            |l, r| ArithmeticKernel::wrapping_trunc_div(l.clone(), r.clone()),
306            |l, r| ArithmeticKernel::wrapping_trunc_div_scalar_lhs(l, r.clone()),
307            |l, r| ArithmeticKernel::wrapping_trunc_div_scalar(l.clone(), r),
308        )
309    }
310
311    fn wrapping_mod(self, rhs: Self) -> Self::Out {
312        apply_binary_kernel_broadcast(
313            self,
314            rhs,
315            |l, r| ArithmeticKernel::wrapping_mod(l.clone(), r.clone()),
316            |l, r| ArithmeticKernel::wrapping_mod_scalar_lhs(l, r.clone()),
317            |l, r| ArithmeticKernel::wrapping_mod_scalar(l.clone(), r),
318        )
319    }
320
321    fn wrapping_add_scalar(self, rhs: Self::Scalar) -> Self::Out {
322        unary_kernel(self, |a| {
323            ArithmeticKernel::wrapping_add_scalar(a.clone(), rhs)
324        })
325    }
326
327    fn wrapping_sub_scalar(self, rhs: Self::Scalar) -> Self::Out {
328        unary_kernel(self, |a| {
329            ArithmeticKernel::wrapping_sub_scalar(a.clone(), rhs)
330        })
331    }
332
333    fn wrapping_sub_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out {
334        unary_kernel(rhs, |a| {
335            ArithmeticKernel::wrapping_sub_scalar_lhs(lhs, a.clone())
336        })
337    }
338
339    fn wrapping_mul_scalar(self, rhs: Self::Scalar) -> Self::Out {
340        unary_kernel(self, |a| {
341            ArithmeticKernel::wrapping_mul_scalar(a.clone(), rhs)
342        })
343    }
344
345    fn wrapping_floor_div_scalar(self, rhs: Self::Scalar) -> Self::Out {
346        unary_kernel(self, |a| {
347            ArithmeticKernel::wrapping_floor_div_scalar(a.clone(), rhs)
348        })
349    }
350
351    fn wrapping_floor_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out {
352        unary_kernel(rhs, |a| {
353            ArithmeticKernel::wrapping_floor_div_scalar_lhs(lhs, a.clone())
354        })
355    }
356
357    fn wrapping_trunc_div_scalar(self, rhs: Self::Scalar) -> Self::Out {
358        unary_kernel(self, |a| {
359            ArithmeticKernel::wrapping_trunc_div_scalar(a.clone(), rhs)
360        })
361    }
362
363    fn wrapping_trunc_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out {
364        unary_kernel(rhs, |a| {
365            ArithmeticKernel::wrapping_trunc_div_scalar_lhs(lhs, a.clone())
366        })
367    }
368
369    fn wrapping_mod_scalar(self, rhs: Self::Scalar) -> Self::Out {
370        unary_kernel(self, |a| {
371            ArithmeticKernel::wrapping_mod_scalar(a.clone(), rhs)
372        })
373    }
374
375    fn wrapping_mod_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out {
376        unary_kernel(rhs, |a| {
377            ArithmeticKernel::wrapping_mod_scalar_lhs(lhs, a.clone())
378        })
379    }
380
381    fn true_div(self, rhs: Self) -> Self::TrueDivOut {
382        apply_binary_kernel_broadcast(
383            self,
384            rhs,
385            |l, r| ArithmeticKernel::true_div(l.clone(), r.clone()),
386            |l, r| ArithmeticKernel::true_div_scalar_lhs(l, r.clone()),
387            |l, r| ArithmeticKernel::true_div_scalar(l.clone(), r),
388        )
389    }
390
391    fn true_div_scalar(self, rhs: Self::Scalar) -> Self::TrueDivOut {
392        unary_kernel(self, |a| ArithmeticKernel::true_div_scalar(a.clone(), rhs))
393    }
394
395    fn true_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::TrueDivOut {
396        unary_kernel(rhs, |a| {
397            ArithmeticKernel::true_div_scalar_lhs(lhs, a.clone())
398        })
399    }
400
401    fn legacy_div(self, rhs: Self) -> Self::Out {
402        apply_binary_kernel_broadcast(
403            self,
404            rhs,
405            |l, r| ArithmeticKernel::legacy_div(l.clone(), r.clone()),
406            |l, r| ArithmeticKernel::legacy_div_scalar_lhs(l, r.clone()),
407            |l, r| ArithmeticKernel::legacy_div_scalar(l.clone(), r),
408        )
409    }
410
411    fn legacy_div_scalar(self, rhs: Self::Scalar) -> Self::Out {
412        unary_kernel(self, |a| {
413            ArithmeticKernel::legacy_div_scalar(a.clone(), rhs)
414        })
415    }
416
417    fn legacy_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out {
418        unary_kernel(rhs, |a| {
419            ArithmeticKernel::legacy_div_scalar_lhs(lhs, a.clone())
420        })
421    }
422}
polars_core/chunked_array/arithmetic/numeric.rs

polars_core/chunked_array/arithmetic/
numeric.rs
