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

use crate::prelude::*;

impl<T> IsIn for ChunkedArray<T>
where
    T: PolarsNumericType,
{
    fn is_in(&self, list_array: &ListChunked) -> Result<BooleanChunked> {
        match list_array.dtype() {
            DataType::List(dt) if self.dtype() == dt => {
                let ca: BooleanChunked = self
                    .into_iter()
                    .zip(list_array.into_iter())
                    .map(|(value, series)| match (value, series) {
                        (val, Some(series)) => {
                            let ca = series.unpack::<T>().unwrap();
                            ca.into_iter().any(|a| a == val)
                        }
                        _ => false,
                    })
                    .collect();
                Ok(ca)
            }
            _ => Err(PolarsError::DataTypeMisMatch(
                format!(
                    "cannot do is_in operation with left a dtype: {:?} and right a dtype {:?}",
                    self.dtype(),
                    list_array.dtype()
                )
                .into(),
            )),
        }
    }
}
impl IsIn for Utf8Chunked {
    fn is_in(&self, list_array: &ListChunked) -> Result<BooleanChunked> {
        match list_array.dtype() {
            DataType::List(dt) if self.dtype() == dt => {
                let ca: BooleanChunked = self
                    .into_iter()
                    .zip(list_array.into_iter())
                    .map(|(value, series)| match (value, series) {
                        (val, Some(series)) => {
                            let ca = series.unpack::<Utf8Type>().unwrap();
                            ca.into_iter().any(|a| a == val)
                        }
                        _ => false,
                    })
                    .collect();
                Ok(ca)
            }
            _ => Err(PolarsError::DataTypeMisMatch(
                format!(
                    "cannot do is_in operation with left a dtype: {:?} and right a dtype {:?}",
                    self.dtype(),
                    list_array.dtype()
                )
                .into(),
            )),
        }
    }
}

impl IsIn for BooleanChunked {
    fn is_in(&self, list_array: &ListChunked) -> Result<BooleanChunked> {
        match list_array.dtype() {
            DataType::List(dt) if self.dtype() == dt => {
                let ca: BooleanChunked = self
                    .into_iter()
                    .zip(list_array.into_iter())
                    .map(|(value, series)| match (value, series) {
                        (val, Some(series)) => {
                            let ca = series.unpack::<BooleanType>().unwrap();
                            ca.into_iter().any(|a| a == val)
                        }
                        _ => false,
                    })
                    .collect();
                Ok(ca)
            }
            _ => Err(PolarsError::DataTypeMisMatch(
                format!(
                    "cannot do is_in operation with left a dtype: {:?} and right a dtype {:?}",
                    self.dtype(),
                    list_array.dtype()
                )
                .into(),
            )),
        }
    }
}

impl IsIn for CategoricalChunked {
    fn is_in(&self, list_array: &ListChunked) -> Result<BooleanChunked> {
        self.cast::<UInt32Type>().unwrap().is_in(list_array)
    }
}

impl IsIn for ListChunked {}