polars_compute/arithmetic/
pl_num.rs

1use core::any::TypeId;
2
3use arrow::types::NativeType;
4use polars_utils::floor_divmod::FloorDivMod;
5
6/// Implements basic arithmetic between scalars with the same behavior as `ArithmeticKernel`.
7///
8/// Note, however, that the user is responsible for setting the validity of
9/// results for e.g. div/mod operations with 0 in the denominator.
10///
11/// This is intended as a low-level utility for custom arithmetic loops
12/// (e.g. in list arithmetic). In most cases prefer using `ArithmeticKernel` or
13/// `ArithmeticChunked` instead.
14pub trait PlNumArithmetic: Sized + Copy + 'static {
15    type TrueDivT: NativeType;
16
17    fn wrapping_abs(self) -> Self;
18    fn wrapping_neg(self) -> Self;
19    fn wrapping_add(self, rhs: Self) -> Self;
20    fn wrapping_sub(self, rhs: Self) -> Self;
21    fn wrapping_mul(self, rhs: Self) -> Self;
22    fn wrapping_floor_div(self, rhs: Self) -> Self;
23    fn wrapping_trunc_div(self, rhs: Self) -> Self;
24    fn wrapping_mod(self, rhs: Self) -> Self;
25
26    fn true_div(self, rhs: Self) -> Self::TrueDivT;
27
28    #[inline(always)]
29    fn legacy_div(self, rhs: Self) -> Self {
30        if TypeId::of::<Self>() == TypeId::of::<Self::TrueDivT>() {
31            let ret = self.true_div(rhs);
32            unsafe { core::mem::transmute_copy(&ret) }
33        } else {
34            self.wrapping_floor_div(rhs)
35        }
36    }
37}
38
39macro_rules! impl_signed_pl_num_arith {
40    ($T:ty) => {
41        impl PlNumArithmetic for $T {
42            type TrueDivT = f64;
43
44            #[inline(always)]
45            fn wrapping_abs(self) -> Self {
46                self.wrapping_abs()
47            }
48
49            #[inline(always)]
50            fn wrapping_neg(self) -> Self {
51                self.wrapping_neg()
52            }
53
54            #[inline(always)]
55            fn wrapping_add(self, rhs: Self) -> Self {
56                self.wrapping_add(rhs)
57            }
58
59            #[inline(always)]
60            fn wrapping_sub(self, rhs: Self) -> Self {
61                self.wrapping_sub(rhs)
62            }
63
64            #[inline(always)]
65            fn wrapping_mul(self, rhs: Self) -> Self {
66                self.wrapping_mul(rhs)
67            }
68
69            #[inline(always)]
70            fn wrapping_floor_div(self, rhs: Self) -> Self {
71                self.wrapping_floor_div_mod(rhs).0
72            }
73
74            #[inline(always)]
75            fn wrapping_trunc_div(self, rhs: Self) -> Self {
76                if rhs != 0 {
77                    self.wrapping_div(rhs)
78                } else {
79                    0
80                }
81            }
82
83            #[inline(always)]
84            fn wrapping_mod(self, rhs: Self) -> Self {
85                self.wrapping_floor_div_mod(rhs).1
86            }
87
88            #[inline(always)]
89            fn true_div(self, rhs: Self) -> Self::TrueDivT {
90                self as f64 / rhs as f64
91            }
92        }
93    };
94}
95
96impl_signed_pl_num_arith!(i8);
97impl_signed_pl_num_arith!(i16);
98impl_signed_pl_num_arith!(i32);
99impl_signed_pl_num_arith!(i64);
100impl_signed_pl_num_arith!(i128);
101
102macro_rules! impl_unsigned_pl_num_arith {
103    ($T:ty) => {
104        impl PlNumArithmetic for $T {
105            type TrueDivT = f64;
106
107            #[inline(always)]
108            fn wrapping_abs(self) -> Self {
109                self
110            }
111
112            #[inline(always)]
113            fn wrapping_neg(self) -> Self {
114                self.wrapping_neg()
115            }
116
117            #[inline(always)]
118            fn wrapping_add(self, rhs: Self) -> Self {
119                self.wrapping_add(rhs)
120            }
121
122            #[inline(always)]
123            fn wrapping_sub(self, rhs: Self) -> Self {
124                self.wrapping_sub(rhs)
125            }
126
127            #[inline(always)]
128            fn wrapping_mul(self, rhs: Self) -> Self {
129                self.wrapping_mul(rhs)
130            }
131
132            #[inline(always)]
133            fn wrapping_floor_div(self, rhs: Self) -> Self {
134                if rhs != 0 {
135                    self / rhs
136                } else {
137                    0
138                }
139            }
140
141            #[inline(always)]
142            fn wrapping_trunc_div(self, rhs: Self) -> Self {
143                self.wrapping_floor_div(rhs)
144            }
145
146            #[inline(always)]
147            fn wrapping_mod(self, rhs: Self) -> Self {
148                if rhs != 0 {
149                    self % rhs
150                } else {
151                    0
152                }
153            }
154
155            #[inline(always)]
156            fn true_div(self, rhs: Self) -> Self::TrueDivT {
157                self as f64 / rhs as f64
158            }
159        }
160    };
161}
162
163impl_unsigned_pl_num_arith!(u8);
164impl_unsigned_pl_num_arith!(u16);
165impl_unsigned_pl_num_arith!(u32);
166impl_unsigned_pl_num_arith!(u64);
167impl_unsigned_pl_num_arith!(u128);
168
169macro_rules! impl_float_pl_num_arith {
170    ($T:ty) => {
171        impl PlNumArithmetic for $T {
172            type TrueDivT = $T;
173
174            #[inline(always)]
175            fn wrapping_abs(self) -> Self {
176                self.abs()
177            }
178
179            #[inline(always)]
180            fn wrapping_neg(self) -> Self {
181                -self
182            }
183
184            #[inline(always)]
185            fn wrapping_add(self, rhs: Self) -> Self {
186                self + rhs
187            }
188
189            #[inline(always)]
190            fn wrapping_sub(self, rhs: Self) -> Self {
191                self - rhs
192            }
193
194            #[inline(always)]
195            fn wrapping_mul(self, rhs: Self) -> Self {
196                self * rhs
197            }
198
199            #[inline(always)]
200            fn wrapping_floor_div(self, rhs: Self) -> Self {
201                let l = self;
202                let r = rhs;
203                (l / r).floor()
204            }
205
206            #[inline(always)]
207            fn wrapping_trunc_div(self, rhs: Self) -> Self {
208                let l = self;
209                let r = rhs;
210                (l / r).trunc()
211            }
212
213            #[inline(always)]
214            fn wrapping_mod(self, rhs: Self) -> Self {
215                let l = self;
216                let r = rhs;
217                l - r * (l / r).floor()
218            }
219
220            #[inline(always)]
221            fn true_div(self, rhs: Self) -> Self::TrueDivT {
222                self / rhs
223            }
224        }
225    };
226}
227
228impl_float_pl_num_arith!(f32);
229impl_float_pl_num_arith!(f64);
polars_compute/arithmetic/pl_num.rs

polars_compute/arithmetic/
pl_num.rs
