feat(rust, python): add arg_min/arg_max for series of dtype boolean (#…

…5592)
pola-rs · Nov 22, 2022 · 52487c2 · 52487c2
1 parent 664a851
commit 52487c2
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/polars/polars-core/src/chunked_array/ops/aggregate.rs b/polars/polars-core/src/chunked_array/ops/aggregate.rs
@@ -846,7 +846,16 @@ where
     }
 }
 
-impl ArgAgg for BooleanChunked {}
+impl ArgAgg for BooleanChunked {
+    fn arg_min(&self) -> Option<usize> {
+        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)))
+    }
+}
 impl ArgAgg for Utf8Chunked {}
 #[cfg(feature = "dtype-binary")]
 impl ArgAgg for BinaryChunked {}

diff --git a/py-polars/tests/unit/test_series.py b/py-polars/tests/unit/test_series.py
@@ -1604,6 +1604,13 @@ def test_arg_min_and_arg_max() -> None:
     assert s.arg_min() == 3
     assert s.arg_max() == 0
 
+    s = pl.Series([None, True, False, True])
+    assert s.arg_min() == 2
+    assert s.arg_max() == 1
+    s = pl.Series([None, None], dtype=pl.Boolean)
+    assert s.arg_min() is None
+    assert s.arg_max() is None
+
 
 def test_is_null_is_not_null() -> None:
     s = pl.Series("a", [1.0, 2.0, 3.0, None])