fix boolean Xor broadcast (#2122)

pola-rs · Dec 22, 2021 · 17c5568 · 17c5568
1 parent ae4384a
commit 17c5568
Showing 1 changed file with 23 additions and 1 deletion.
diff --git a/polars/polars-core/src/chunked_array/bitwise.rs b/polars/polars-core/src/chunked_array/bitwise.rs
@@ -2,7 +2,7 @@ use super::*;
 use crate::utils::{align_chunks_binary, combine_validities, CustomIterTools};
 use arrow::bitmap::MutableBitmap;
 use arrow::compute;
-use std::ops::{BitAnd, BitOr, BitXor};
+use std::ops::{BitAnd, BitOr, BitXor, Not};
 
 impl<T> BitAnd for &ChunkedArray<T>
 where
@@ -150,6 +150,28 @@ impl BitXor for &BooleanChunked {
     type Output = BooleanChunked;
 
     fn bitxor(self, rhs: Self) -> Self::Output {
+        if self.len() == 1 {
+            return match self.get(0) {
+                Some(true) => {
+                    let mut rhs = rhs.not();
+                    rhs.rename(self.name());
+                    rhs
+                }
+                Some(false) => {
+                    let mut rhs = rhs.clone();
+                    rhs.rename(self.name());
+                    rhs
+                }
+                None => &self.expand_at_index(0, rhs.len()) | rhs,
+            };
+        } else if rhs.len() == 1 {
+            return match rhs.get(0) {
+                Some(true) => self.not(),
+                Some(false) => self.clone(),
+                None => &rhs.expand_at_index(0, self.len()) | self,
+            };
+        }
+
         let (l, r) = align_chunks_binary(self, rhs);
         let chunks = l
             .downcast_iter()