polars_core/chunked_array/arithmetic/
numeric.rs

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 checked_mul_scalar(self, rhs: Self::Scalar) -> Self::Out;
80
81    fn true_div(self, rhs: Self) -> Self::TrueDivOut;
82    fn true_div_scalar(self, rhs: Self::Scalar) -> Self::TrueDivOut;
83    fn true_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::TrueDivOut;
84
85    // TODO: remove these.
86    // These are flooring division for integer types, true division for floating point types.
87    fn legacy_div(self, rhs: Self) -> Self::Out;
88    fn legacy_div_scalar(self, rhs: Self::Scalar) -> Self::Out;
89    fn legacy_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out;
90}
91
92impl<T: PolarsNumericType> ArithmeticChunked for ChunkedArray<T> {
93    type Scalar = T::Native;
94    type Out = ChunkedArray<T>;
95    type TrueDivOut = ChunkedArray<<T::Native as NumericNative>::TrueDivPolarsType>;
96
97    fn wrapping_abs(self) -> Self::Out {
98        unary_kernel_owned(self, ArithmeticKernel::wrapping_abs)
99    }
100
101    fn wrapping_neg(self) -> Self::Out {
102        unary_kernel_owned(self, ArithmeticKernel::wrapping_neg)
103    }
104
105    fn wrapping_add(self, rhs: Self) -> Self::Out {
106        apply_binary_kernel_broadcast_owned(
107            self,
108            rhs,
109            ArithmeticKernel::wrapping_add,
110            |l, r| ArithmeticKernel::wrapping_add_scalar(r, l),
111            ArithmeticKernel::wrapping_add_scalar,
112        )
113    }
114
115    fn wrapping_sub(self, rhs: Self) -> Self::Out {
116        apply_binary_kernel_broadcast_owned(
117            self,
118            rhs,
119            ArithmeticKernel::wrapping_sub,
120            ArithmeticKernel::wrapping_sub_scalar_lhs,
121            ArithmeticKernel::wrapping_sub_scalar,
122        )
123    }
124
125    fn wrapping_mul(self, rhs: Self) -> Self::Out {
126        apply_binary_kernel_broadcast_owned(
127            self,
128            rhs,
129            ArithmeticKernel::wrapping_mul,
130            |l, r| ArithmeticKernel::wrapping_mul_scalar(r, l),
131            ArithmeticKernel::wrapping_mul_scalar,
132        )
133    }
134
135    fn wrapping_floor_div(self, rhs: Self) -> Self::Out {
136        apply_binary_kernel_broadcast_owned(
137            self,
138            rhs,
139            ArithmeticKernel::wrapping_floor_div,
140            ArithmeticKernel::wrapping_floor_div_scalar_lhs,
141            ArithmeticKernel::wrapping_floor_div_scalar,
142        )
143    }
144
145    fn wrapping_trunc_div(self, rhs: Self) -> Self::Out {
146        apply_binary_kernel_broadcast_owned(
147            self,
148            rhs,
149            ArithmeticKernel::wrapping_trunc_div,
150            ArithmeticKernel::wrapping_trunc_div_scalar_lhs,
151            ArithmeticKernel::wrapping_trunc_div_scalar,
152        )
153    }
154
155    fn wrapping_mod(self, rhs: Self) -> Self::Out {
156        apply_binary_kernel_broadcast_owned(
157            self,
158            rhs,
159            ArithmeticKernel::wrapping_mod,
160            ArithmeticKernel::wrapping_mod_scalar_lhs,
161            ArithmeticKernel::wrapping_mod_scalar,
162        )
163    }
164
165    fn wrapping_add_scalar(self, rhs: Self::Scalar) -> Self::Out {
166        unary_kernel_owned(self, |a| ArithmeticKernel::wrapping_add_scalar(a, rhs))
167    }
168
169    fn wrapping_sub_scalar(self, rhs: Self::Scalar) -> Self::Out {
170        unary_kernel_owned(self, |a| ArithmeticKernel::wrapping_sub_scalar(a, rhs))
171    }
172
173    fn wrapping_sub_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out {
174        unary_kernel_owned(rhs, |a| ArithmeticKernel::wrapping_sub_scalar_lhs(lhs, a))
175    }
176
177    fn wrapping_mul_scalar(self, rhs: Self::Scalar) -> Self::Out {
178        unary_kernel_owned(self, |a| ArithmeticKernel::wrapping_mul_scalar(a, rhs))
179    }
180
181    fn wrapping_floor_div_scalar(self, rhs: Self::Scalar) -> Self::Out {
182        unary_kernel_owned(self, |a| {
183            ArithmeticKernel::wrapping_floor_div_scalar(a, rhs)
184        })
185    }
186
187    fn wrapping_floor_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out {
188        unary_kernel_owned(rhs, |a| {
189            ArithmeticKernel::wrapping_floor_div_scalar_lhs(lhs, a)
190        })
191    }
192
193    fn wrapping_trunc_div_scalar(self, rhs: Self::Scalar) -> Self::Out {
194        unary_kernel_owned(self, |a| {
195            ArithmeticKernel::wrapping_trunc_div_scalar(a, rhs)
196        })
197    }
198
199    fn wrapping_trunc_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out {
200        unary_kernel_owned(rhs, |a| {
201            ArithmeticKernel::wrapping_trunc_div_scalar_lhs(lhs, a)
202        })
203    }
204
205    fn wrapping_mod_scalar(self, rhs: Self::Scalar) -> Self::Out {
206        unary_kernel_owned(self, |a| ArithmeticKernel::wrapping_mod_scalar(a, rhs))
207    }
208
209    fn wrapping_mod_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out {
210        unary_kernel_owned(rhs, |a| ArithmeticKernel::wrapping_mod_scalar_lhs(lhs, a))
211    }
212
213    fn checked_mul_scalar(self, rhs: Self::Scalar) -> Self::Out {
214        unary_kernel_owned(self, |a| ArithmeticKernel::checked_mul_scalar(a, rhs))
215    }
216
217    fn true_div(self, rhs: Self) -> Self::TrueDivOut {
218        apply_binary_kernel_broadcast_owned(
219            self,
220            rhs,
221            ArithmeticKernel::true_div,
222            ArithmeticKernel::true_div_scalar_lhs,
223            ArithmeticKernel::true_div_scalar,
224        )
225    }
226
227    fn true_div_scalar(self, rhs: Self::Scalar) -> Self::TrueDivOut {
228        unary_kernel_owned(self, |a| ArithmeticKernel::true_div_scalar(a, rhs))
229    }
230
231    fn true_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::TrueDivOut {
232        unary_kernel_owned(rhs, |a| ArithmeticKernel::true_div_scalar_lhs(lhs, a))
233    }
234
235    fn legacy_div(self, rhs: Self) -> Self::Out {
236        apply_binary_kernel_broadcast_owned(
237            self,
238            rhs,
239            ArithmeticKernel::legacy_div,
240            ArithmeticKernel::legacy_div_scalar_lhs,
241            ArithmeticKernel::legacy_div_scalar,
242        )
243    }
244
245    fn legacy_div_scalar(self, rhs: Self::Scalar) -> Self::Out {
246        unary_kernel_owned(self, |a| ArithmeticKernel::legacy_div_scalar(a, rhs))
247    }
248
249    fn legacy_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out {
250        unary_kernel_owned(rhs, |a| ArithmeticKernel::legacy_div_scalar_lhs(lhs, a))
251    }
252}
253
254impl<T: PolarsNumericType> ArithmeticChunked for &ChunkedArray<T> {
255    type Scalar = T::Native;
256    type Out = ChunkedArray<T>;
257    type TrueDivOut = ChunkedArray<<T::Native as NumericNative>::TrueDivPolarsType>;
258
259    fn wrapping_abs(self) -> Self::Out {
260        unary_kernel(self, |a| ArithmeticKernel::wrapping_abs(a.clone()))
261    }
262
263    fn wrapping_neg(self) -> Self::Out {
264        unary_kernel(self, |a| ArithmeticKernel::wrapping_neg(a.clone()))
265    }
266
267    fn wrapping_add(self, rhs: Self) -> Self::Out {
268        apply_binary_kernel_broadcast(
269            self,
270            rhs,
271            |l, r| ArithmeticKernel::wrapping_add(l.clone(), r.clone()),
272            |l, r| ArithmeticKernel::wrapping_add_scalar(r.clone(), l),
273            |l, r| ArithmeticKernel::wrapping_add_scalar(l.clone(), r),
274        )
275    }
276
277    fn wrapping_sub(self, rhs: Self) -> Self::Out {
278        apply_binary_kernel_broadcast(
279            self,
280            rhs,
281            |l, r| ArithmeticKernel::wrapping_sub(l.clone(), r.clone()),
282            |l, r| ArithmeticKernel::wrapping_sub_scalar_lhs(l, r.clone()),
283            |l, r| ArithmeticKernel::wrapping_sub_scalar(l.clone(), r),
284        )
285    }
286
287    fn wrapping_mul(self, rhs: Self) -> Self::Out {
288        apply_binary_kernel_broadcast(
289            self,
290            rhs,
291            |l, r| ArithmeticKernel::wrapping_mul(l.clone(), r.clone()),
292            |l, r| ArithmeticKernel::wrapping_mul_scalar(r.clone(), l),
293            |l, r| ArithmeticKernel::wrapping_mul_scalar(l.clone(), r),
294        )
295    }
296
297    fn wrapping_floor_div(self, rhs: Self) -> Self::Out {
298        apply_binary_kernel_broadcast(
299            self,
300            rhs,
301            |l, r| ArithmeticKernel::wrapping_floor_div(l.clone(), r.clone()),
302            |l, r| ArithmeticKernel::wrapping_floor_div_scalar_lhs(l, r.clone()),
303            |l, r| ArithmeticKernel::wrapping_floor_div_scalar(l.clone(), r),
304        )
305    }
306
307    fn wrapping_trunc_div(self, rhs: Self) -> Self::Out {
308        apply_binary_kernel_broadcast(
309            self,
310            rhs,
311            |l, r| ArithmeticKernel::wrapping_trunc_div(l.clone(), r.clone()),
312            |l, r| ArithmeticKernel::wrapping_trunc_div_scalar_lhs(l, r.clone()),
313            |l, r| ArithmeticKernel::wrapping_trunc_div_scalar(l.clone(), r),
314        )
315    }
316
317    fn wrapping_mod(self, rhs: Self) -> Self::Out {
318        apply_binary_kernel_broadcast(
319            self,
320            rhs,
321            |l, r| ArithmeticKernel::wrapping_mod(l.clone(), r.clone()),
322            |l, r| ArithmeticKernel::wrapping_mod_scalar_lhs(l, r.clone()),
323            |l, r| ArithmeticKernel::wrapping_mod_scalar(l.clone(), r),
324        )
325    }
326
327    fn wrapping_add_scalar(self, rhs: Self::Scalar) -> Self::Out {
328        unary_kernel(self, |a| {
329            ArithmeticKernel::wrapping_add_scalar(a.clone(), rhs)
330        })
331    }
332
333    fn wrapping_sub_scalar(self, rhs: Self::Scalar) -> Self::Out {
334        unary_kernel(self, |a| {
335            ArithmeticKernel::wrapping_sub_scalar(a.clone(), rhs)
336        })
337    }
338
339    fn wrapping_sub_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out {
340        unary_kernel(rhs, |a| {
341            ArithmeticKernel::wrapping_sub_scalar_lhs(lhs, a.clone())
342        })
343    }
344
345    fn wrapping_mul_scalar(self, rhs: Self::Scalar) -> Self::Out {
346        unary_kernel(self, |a| {
347            ArithmeticKernel::wrapping_mul_scalar(a.clone(), rhs)
348        })
349    }
350
351    fn wrapping_floor_div_scalar(self, rhs: Self::Scalar) -> Self::Out {
352        unary_kernel(self, |a| {
353            ArithmeticKernel::wrapping_floor_div_scalar(a.clone(), rhs)
354        })
355    }
356
357    fn wrapping_floor_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out {
358        unary_kernel(rhs, |a| {
359            ArithmeticKernel::wrapping_floor_div_scalar_lhs(lhs, a.clone())
360        })
361    }
362
363    fn wrapping_trunc_div_scalar(self, rhs: Self::Scalar) -> Self::Out {
364        unary_kernel(self, |a| {
365            ArithmeticKernel::wrapping_trunc_div_scalar(a.clone(), rhs)
366        })
367    }
368
369    fn wrapping_trunc_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out {
370        unary_kernel(rhs, |a| {
371            ArithmeticKernel::wrapping_trunc_div_scalar_lhs(lhs, a.clone())
372        })
373    }
374
375    fn wrapping_mod_scalar(self, rhs: Self::Scalar) -> Self::Out {
376        unary_kernel(self, |a| {
377            ArithmeticKernel::wrapping_mod_scalar(a.clone(), rhs)
378        })
379    }
380
381    fn wrapping_mod_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out {
382        unary_kernel(rhs, |a| {
383            ArithmeticKernel::wrapping_mod_scalar_lhs(lhs, a.clone())
384        })
385    }
386
387    fn checked_mul_scalar(self, rhs: Self::Scalar) -> Self::Out {
388        unary_kernel(self, |a| {
389            ArithmeticKernel::checked_mul_scalar(a.clone(), rhs)
390        })
391    }
392
393    fn true_div(self, rhs: Self) -> Self::TrueDivOut {
394        apply_binary_kernel_broadcast(
395            self,
396            rhs,
397            |l, r| ArithmeticKernel::true_div(l.clone(), r.clone()),
398            |l, r| ArithmeticKernel::true_div_scalar_lhs(l, r.clone()),
399            |l, r| ArithmeticKernel::true_div_scalar(l.clone(), r),
400        )
401    }
402
403    fn true_div_scalar(self, rhs: Self::Scalar) -> Self::TrueDivOut {
404        unary_kernel(self, |a| ArithmeticKernel::true_div_scalar(a.clone(), rhs))
405    }
406
407    fn true_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::TrueDivOut {
408        unary_kernel(rhs, |a| {
409            ArithmeticKernel::true_div_scalar_lhs(lhs, a.clone())
410        })
411    }
412
413    fn legacy_div(self, rhs: Self) -> Self::Out {
414        apply_binary_kernel_broadcast(
415            self,
416            rhs,
417            |l, r| ArithmeticKernel::legacy_div(l.clone(), r.clone()),
418            |l, r| ArithmeticKernel::legacy_div_scalar_lhs(l, r.clone()),
419            |l, r| ArithmeticKernel::legacy_div_scalar(l.clone(), r),
420        )
421    }
422
423    fn legacy_div_scalar(self, rhs: Self::Scalar) -> Self::Out {
424        unary_kernel(self, |a| {
425            ArithmeticKernel::legacy_div_scalar(a.clone(), rhs)
426        })
427    }
428
429    fn legacy_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self::Out {
430        unary_kernel(rhs, |a| {
431            ArithmeticKernel::legacy_div_scalar_lhs(lhs, a.clone())
432        })
433    }
434}
polars_core/chunked_array/arithmetic/numeric.rs

polars_core/chunked_array/arithmetic/
numeric.rs
