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 { self.wrapping_div(rhs) } else { 0 }
77            }
78
79            #[inline(always)]
80            fn wrapping_mod(self, rhs: Self) -> Self {
81                self.wrapping_floor_div_mod(rhs).1
82            }
83
84            #[inline(always)]
85            fn true_div(self, rhs: Self) -> Self::TrueDivT {
86                self as f64 / rhs as f64
87            }
88        }
89    };
90}
91
92impl_signed_pl_num_arith!(i8);
93impl_signed_pl_num_arith!(i16);
94impl_signed_pl_num_arith!(i32);
95impl_signed_pl_num_arith!(i64);
96impl_signed_pl_num_arith!(i128);
97
98macro_rules! impl_unsigned_pl_num_arith {
99    ($T:ty) => {
100        impl PlNumArithmetic for $T {
101            type TrueDivT = f64;
102
103            #[inline(always)]
104            fn wrapping_abs(self) -> Self {
105                self
106            }
107
108            #[inline(always)]
109            fn wrapping_neg(self) -> Self {
110                self.wrapping_neg()
111            }
112
113            #[inline(always)]
114            fn wrapping_add(self, rhs: Self) -> Self {
115                self.wrapping_add(rhs)
116            }
117
118            #[inline(always)]
119            fn wrapping_sub(self, rhs: Self) -> Self {
120                self.wrapping_sub(rhs)
121            }
122
123            #[inline(always)]
124            fn wrapping_mul(self, rhs: Self) -> Self {
125                self.wrapping_mul(rhs)
126            }
127
128            #[inline(always)]
129            fn wrapping_floor_div(self, rhs: Self) -> Self {
130                if rhs != 0 { self / rhs } else { 0 }
131            }
132
133            #[inline(always)]
134            fn wrapping_trunc_div(self, rhs: Self) -> Self {
135                self.wrapping_floor_div(rhs)
136            }
137
138            #[inline(always)]
139            fn wrapping_mod(self, rhs: Self) -> Self {
140                if rhs != 0 { self % rhs } else { 0 }
141            }
142
143            #[inline(always)]
144            fn true_div(self, rhs: Self) -> Self::TrueDivT {
145                self as f64 / rhs as f64
146            }
147        }
148    };
149}
150
151impl_unsigned_pl_num_arith!(u8);
152impl_unsigned_pl_num_arith!(u16);
153impl_unsigned_pl_num_arith!(u32);
154impl_unsigned_pl_num_arith!(u64);
155impl_unsigned_pl_num_arith!(u128);
156
157macro_rules! impl_float_pl_num_arith {
158    ($T:ty) => {
159        impl PlNumArithmetic for $T {
160            type TrueDivT = $T;
161
162            #[inline(always)]
163            fn wrapping_abs(self) -> Self {
164                self.abs()
165            }
166
167            #[inline(always)]
168            fn wrapping_neg(self) -> Self {
169                -self
170            }
171
172            #[inline(always)]
173            fn wrapping_add(self, rhs: Self) -> Self {
174                self + rhs
175            }
176
177            #[inline(always)]
178            fn wrapping_sub(self, rhs: Self) -> Self {
179                self - rhs
180            }
181
182            #[inline(always)]
183            fn wrapping_mul(self, rhs: Self) -> Self {
184                self * rhs
185            }
186
187            #[inline(always)]
188            fn wrapping_floor_div(self, rhs: Self) -> Self {
189                let l = self;
190                let r = rhs;
191                (l / r).floor()
192            }
193
194            #[inline(always)]
195            fn wrapping_trunc_div(self, rhs: Self) -> Self {
196                let l = self;
197                let r = rhs;
198                (l / r).trunc()
199            }
200
201            #[inline(always)]
202            fn wrapping_mod(self, rhs: Self) -> Self {
203                let l = self;
204                let r = rhs;
205                l - r * (l / r).floor()
206            }
207
208            #[inline(always)]
209            fn true_div(self, rhs: Self) -> Self::TrueDivT {
210                self / rhs
211            }
212        }
213    };
214}
215
216impl_float_pl_num_arith!(f32);
217impl_float_pl_num_arith!(f64);
polars_compute/arithmetic/pl_num.rs

polars_compute/arithmetic/
pl_num.rs
