polars_core/series/arithmetic/
owned.rs1use super::*;
2#[cfg(feature = "performant")]
3use crate::utils::align_chunks_binary_owned_series;
4
5#[cfg(feature = "performant")]
6pub fn coerce_lhs_rhs_owned(lhs: Series, rhs: Series) -> PolarsResult<(Series, Series)> {
7 let dtype = try_get_supertype(lhs.dtype(), rhs.dtype())?;
8 let left = if lhs.dtype() == &dtype {
9 lhs
10 } else {
11 lhs.cast(&dtype)?
12 };
13 let right = if rhs.dtype() == &dtype {
14 rhs
15 } else {
16 rhs.cast(&dtype)?
17 };
18 Ok((left, right))
19}
20
21fn is_eligible(lhs: &DataType, rhs: &DataType) -> bool {
22 !lhs.is_logical()
23 && lhs.to_physical().is_primitive_numeric()
24 && rhs.to_physical().is_primitive_numeric()
25}
26
27#[cfg(feature = "performant")]
28fn apply_operation_mut<T, F>(mut lhs: Series, mut rhs: Series, op: F) -> Series
29where
30 T: PolarsNumericType,
31 F: Fn(ChunkedArray<T>, ChunkedArray<T>) -> ChunkedArray<T> + Copy,
32 ChunkedArray<T>: IntoSeries,
33{
34 let lhs_ca: &mut ChunkedArray<T> = lhs._get_inner_mut().as_mut();
35 let rhs_ca: &mut ChunkedArray<T> = rhs._get_inner_mut().as_mut();
36
37 let lhs = std::mem::take(lhs_ca);
38 let rhs = std::mem::take(rhs_ca);
39
40 op(lhs, rhs).into_series()
41}
42
43macro_rules! impl_operation {
44 ($operation:ident, $method:ident, $function:expr) => {
45 impl $operation for Series {
46 type Output = PolarsResult<Series>;
47
48 fn $method(self, rhs: Self) -> Self::Output {
49 #[cfg(feature = "performant")]
50 {
51 if is_eligible(self.dtype(), rhs.dtype()) {
53 let (lhs, rhs) = coerce_lhs_rhs_owned(self, rhs).unwrap();
54 let (lhs, rhs) = align_chunks_binary_owned_series(lhs, rhs);
55 use DataType::*;
56 Ok(match lhs.dtype() {
57 #[cfg(feature = "dtype-i8")]
58 Int8 => apply_operation_mut::<Int8Type, _>(lhs, rhs, $function),
59 #[cfg(feature = "dtype-i16")]
60 Int16 => apply_operation_mut::<Int16Type, _>(lhs, rhs, $function),
61 Int32 => apply_operation_mut::<Int32Type, _>(lhs, rhs, $function),
62 Int64 => apply_operation_mut::<Int64Type, _>(lhs, rhs, $function),
63 #[cfg(feature = "dtype-u8")]
64 UInt8 => apply_operation_mut::<UInt8Type, _>(lhs, rhs, $function),
65 #[cfg(feature = "dtype-u16")]
66 UInt16 => apply_operation_mut::<UInt16Type, _>(lhs, rhs, $function),
67 UInt32 => apply_operation_mut::<UInt32Type, _>(lhs, rhs, $function),
68 UInt64 => apply_operation_mut::<UInt64Type, _>(lhs, rhs, $function),
69 Float32 => apply_operation_mut::<Float32Type, _>(lhs, rhs, $function),
70 Float64 => apply_operation_mut::<Float64Type, _>(lhs, rhs, $function),
71 _ => unreachable!(),
72 })
73 } else {
74 (&self).$method(&rhs)
75 }
76 }
77 #[cfg(not(feature = "performant"))]
78 {
79 (&self).$method(&rhs)
80 }
81 }
82 }
83 };
84}
85
86impl_operation!(Add, add, |a, b| a.add(b));
87impl_operation!(Sub, sub, |a, b| a.sub(b));
88impl_operation!(Mul, mul, |a, b| a.mul(b));
89impl_operation!(Div, div, |a, b| a.div(b));
90
91impl Series {
92 pub fn try_add_owned(self, other: Self) -> PolarsResult<Self> {
93 if is_eligible(self.dtype(), other.dtype()) {
94 self + other
95 } else {
96 std::ops::Add::add(&self, &other)
97 }
98 }
99
100 pub fn try_sub_owned(self, other: Self) -> PolarsResult<Self> {
101 if is_eligible(self.dtype(), other.dtype()) {
102 self - other
103 } else {
104 std::ops::Sub::sub(&self, &other)
105 }
106 }
107
108 pub fn try_mul_owned(self, other: Self) -> PolarsResult<Self> {
109 if is_eligible(self.dtype(), other.dtype()) {
110 self * other
111 } else {
112 std::ops::Mul::mul(&self, &other)
113 }
114 }
115}