fix(rust, python): fix regression in boolean nulls comparison (pola-r…

…s#9142)
c-peters · Jul 14, 2023 · 3d3ce52 · 3d3ce52
1 parent bcd4082
commit 3d3ce52
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 3 deletions.
diff --git a/polars/polars-core/src/chunked_array/comparison/mod.rs b/polars/polars-core/src/chunked_array/comparison/mod.rs
@@ -298,9 +298,29 @@ impl ChunkCompare<&BooleanChunked> for BooleanChunked {
             (_, 1) => {
                 if let Some(value) = rhs.get(0) {
                     if value {
-                        self.clone()
+                        if self.null_count() == 0 {
+                            self.clone()
+                        } else {
+                            self.apply_kernel(&|arr| {
+                                if let Some(validity) = arr.validity() {
+                                    Box::new(BooleanArray::from_data_default(
+                                        arr.values() & validity,
+                                        None,
+                                    )) as ArrayRef
+                                } else {
+                                    Box::new(arr.clone())
+                                }
+                            })
+                        }
                     } else {
-                        self.not()
+                        self.apply_kernel(&|arr| {
+                            let bitmap = if let Some(validity) = arr.validity() {
+                                arr.values() ^ validity
+                            } else {
+                                arr.values().not()
+                            };
+                            Box::new(BooleanArray::from_data_default(bitmap, None)) as ArrayRef
+                        })
                     }
                 } else {
                     self.is_null()

diff --git a/polars/polars-core/src/chunked_array/from.rs b/polars/polars-core/src/chunked_array/from.rs
@@ -85,6 +85,14 @@ where
             dtype @ DataType::Array(_, _) => from_chunks_list_dtype(&mut chunks, dtype),
             dt => dt,
         };
+        // assertions in debug mode
+        // that check if the data types in the arrays are as expected
+        #[cfg(debug_assertions)]
+        {
+            if !chunks.is_empty() && dtype.is_primitive() {
+                assert_eq!(chunks[0].data_type(), &dtype.to_physical().to_arrow())
+            }
+        }
         let field = Arc::new(Field::new(name, dtype));
         let mut out = ChunkedArray {
             field,

diff --git a/polars/polars-ops/src/chunked_array/list/count.rs b/polars/polars-ops/src/chunked_array/list/count.rs
@@ -29,8 +29,9 @@ pub fn list_count_match(ca: &ListChunked, value: AnyValue) -> PolarsResult<Serie
     let value = Series::new("", [value]);
 
     let ca = ca.apply_to_inner(&|s| {
-        ChunkCompare::<&Series>::equal(&s, &value).map(|ca| ca.into_series())
+        ChunkCompare::<&Series>::equal_missing(&s, &value).map(|ca| ca.into_series())
     })?;
+    dbg!(&ca);
     let out = count_boolean_bits(&ca);
     Ok(out.into_series())
 }

diff --git a/py-polars/tests/unit/namespaces/test_list.py b/py-polars/tests/unit/namespaces/test_list.py
@@ -448,3 +448,9 @@ def test_list_tail_underflow_9087() -> None:
     assert pl.Series([["a", "b", "c"]]).list.tail(pl.lit(1, pl.UInt32)).to_list() == [
         ["c"]
     ]
+
+
+def test_list_count_match_boolean_nulls_9141() -> None:
+    a = pl.DataFrame({"a": [[True, None, False]]})
+
+    assert a.select(pl.col("a").list.count_match(True))["a"].to_list() == [1]