polars_compute/arithmetic/
pl_num.rs

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

polars_compute/arithmetic/
pl_num.rs
