polars_core/chunked_array/
bitwise.rs

1use std::ops::{BitAnd, BitOr, BitXor};
2
3use arrow::compute::bitwise;
4use arrow::compute::utils::combine_validities_and;
5
6use super::*;
7use crate::chunked_array::arity::apply_binary_kernel_broadcast;
8
9impl<T> BitAnd for &ChunkedArray<T>
10where
11    T: PolarsIntegerType,
12    T::Native: BitAnd<Output = T::Native>,
13{
14    type Output = ChunkedArray<T>;
15
16    fn bitand(self, rhs: Self) -> Self::Output {
17        apply_binary_kernel_broadcast(
18            self,
19            rhs,
20            bitwise::and,
21            |l, r| bitwise::and_scalar(r, &l),
22            |l, r| bitwise::and_scalar(l, &r),
23        )
24    }
25}
26
27impl<T> BitOr for &ChunkedArray<T>
28where
29    T: PolarsIntegerType,
30    T::Native: BitOr<Output = T::Native>,
31{
32    type Output = ChunkedArray<T>;
33
34    fn bitor(self, rhs: Self) -> Self::Output {
35        apply_binary_kernel_broadcast(
36            self,
37            rhs,
38            bitwise::or,
39            |l, r| bitwise::or_scalar(r, &l),
40            |l, r| bitwise::or_scalar(l, &r),
41        )
42    }
43}
44
45impl<T> BitXor for &ChunkedArray<T>
46where
47    T: PolarsIntegerType,
48    T::Native: BitXor<Output = T::Native>,
49{
50    type Output = ChunkedArray<T>;
51
52    fn bitxor(self, rhs: Self) -> Self::Output {
53        apply_binary_kernel_broadcast(
54            self,
55            rhs,
56            bitwise::xor,
57            |l, r| bitwise::xor_scalar(r, &l),
58            |l, r| bitwise::xor_scalar(l, &r),
59        )
60    }
61}
62
63impl BitOr for &BooleanChunked {
64    type Output = BooleanChunked;
65
66    fn bitor(self, rhs: Self) -> Self::Output {
67        match (self.len(), rhs.len()) {
68            // make sure that we fall through if both are equal unit lengths
69            // otherwise we stackoverflow
70            (1, 1) => {},
71            (1, _) => {
72                return match self.get(0) {
73                    Some(true) => BooleanChunked::full(self.name().clone(), true, rhs.len()),
74                    Some(false) => {
75                        let mut rhs = rhs.clone();
76                        rhs.rename(self.name().clone());
77                        rhs
78                    },
79                    None => &self.new_from_index(0, rhs.len()) | rhs,
80                };
81            },
82            (_, 1) => {
83                return match rhs.get(0) {
84                    Some(true) => BooleanChunked::full(self.name().clone(), true, self.len()),
85                    Some(false) => self.clone(),
86                    None => self | &rhs.new_from_index(0, self.len()),
87                };
88            },
89            _ => {},
90        }
91
92        arity::binary(self, rhs, polars_compute::boolean::or)
93    }
94}
95
96impl BitOr for BooleanChunked {
97    type Output = BooleanChunked;
98
99    fn bitor(self, rhs: Self) -> Self::Output {
100        (&self).bitor(&rhs)
101    }
102}
103
104impl BitXor for &BooleanChunked {
105    type Output = BooleanChunked;
106
107    fn bitxor(self, rhs: Self) -> Self::Output {
108        if let Some((scalar, other_ca)) = match (self.len(), rhs.len()) {
109            // make sure that we fall through if both are equal unit lengths
110            // otherwise we stackoverflow
111            (1, 1) => None,
112            (1, _) => Some((self.get(0), rhs)),
113            (_, 1) => Some((rhs.get(0), self)),
114            _ => None,
115        } {
116            match scalar {
117                Some(false) => other_ca.clone(),
118                None => BooleanChunked::full_null(self.name().clone(), other_ca.len()),
119                Some(true) => !other_ca,
120            }
121        } else {
122            arity::binary(self, rhs, |l_arr, r_arr| {
123                let validity = combine_validities_and(l_arr.validity(), r_arr.validity());
124                let values = l_arr.values() ^ r_arr.values();
125                BooleanArray::from_data_default(values, validity)
126            })
127        }
128    }
129}
130
131impl BitXor for BooleanChunked {
132    type Output = BooleanChunked;
133
134    fn bitxor(self, rhs: Self) -> Self::Output {
135        (&self).bitxor(&rhs)
136    }
137}
138
139impl BitAnd for &BooleanChunked {
140    type Output = BooleanChunked;
141
142    fn bitand(self, rhs: Self) -> Self::Output {
143        match (self.len(), rhs.len()) {
144            // make sure that we fall through if both are equal unit lengths
145            // otherwise we stackoverflow
146            (1, 1) => {},
147            (1, _) => {
148                return match self.get(0) {
149                    Some(true) => rhs.clone().with_name(self.name().clone()),
150                    Some(false) => BooleanChunked::full(self.name().clone(), false, rhs.len()),
151                    None => &self.new_from_index(0, rhs.len()) & rhs,
152                };
153            },
154            (_, 1) => {
155                return match rhs.get(0) {
156                    Some(true) => self.clone(),
157                    Some(false) => BooleanChunked::full(self.name().clone(), false, self.len()),
158                    None => self & &rhs.new_from_index(0, self.len()),
159                };
160            },
161            _ => {},
162        }
163
164        arity::binary(self, rhs, polars_compute::boolean::and)
165    }
166}
167
168impl BitAnd for BooleanChunked {
169    type Output = BooleanChunked;
170
171    fn bitand(self, rhs: Self) -> Self::Output {
172        (&self).bitand(&rhs)
173    }
174}
175
176#[cfg(test)]
177mod test {
178    use super::*;
179
180    #[test]
181    fn guard_so_issue_2494() {
182        // this cause a stack overflow
183        let a = BooleanChunked::new(PlSmallStr::from_static("a"), [None]);
184        let b = BooleanChunked::new(PlSmallStr::from_static("b"), [None]);
185
186        assert_eq!((&a).bitand(&b).null_count(), 1);
187        assert_eq!((&a).bitor(&b).null_count(), 1);
188        assert_eq!((&a).bitxor(&b).null_count(), 1);
189    }
190}
polars_core/chunked_array/bitwise.rs

polars_core/chunked_array/
bitwise.rs
