feat(rust, python): support is_in for boolean dtype (pola-rs#5682)

zundertj · Jan 7, 2023 · b11d7f3 · b11d7f3
1 parent 2b406e5
commit b11d7f3
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/polars/polars-core/src/chunked_array/ops/is_in.rs b/polars/polars-core/src/chunked_array/ops/is_in.rs
@@ -329,6 +329,12 @@ impl IsIn for BooleanChunked {
                 ca.rename(self.name());
                 Ok(ca)
             }
+            DataType::Boolean => {
+                let other = other.bool().unwrap();
+                let has_true = other.any();
+                let has_false = !other.all();
+                Ok(self.apply(|v| if v { has_true } else { has_false }))
+            }
             _ => Err(PolarsError::SchemaMisMatch(
                 format!(
                     "cannot do is_in operation with left a dtype: {:?} and right a dtype {:?}",

diff --git a/py-polars/tests/unit/test_filter.py b/py-polars/tests/unit/test_filter.py
@@ -69,3 +69,11 @@ def test_filter_aggregation_any() -> None:
         "any": [[False, True, True], [True]],
         "filtered": [[3, 4], [2]],
     }
+
+
+def test_is_in_bool() -> None:
+    bool_value_to_filter_on = [True, None]
+    df = pl.DataFrame({"A": [True, False, None]})
+    assert df.filter(pl.col("A").is_in(bool_value_to_filter_on)).to_dict(False) == {
+        "A": [True, False]
+    }