fix: Fix boolean distinct (#16765)

pola-rs · Jun 6, 2024 · 5a0c803 · 5a0c803
1 parent 27abc9d
commit 5a0c803
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/crates/polars-compute/src/distinct_count.rs b/crates/polars-compute/src/distinct_count.rs
@@ -12,7 +12,13 @@ impl DistinctCountKernel for BooleanArray {
             return 0;
         }
 
-        let unset_bits = self.values().unset_bits();
-        2 - usize::from(unset_bits == 0 || unset_bits == self.values().len())
+        if self.null_count() == 0 {
+            let unset_bits = self.values().unset_bits();
+            2 - usize::from(unset_bits == 0 || unset_bits == self.values().len())
+        } else {
+            let values = self.values() & self.validity().unwrap();
+            let unset_bits = self.values().unset_bits();
+            3 - usize::from(unset_bits == 0 || unset_bits == values.len())
+        }
     }
 }
diff --git a/crates/polars-parquet/src/arrow/write/boolean/basic.rs b/crates/polars-parquet/src/arrow/write/boolean/basic.rs
@@ -91,7 +91,11 @@ pub(super) fn build_statistics(
         null_count: options.null_count.then(|| array.null_count() as i64),
         distinct_count: options
             .distinct_count
-            .then(|| array.distinct_count().try_into().ok())
+            .then(|| {
+                (array.distinct_count() - ((array.null_count() > 0) as usize))
+                    .try_into()
+                    .ok()
+            })
             .flatten(),
         max_value: options
             .max_value