polars_core/frame/column/
compare.rs

1use polars_error::PolarsResult;
2
3use super::{BooleanChunked, ChunkCompareEq, ChunkCompareIneq, ChunkExpandAtIndex, Column, Series};
4
5macro_rules! column_element_wise_broadcasting {
6    ($lhs:expr, $rhs:expr, $op:expr) => {
7        match ($lhs, $rhs) {
8            (Column::Series(lhs), Column::Scalar(rhs)) => $op(lhs, &rhs.as_single_value_series()),
9            (Column::Scalar(lhs), Column::Series(rhs)) => $op(&lhs.as_single_value_series(), rhs),
10            (Column::Scalar(lhs), Column::Scalar(rhs)) => {
11                $op(&lhs.as_single_value_series(), &rhs.as_single_value_series()).map(|ca| {
12                    if ca.len() == 0 {
13                        ca
14                    } else {
15                        ca.new_from_index(0, lhs.len())
16                    }
17                })
18            },
19            (lhs, rhs) => $op(lhs.as_materialized_series(), rhs.as_materialized_series()),
20        }
21    };
22}
23
24impl ChunkCompareEq<&Column> for Column {
25    type Item = PolarsResult<BooleanChunked>;
26
27    /// Create a boolean mask by checking for equality.
28    #[inline]
29    fn equal(&self, rhs: &Column) -> PolarsResult<BooleanChunked> {
30        column_element_wise_broadcasting!(self, rhs, <Series as ChunkCompareEq<&Series>>::equal)
31    }
32
33    /// Create a boolean mask by checking for equality.
34    #[inline]
35    fn equal_missing(&self, rhs: &Column) -> PolarsResult<BooleanChunked> {
36        column_element_wise_broadcasting!(
37            self,
38            rhs,
39            <Series as ChunkCompareEq<&Series>>::equal_missing
40        )
41    }
42
43    /// Create a boolean mask by checking for inequality.
44    #[inline]
45    fn not_equal(&self, rhs: &Column) -> PolarsResult<BooleanChunked> {
46        column_element_wise_broadcasting!(self, rhs, <Series as ChunkCompareEq<&Series>>::not_equal)
47    }
48
49    /// Create a boolean mask by checking for inequality.
50    #[inline]
51    fn not_equal_missing(&self, rhs: &Column) -> PolarsResult<BooleanChunked> {
52        column_element_wise_broadcasting!(
53            self,
54            rhs,
55            <Series as ChunkCompareEq<&Series>>::not_equal_missing
56        )
57    }
58}
59
60impl ChunkCompareIneq<&Column> for Column {
61    type Item = PolarsResult<BooleanChunked>;
62
63    /// Create a boolean mask by checking if self > rhs.
64    #[inline]
65    fn gt(&self, rhs: &Column) -> PolarsResult<BooleanChunked> {
66        column_element_wise_broadcasting!(self, rhs, <Series as ChunkCompareIneq<&Series>>::gt)
67    }
68
69    /// Create a boolean mask by checking if self >= rhs.
70    #[inline]
71    fn gt_eq(&self, rhs: &Column) -> PolarsResult<BooleanChunked> {
72        column_element_wise_broadcasting!(self, rhs, <Series as ChunkCompareIneq<&Series>>::gt_eq)
73    }
74
75    /// Create a boolean mask by checking if self < rhs.
76    #[inline]
77    fn lt(&self, rhs: &Column) -> PolarsResult<BooleanChunked> {
78        column_element_wise_broadcasting!(self, rhs, <Series as ChunkCompareIneq<&Series>>::lt)
79    }
80
81    /// Create a boolean mask by checking if self <= rhs.
82    #[inline]
83    fn lt_eq(&self, rhs: &Column) -> PolarsResult<BooleanChunked> {
84        column_element_wise_broadcasting!(self, rhs, <Series as ChunkCompareIneq<&Series>>::lt_eq)
85    }
86}