polars_core/chunked_array/
bitwise.rs

1use std::ops::{BitAnd, BitOr, BitXor};
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        if let Some((scalar, other_ca)) = 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) => None,
113            (1, _) => Some((self.get(0), rhs)),
114            (_, 1) => Some((rhs.get(0), self)),
115            _ => None,
116        } {
117            match scalar {
118                Some(false) => other_ca.clone(),
119                None => BooleanChunked::full_null(self.name().clone(), other_ca.len()),
120                Some(true) => !other_ca,
121            }
122        } else {
123            arity::binary(self, rhs, |l_arr, r_arr| {
124                let validity = combine_validities_and(l_arr.validity(), r_arr.validity());
125                let values = l_arr.values() ^ r_arr.values();
126                BooleanArray::from_data_default(values, validity)
127            })
128        }
129    }
130}
131
132impl BitXor for BooleanChunked {
133    type Output = BooleanChunked;
134
135    fn bitxor(self, rhs: Self) -> Self::Output {
136        (&self).bitxor(&rhs)
137    }
138}
139
140impl BitAnd for &BooleanChunked {
141    type Output = BooleanChunked;
142
143    fn bitand(self, rhs: Self) -> Self::Output {
144        match (self.len(), rhs.len()) {
145            // make sure that we fall through if both are equal unit lengths
146            // otherwise we stackoverflow
147            (1, 1) => {},
148            (1, _) => {
149                return match self.get(0) {
150                    Some(true) => rhs.clone().with_name(self.name().clone()),
151                    Some(false) => BooleanChunked::full(self.name().clone(), false, rhs.len()),
152                    None => &self.new_from_index(0, rhs.len()) & rhs,
153                };
154            },
155            (_, 1) => {
156                return match rhs.get(0) {
157                    Some(true) => self.clone(),
158                    Some(false) => BooleanChunked::full(self.name().clone(), false, self.len()),
159                    None => self & &rhs.new_from_index(0, self.len()),
160                };
161            },
162            _ => {},
163        }
164
165        arity::binary(self, rhs, compute::boolean_kleene::and)
166    }
167}
168
169impl BitAnd for BooleanChunked {
170    type Output = BooleanChunked;
171
172    fn bitand(self, rhs: Self) -> Self::Output {
173        (&self).bitand(&rhs)
174    }
175}
176
177#[cfg(test)]
178mod test {
179    use super::*;
180
181    #[test]
182    fn guard_so_issue_2494() {
183        // this cause a stack overflow
184        let a = BooleanChunked::new(PlSmallStr::from_static("a"), [None]);
185        let b = BooleanChunked::new(PlSmallStr::from_static("b"), [None]);
186
187        assert_eq!((&a).bitand(&b).null_count(), 1);
188        assert_eq!((&a).bitor(&b).null_count(), 1);
189        assert_eq!((&a).bitxor(&b).null_count(), 1);
190    }
191}
polars_core/chunked_array/bitwise.rs

polars_core/chunked_array/
bitwise.rs
