update booelan aggregates and ensure they return IdxSize (#3563)

pola-rs · Jun 3, 2022 · e5f3173 · e5f3173
1 parent 322f1ed
commit e5f3173
Showing 1 changed file with 29 additions and 35 deletions.
diff --git a/polars/polars-core/src/chunked_array/ops/aggregate.rs b/polars/polars-core/src/chunked_array/ops/aggregate.rs
@@ -451,52 +451,46 @@ impl ChunkVar<Series> for ListChunked {}
 impl<T> ChunkVar<Series> for ObjectChunked<T> {}
 impl ChunkVar<bool> for BooleanChunked {}
 
-fn min_max_helper(ca: &BooleanChunked, min: bool) -> u32 {
-    ca.into_iter().fold(0, |acc: u32, x| match x {
-        Some(v) => {
-            let v = v as u32;
-            if min {
-                if acc < v {
-                    acc
-                } else {
-                    v
-                }
-            } else if acc > v {
-                acc
-            } else {
-                v
-            }
-        }
-        None => acc,
-    })
-}
-
 /// Booleans are casted to 1 or 0.
-impl ChunkAgg<u32> for BooleanChunked {
+impl ChunkAgg<IdxSize> for BooleanChunked {
     /// Returns `None` if the array is empty or only contains null values.
-    fn sum(&self) -> Option<u32> {
+    fn sum(&self) -> Option<IdxSize> {
         if self.is_empty() {
-            return None;
+            None
+        } else {
+            Some(
+                self.downcast_iter()
+                    .map(|arr| match arr.validity() {
+                        Some(validity) => {
+                            (arr.len() - (validity & arr.values()).null_count()) as IdxSize
+                        }
+                        None => (arr.len() - arr.values().null_count()) as IdxSize,
+                    })
+                    .sum(),
+            )
         }
-        let sum = self.into_iter().fold(0, |acc: u32, x| match x {
-            Some(v) => acc + v as u32,
-            None => acc,
-        });
-        Some(sum)
     }
 
-    fn min(&self) -> Option<u32> {
+    fn min(&self) -> Option<IdxSize> {
         if self.is_empty() {
             return None;
         }
-        Some(min_max_helper(self, true))
+        if self.all() {
+            Some(1)
+        } else {
+            Some(0)
+        }
     }
 
-    fn max(&self) -> Option<u32> {
+    fn max(&self) -> Option<IdxSize> {
         if self.is_empty() {
             return None;
         }
-        Some(min_max_helper(self, false))
+        if self.any() {
+            Some(1)
+        } else {
+            Some(0)
+        }
     }
 }
 
@@ -739,19 +733,19 @@ impl QuantileAggSeries for Utf8Chunked {
 impl ChunkAggSeries for BooleanChunked {
     fn sum_as_series(&self) -> Series {
         let v = ChunkAgg::sum(self);
-        let mut ca: UInt32Chunked = [v].iter().copied().collect();
+        let mut ca: IdxCa = [v].iter().copied().collect();
         ca.rename(self.name());
         ca.into_series()
     }
     fn max_as_series(&self) -> Series {
         let v = ChunkAgg::max(self);
-        let mut ca: UInt32Chunked = [v].iter().copied().collect();
+        let mut ca: IdxCa = [v].iter().copied().collect();
         ca.rename(self.name());
         ca.into_series()
     }
     fn min_as_series(&self) -> Series {
         let v = ChunkAgg::min(self);
-        let mut ca: UInt32Chunked = [v].iter().copied().collect();
+        let mut ca: IdxCa = [v].iter().copied().collect();
         ca.rename(self.name());
         ca.into_series()
     }