polars_core/frame/column/
arithmetic.rs

1use num_traits::{Num, NumCast};
2use polars_error::PolarsResult;
3
4use super::{Column, ScalarColumn, Series};
5
6fn num_op_with_broadcast<T: Num + NumCast, F: Fn(&Series, T) -> Series>(
7    c: &'_ Column,
8    n: T,
9    op: F,
10) -> Column {
11    match c {
12        Column::Series(s) => op(s, n).into(),
13        // @partition-opt
14        Column::Partitioned(s) => op(s.as_materialized_series(), n).into(),
15        Column::Scalar(s) => {
16            ScalarColumn::from_single_value_series(op(&s.as_single_value_series(), n), s.len())
17                .into()
18        },
19    }
20}
21
22macro_rules! broadcastable_ops {
23    ($(($trait:ident, $op:ident))+) => {
24        $(
25        impl std::ops::$trait for Column {
26            type Output = PolarsResult<Column>;
27
28            #[inline]
29            fn $op(self, rhs: Self) -> Self::Output {
30                self.try_apply_broadcasting_binary_elementwise(&rhs, |l, r| l.$op(r))
31            }
32        }
33
34        impl std::ops::$trait for &Column {
35            type Output = PolarsResult<Column>;
36
37            #[inline]
38            fn $op(self, rhs: Self) -> Self::Output {
39                self.try_apply_broadcasting_binary_elementwise(rhs, |l, r| l.$op(r))
40            }
41        }
42        )+
43    }
44}
45
46macro_rules! broadcastable_num_ops {
47    ($(($trait:ident, $op:ident))+) => {
48        $(
49        impl<T> std::ops::$trait::<T> for Column
50        where
51            T: Num + NumCast,
52        {
53            type Output = Self;
54
55            #[inline]
56            fn $op(self, rhs: T) -> Self::Output {
57                num_op_with_broadcast(&self, rhs, |l, r| l.$op(r))
58            }
59        }
60
61        impl<T> std::ops::$trait::<T> for &Column
62        where
63            T: Num + NumCast,
64        {
65            type Output = Column;
66
67            #[inline]
68            fn $op(self, rhs: T) -> Self::Output {
69                num_op_with_broadcast(self, rhs, |l, r| l.$op(r))
70            }
71        }
72        )+
73    };
74}
75
76broadcastable_ops! {
77    (Add, add)
78    (Sub, sub)
79    (Mul, mul)
80    (Div, div)
81    (Rem, rem)
82    (BitAnd, bitand)
83    (BitOr, bitor)
84    (BitXor, bitxor)
85}
86
87broadcastable_num_ops! {
88    (Add, add)
89    (Sub, sub)
90    (Mul, mul)
91    (Div, div)
92    (Rem, rem)
93}
polars_core/frame/column/arithmetic.rs

polars_core/frame/column/
arithmetic.rs
