polars_core/chunked_array/
bitwise.rs

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

polars_core/chunked_array/
bitwise.rs
