fix(rust, python): fix boolean arg-max if all equal (#5680)

pola-rs · Nov 30, 2022 · ef53b62 · ef53b62
1 parent 264da15
commit ef53b62
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 4 deletions.
diff --git a/polars/polars-core/src/chunked_array/ops/aggregate.rs b/polars/polars-core/src/chunked_array/ops/aggregate.rs
@@ -848,12 +848,24 @@ where
 
 impl ArgAgg for BooleanChunked {
     fn arg_min(&self) -> Option<usize> {
-        self.into_iter()
-            .position(|opt_val| matches!(opt_val, Some(false)))
+        if self.is_empty() || self.null_count() == self.len() {
+            None
+        } else if self.all() {
+            Some(0)
+        } else {
+            self.into_iter()
+                .position(|opt_val| matches!(opt_val, Some(false)))
+        }
     }
     fn arg_max(&self) -> Option<usize> {
-        self.into_iter()
-            .position(|opt_val| matches!(opt_val, Some(true)))
+        if self.is_empty() || self.null_count() == self.len() {
+            None
+        } else if self.any() {
+            self.into_iter()
+                .position(|opt_val| matches!(opt_val, Some(true)))
+        } else {
+            Some(0)
+        }
     }
 }
 impl ArgAgg for Utf8Chunked {}

diff --git a/py-polars/tests/unit/test_series.py b/py-polars/tests/unit/test_series.py
@@ -1610,6 +1610,12 @@ def test_arg_min_and_arg_max() -> None:
     s = pl.Series([None, None], dtype=pl.Boolean)
     assert s.arg_min() is None
     assert s.arg_max() is None
+    s = pl.Series([True, True])
+    assert s.arg_min() == 0
+    assert s.arg_max() == 0
+    s = pl.Series([False, False])
+    assert s.arg_min() == 0
+    assert s.arg_max() == 0
 
 
 def test_is_null_is_not_null() -> None: