polars_core/frame/column/
compare.rs

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
use polars_error::PolarsResult;

use super::{BooleanChunked, ChunkCompareEq, ChunkCompareIneq, ChunkExpandAtIndex, Column, Series};

macro_rules! column_element_wise_broadcasting {
    ($lhs:expr, $rhs:expr, $op:expr) => {
        match ($lhs, $rhs) {
            (Column::Series(lhs), Column::Scalar(rhs)) => $op(lhs, &rhs.as_single_value_series()),
            (Column::Scalar(lhs), Column::Series(rhs)) => $op(&lhs.as_single_value_series(), rhs),
            (Column::Scalar(lhs), Column::Scalar(rhs)) => {
                $op(&lhs.as_single_value_series(), &rhs.as_single_value_series()).map(|ca| {
                    if ca.len() == 0 {
                        ca
                    } else {
                        ca.new_from_index(0, lhs.len())
                    }
                })
            },
            (lhs, rhs) => $op(lhs.as_materialized_series(), rhs.as_materialized_series()),
        }
    };
}

impl ChunkCompareEq<&Column> for Column {
    type Item = PolarsResult<BooleanChunked>;

    /// Create a boolean mask by checking for equality.
    #[inline]
    fn equal(&self, rhs: &Column) -> PolarsResult<BooleanChunked> {
        column_element_wise_broadcasting!(self, rhs, <Series as ChunkCompareEq<&Series>>::equal)
    }

    /// Create a boolean mask by checking for equality.
    #[inline]
    fn equal_missing(&self, rhs: &Column) -> PolarsResult<BooleanChunked> {
        column_element_wise_broadcasting!(
            self,
            rhs,
            <Series as ChunkCompareEq<&Series>>::equal_missing
        )
    }

    /// Create a boolean mask by checking for inequality.
    #[inline]
    fn not_equal(&self, rhs: &Column) -> PolarsResult<BooleanChunked> {
        column_element_wise_broadcasting!(self, rhs, <Series as ChunkCompareEq<&Series>>::not_equal)
    }

    /// Create a boolean mask by checking for inequality.
    #[inline]
    fn not_equal_missing(&self, rhs: &Column) -> PolarsResult<BooleanChunked> {
        column_element_wise_broadcasting!(
            self,
            rhs,
            <Series as ChunkCompareEq<&Series>>::not_equal_missing
        )
    }
}

impl ChunkCompareIneq<&Column> for Column {
    type Item = PolarsResult<BooleanChunked>;

    /// Create a boolean mask by checking if self > rhs.
    #[inline]
    fn gt(&self, rhs: &Column) -> PolarsResult<BooleanChunked> {
        column_element_wise_broadcasting!(self, rhs, <Series as ChunkCompareIneq<&Series>>::gt)
    }

    /// Create a boolean mask by checking if self >= rhs.
    #[inline]
    fn gt_eq(&self, rhs: &Column) -> PolarsResult<BooleanChunked> {
        column_element_wise_broadcasting!(self, rhs, <Series as ChunkCompareIneq<&Series>>::gt_eq)
    }

    /// Create a boolean mask by checking if self < rhs.
    #[inline]
    fn lt(&self, rhs: &Column) -> PolarsResult<BooleanChunked> {
        column_element_wise_broadcasting!(self, rhs, <Series as ChunkCompareIneq<&Series>>::lt)
    }

    /// Create a boolean mask by checking if self <= rhs.
    #[inline]
    fn lt_eq(&self, rhs: &Column) -> PolarsResult<BooleanChunked> {
        column_element_wise_broadcasting!(self, rhs, <Series as ChunkCompareIneq<&Series>>::lt_eq)
    }
}