1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

use super::*;
#[cfg(feature = "performant")]
use crate::utils::align_chunks_binary_owned_series;

#[cfg(feature = "performant")]
pub fn coerce_lhs_rhs_owned(lhs: Series, rhs: Series) -> PolarsResult<(Series, Series)> {
    let dtype = try_get_supertype(lhs.dtype(), rhs.dtype())?;
    let left = if lhs.dtype() == &dtype {
        lhs
    } else {
        lhs.cast(&dtype)?
    };
    let right = if rhs.dtype() == &dtype {
        rhs
    } else {
        rhs.cast(&dtype)?
    };
    Ok((left, right))
}

fn is_eligible(lhs: &DataType, rhs: &DataType) -> bool {
    !lhs.is_logical() && lhs.to_physical().is_numeric() && rhs.to_physical().is_numeric()
}

#[cfg(feature = "performant")]
fn apply_operation_mut<T, F>(mut lhs: Series, mut rhs: Series, op: F) -> Series
where
    T: PolarsNumericType,
    F: Fn(ChunkedArray<T>, ChunkedArray<T>) -> ChunkedArray<T> + Copy,
    ChunkedArray<T>: IntoSeries,
{
    let lhs_ca: &mut ChunkedArray<T> = lhs._get_inner_mut().as_mut();
    let rhs_ca: &mut ChunkedArray<T> = rhs._get_inner_mut().as_mut();

    let lhs = std::mem::take(lhs_ca);
    let rhs = std::mem::take(rhs_ca);

    op(lhs, rhs).into_series()
}

macro_rules! impl_operation {
    ($operation:ident, $method:ident, $function:expr) => {
        impl $operation for Series {
            type Output = Series;

            fn $method(self, rhs: Self) -> Self::Output {
                #[cfg(feature = "performant")]
                {
                    // only physical numeric values take the mutable path
                    if is_eligible(self.dtype(), rhs.dtype()) {
                        let (lhs, rhs) = coerce_lhs_rhs_owned(self, rhs).unwrap();
                        let (lhs, rhs) = align_chunks_binary_owned_series(lhs, rhs);
                        use DataType::*;
                        match lhs.dtype() {
                            #[cfg(feature = "dtype-i8")]
                            Int8 => apply_operation_mut::<Int8Type, _>(lhs, rhs, $function),
                            #[cfg(feature = "dtype-i16")]
                            Int16 => apply_operation_mut::<Int16Type, _>(lhs, rhs, $function),
                            Int32 => apply_operation_mut::<Int32Type, _>(lhs, rhs, $function),
                            Int64 => apply_operation_mut::<Int64Type, _>(lhs, rhs, $function),
                            #[cfg(feature = "dtype-u8")]
                            UInt8 => apply_operation_mut::<UInt8Type, _>(lhs, rhs, $function),
                            #[cfg(feature = "dtype-u16")]
                            UInt16 => apply_operation_mut::<UInt16Type, _>(lhs, rhs, $function),
                            UInt32 => apply_operation_mut::<UInt32Type, _>(lhs, rhs, $function),
                            UInt64 => apply_operation_mut::<UInt64Type, _>(lhs, rhs, $function),
                            Float32 => apply_operation_mut::<Float32Type, _>(lhs, rhs, $function),
                            Float64 => apply_operation_mut::<Float64Type, _>(lhs, rhs, $function),
                            _ => unreachable!(),
                        }
                    } else {
                        (&self).$method(&rhs)
                    }
                }
                #[cfg(not(feature = "performant"))]
                {
                    (&self).$method(&rhs)
                }
            }
        }
    };
}

impl_operation!(Add, add, |a, b| a.add(b));
impl_operation!(Sub, sub, |a, b| a.sub(b));
impl_operation!(Mul, mul, |a, b| a.mul(b));
impl_operation!(Div, div, |a, b| a.div(b));

impl Series {
    pub fn try_add_owned(self, other: Self) -> PolarsResult<Self> {
        if is_eligible(self.dtype(), other.dtype()) {
            Ok(self + other)
        } else {
            self.try_add(&other)
        }
    }

    pub fn try_sub_owned(self, other: Self) -> PolarsResult<Self> {
        if is_eligible(self.dtype(), other.dtype()) {
            Ok(self - other)
        } else {
            self.try_sub(&other)
        }
    }

    pub fn try_mul_owned(self, other: Self) -> PolarsResult<Self> {
        if is_eligible(self.dtype(), other.dtype()) {
            Ok(self * other)
        } else {
            self.try_mul(&other)
        }
    }
}