polars_compute/arithmetic/
mod.rs

1use std::any::TypeId;
2
3use arrow::array::{Array, PrimitiveArray};
4use arrow::types::NativeType;
5
6// Low-level comparison kernel.
7pub trait ArithmeticKernel: Sized + Array {
8    type Scalar;
9    type TrueDivT: NativeType;
10
11    fn wrapping_abs(self) -> Self;
12    fn wrapping_neg(self) -> Self;
13    fn wrapping_add(self, rhs: Self) -> Self;
14    fn wrapping_sub(self, rhs: Self) -> Self;
15    fn wrapping_mul(self, rhs: Self) -> Self;
16    fn wrapping_floor_div(self, rhs: Self) -> Self;
17    fn wrapping_trunc_div(self, rhs: Self) -> Self;
18    fn wrapping_mod(self, rhs: Self) -> Self;
19
20    fn wrapping_add_scalar(self, rhs: Self::Scalar) -> Self;
21    fn wrapping_sub_scalar(self, rhs: Self::Scalar) -> Self;
22    fn wrapping_sub_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self;
23    fn wrapping_mul_scalar(self, rhs: Self::Scalar) -> Self;
24    fn wrapping_floor_div_scalar(self, rhs: Self::Scalar) -> Self;
25    fn wrapping_floor_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self;
26    fn wrapping_trunc_div_scalar(self, rhs: Self::Scalar) -> Self;
27    fn wrapping_trunc_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self;
28    fn wrapping_mod_scalar(self, rhs: Self::Scalar) -> Self;
29    fn wrapping_mod_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self;
30
31    fn true_div(self, rhs: Self) -> PrimitiveArray<Self::TrueDivT>;
32    fn true_div_scalar(self, rhs: Self::Scalar) -> PrimitiveArray<Self::TrueDivT>;
33    fn true_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> PrimitiveArray<Self::TrueDivT>;
34
35    // TODO: remove these.
36    // These are flooring division for integer types, true division for floating point types.
37    fn legacy_div(self, rhs: Self) -> Self {
38        if TypeId::of::<Self>() == TypeId::of::<PrimitiveArray<Self::TrueDivT>>() {
39            let ret = self.true_div(rhs);
40            unsafe {
41                let cast_ret = std::mem::transmute_copy(&ret);
42                std::mem::forget(ret);
43                cast_ret
44            }
45        } else {
46            self.wrapping_floor_div(rhs)
47        }
48    }
49    fn legacy_div_scalar(self, rhs: Self::Scalar) -> Self {
50        if TypeId::of::<Self>() == TypeId::of::<PrimitiveArray<Self::TrueDivT>>() {
51            let ret = self.true_div_scalar(rhs);
52            unsafe {
53                let cast_ret = std::mem::transmute_copy(&ret);
54                std::mem::forget(ret);
55                cast_ret
56            }
57        } else {
58            self.wrapping_floor_div_scalar(rhs)
59        }
60    }
61
62    fn legacy_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self {
63        if TypeId::of::<Self>() == TypeId::of::<PrimitiveArray<Self::TrueDivT>>() {
64            let ret = ArithmeticKernel::true_div_scalar_lhs(lhs, rhs);
65            unsafe {
66                let cast_ret = std::mem::transmute_copy(&ret);
67                std::mem::forget(ret);
68                cast_ret
69            }
70        } else {
71            ArithmeticKernel::wrapping_floor_div_scalar_lhs(lhs, rhs)
72        }
73    }
74}
75
76// Proxy trait so one can bound T: HasPrimitiveArithmeticKernel. Sadly Rust
77// doesn't support adding supertraits for other types.
78#[allow(private_bounds)]
79pub trait HasPrimitiveArithmeticKernel: NativeType + PrimitiveArithmeticKernelImpl {}
80impl<T: NativeType + PrimitiveArithmeticKernelImpl> HasPrimitiveArithmeticKernel for T {}
81
82use PrimitiveArray as PArr;
83
84#[doc(hidden)]
85pub trait PrimitiveArithmeticKernelImpl: NativeType {
86    type TrueDivT: NativeType;
87
88    fn prim_wrapping_abs(lhs: PArr<Self>) -> PArr<Self>;
89    fn prim_wrapping_neg(lhs: PArr<Self>) -> PArr<Self>;
90    fn prim_wrapping_add(lhs: PArr<Self>, rhs: PArr<Self>) -> PArr<Self>;
91    fn prim_wrapping_sub(lhs: PArr<Self>, rhs: PArr<Self>) -> PArr<Self>;
92    fn prim_wrapping_mul(lhs: PArr<Self>, rhs: PArr<Self>) -> PArr<Self>;
93    fn prim_wrapping_floor_div(lhs: PArr<Self>, rhs: PArr<Self>) -> PArr<Self>;
94    fn prim_wrapping_trunc_div(lhs: PArr<Self>, rhs: PArr<Self>) -> PArr<Self>;
95    fn prim_wrapping_mod(lhs: PArr<Self>, rhs: PArr<Self>) -> PArr<Self>;
96
97    fn prim_wrapping_add_scalar(lhs: PArr<Self>, rhs: Self) -> PArr<Self>;
98    fn prim_wrapping_sub_scalar(lhs: PArr<Self>, rhs: Self) -> PArr<Self>;
99    fn prim_wrapping_sub_scalar_lhs(lhs: Self, rhs: PArr<Self>) -> PArr<Self>;
100    fn prim_wrapping_mul_scalar(lhs: PArr<Self>, rhs: Self) -> PArr<Self>;
101    fn prim_wrapping_floor_div_scalar(lhs: PArr<Self>, rhs: Self) -> PArr<Self>;
102    fn prim_wrapping_floor_div_scalar_lhs(lhs: Self, rhs: PArr<Self>) -> PArr<Self>;
103    fn prim_wrapping_trunc_div_scalar(lhs: PArr<Self>, rhs: Self) -> PArr<Self>;
104    fn prim_wrapping_trunc_div_scalar_lhs(lhs: Self, rhs: PArr<Self>) -> PArr<Self>;
105    fn prim_wrapping_mod_scalar(lhs: PArr<Self>, rhs: Self) -> PArr<Self>;
106    fn prim_wrapping_mod_scalar_lhs(lhs: Self, rhs: PArr<Self>) -> PArr<Self>;
107
108    fn prim_true_div(lhs: PArr<Self>, rhs: PArr<Self>) -> PArr<Self::TrueDivT>;
109    fn prim_true_div_scalar(lhs: PArr<Self>, rhs: Self) -> PArr<Self::TrueDivT>;
110    fn prim_true_div_scalar_lhs(lhs: Self, rhs: PArr<Self>) -> PArr<Self::TrueDivT>;
111}
112
113#[rustfmt::skip]
114impl<T: HasPrimitiveArithmeticKernel> ArithmeticKernel for PrimitiveArray<T> {
115    type Scalar = T;
116    type TrueDivT = T::TrueDivT;
117
118    fn wrapping_abs(self) -> Self { T::prim_wrapping_abs(self) }
119    fn wrapping_neg(self) -> Self { T::prim_wrapping_neg(self) }
120    fn wrapping_add(self, rhs: Self) -> Self { T::prim_wrapping_add(self, rhs) }
121    fn wrapping_sub(self, rhs: Self) -> Self { T::prim_wrapping_sub(self, rhs) }
122    fn wrapping_mul(self, rhs: Self) -> Self { T::prim_wrapping_mul(self, rhs) }
123    fn wrapping_floor_div(self, rhs: Self) -> Self { T::prim_wrapping_floor_div(self, rhs) }
124    fn wrapping_trunc_div(self, rhs: Self) -> Self { T::prim_wrapping_trunc_div(self, rhs) }
125    fn wrapping_mod(self, rhs: Self) -> Self { T::prim_wrapping_mod(self, rhs) }
126
127    fn wrapping_add_scalar(self, rhs: Self::Scalar) -> Self { T::prim_wrapping_add_scalar(self, rhs) }
128    fn wrapping_sub_scalar(self, rhs: Self::Scalar) -> Self { T::prim_wrapping_sub_scalar(self, rhs) }
129    fn wrapping_sub_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self { T::prim_wrapping_sub_scalar_lhs(lhs, rhs) }
130    fn wrapping_mul_scalar(self, rhs: Self::Scalar) -> Self { T::prim_wrapping_mul_scalar(self, rhs) }
131    fn wrapping_floor_div_scalar(self, rhs: Self::Scalar) -> Self { T::prim_wrapping_floor_div_scalar(self, rhs) }
132    fn wrapping_floor_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self { T::prim_wrapping_floor_div_scalar_lhs(lhs, rhs) }
133    fn wrapping_trunc_div_scalar(self, rhs: Self::Scalar) -> Self { T::prim_wrapping_trunc_div_scalar(self, rhs) }
134    fn wrapping_trunc_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self { T::prim_wrapping_trunc_div_scalar_lhs(lhs, rhs) }
135    fn wrapping_mod_scalar(self, rhs: Self::Scalar) -> Self { T::prim_wrapping_mod_scalar(self, rhs) }
136    fn wrapping_mod_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> Self { T::prim_wrapping_mod_scalar_lhs(lhs, rhs) }
137
138    fn true_div(self, rhs: Self) -> PrimitiveArray<Self::TrueDivT> { T::prim_true_div(self, rhs) }
139    fn true_div_scalar(self, rhs: Self::Scalar) -> PrimitiveArray<Self::TrueDivT> { T::prim_true_div_scalar(self, rhs) }
140    fn true_div_scalar_lhs(lhs: Self::Scalar, rhs: Self) -> PrimitiveArray<Self::TrueDivT> { T::prim_true_div_scalar_lhs(lhs, rhs) }
141}
142
143mod float;
144pub mod pl_num;
145mod signed;
146mod unsigned;
polars_compute/arithmetic/mod.rs

polars_compute/arithmetic/
mod.rs
