ENH/BUG: Use Kleene logic for groupby any/all

asv_bench/benchmarks/groupby.py

@@ -480,7 +480,19 @@ class GroupByCythonAgg:
     param_names = ["dtype", "method"]
     params = [
         ["float64"],
-        ["sum", "prod", "min", "max", "mean", "median", "var", "first", "last"],
+        [
+            "sum",
+            "prod",
+            "min",
+            "max",
+            "mean",
+            "median",
+            "var",
+            "first",
+            "last",
+            "any",
+            "all",
+        ],
     ]
 
     def setup(self, dtype, method):

doc/source/whatsnew/v1.3.0.rst

@@ -195,6 +195,8 @@ Other enhancements
 - :class:`RangeIndex` can now be constructed by passing a ``range`` object directly e.g. ``pd.RangeIndex(range(3))`` (:issue:`12067`)
 - :meth:`round` being enabled for the nullable integer and floating dtypes (:issue:`38844`)
 - :meth:`pandas.read_csv` and :meth:`pandas.read_json` expose the argument ``encoding_errors`` to control how encoding errors are handled (:issue:`39450`)
+- :class:`SeriesGroupBy` and :class:`DataFrameGroupBy` return a ``BooleanDType`` result using Kleene logic for methods ``any`` and ``all`` with nullable data types (:issue:`37506`, :issue:`40585`, :issue:`33449`)
+-
 
 .. ---------------------------------------------------------------------------
 

pandas/_libs/groupby.pyx

@@ -393,9 +393,11 @@ def group_any_all(uint8_t[::1] out,
                   const intp_t[:] labels,
                   const uint8_t[::1] mask,
                   str val_test,
-                  bint skipna) -> None:
+                  bint skipna,
+                  bint use_kleene_logic = False) -> None:
     """
-    Aggregated boolean values to show truthfulness of group elements.
+    Aggregated boolean values to show truthfulness of group elements. If the
+    input is a nullable type, Kleene logic will be used.
 
     Parameters
     ----------
@@ -412,11 +414,14 @@ def group_any_all(uint8_t[::1] out,
         String object dictating whether to use any or all truth testing
     skipna : bool
         Flag to ignore nan values during truth testing
+    use_kleene_logic : bool, default False
+        Whether or not to compute the result using Kleene logic
 
     Notes
     -----
     This method modifies the `out` parameter rather than returning an object.
-    The returned values will either be 0 or 1 (False or True, respectively).
+    The returned values will either be 0, 1 (False or True, respectively), or
+    2 to signify a masked position in the case of a nullable input.
     """
     cdef:
         Py_ssize_t i, N = len(labels)
@@ -444,6 +449,16 @@ def group_any_all(uint8_t[::1] out,
             if lab < 0 or (skipna and mask[i]):
                 continue
 
+            if use_kleene_logic and mask[i]:
+                # Set the position as masked if `out[lab] != flag_val`, which
+                # would indicate True/False has not yet been seen for any/all,
+                # so by Kleene logic the result is currently unknown
+                if out[lab] != flag_val:
+                    out[lab] = 2
+                continue
+
+            # If True and 'any' or False and 'all', the result is
+            # already determined
             if values[i] == flag_val:
                 out[lab] = flag_val
 

pandas/core/groupby/groupby.py

@@ -80,6 +80,8 @@ class providing the base-class of operations.
     Categorical,
     ExtensionArray,
 )
+from pandas.core.arrays.boolean import BooleanArray
+from pandas.core.arrays.masked import BaseMaskedArray
 from pandas.core.base import (
     DataError,
     PandasObject,
@@ -1413,16 +1415,25 @@ def _bool_agg(self, val_test, skipna):
         Shared func to call any / all Cython GroupBy implementations.
         """
 
-        def objs_to_bool(vals: np.ndarray) -> tuple[np.ndarray, type]:
+        def objs_to_bool(vals: ArrayLike) -> tuple[np.ndarray, type]:
             if is_object_dtype(vals):
                 vals = np.array([bool(x) for x in vals])
+            elif isinstance(vals, BaseMaskedArray):
+                vals = vals._data.astype(bool, copy=False)
             else:
                 vals = vals.astype(bool)
 
             return vals.view(np.uint8), bool
 
-        def result_to_bool(result: np.ndarray, inference: type) -> np.ndarray:
-            return result.astype(inference, copy=False)
+        def result_to_bool(
+            result: np.ndarray,
+            inference: type,
+            result_is_nullable: bool = False,
+        ) -> ArrayLike:
+            if result_is_nullable:
+                return BooleanArray(result.astype(bool, copy=False), result == 2)
+            else:
+                return result.astype(inference, copy=False)
 
         return self._get_cythonized_result(
             "group_any_all",
@@ -2735,6 +2746,10 @@ def _get_cythonized_result(
                 mask = isna(values).view(np.uint8)
                 func = partial(func, mask)
 
+            if how == "group_any_all" and isinstance(values, BaseMaskedArray):
+                post_processing = partial(post_processing, result_is_nullable=True)
+                func = partial(func, use_kleene_logic=True)
+
             if needs_ngroups:
                 func = partial(func, ngroups)
 

pandas/tests/groupby/test_any_all.py

@@ -3,6 +3,7 @@
 import numpy as np
 import pytest
 
+import pandas as pd
 from pandas import (
     DataFrame,
     Index,
@@ -68,3 +69,43 @@ def test_bool_aggs_dup_column_labels(bool_agg_func):
 
     expected = df
     tm.assert_frame_equal(result, expected)
+
+
+@pytest.mark.parametrize("bool_agg_func", ["any", "all"])
+@pytest.mark.parametrize("skipna", [True, False])
+@pytest.mark.parametrize(
+    # expected indexed as [skipna][bool_agg_func == "all"]
+    "data,expected",
+    [
+        ([False, False, False], [[False, False], [False, False]]),
+        ([True, True, True], [[True, True], [True, True]]),
+        ([pd.NA, pd.NA, pd.NA], [[pd.NA, pd.NA], [False, True]]),
+        ([False, pd.NA, False], [[pd.NA, False], [False, False]]),
+        ([True, pd.NA, True], [[True, pd.NA], [True, True]]),
+        ([True, pd.NA, False], [[True, False], [True, False]]),
+    ],
+)
+def test_masked_kleene_logic(bool_agg_func, data, expected, skipna):
+    # GH#37506
+    df = DataFrame(data, dtype="boolean")
+    expected = DataFrame(
+        [expected[skipna][bool_agg_func == "all"]], dtype="boolean", index=[1]
+    )
+
+    result = df.groupby([1, 1, 1]).agg(bool_agg_func, skipna=skipna)
+    tm.assert_frame_equal(result, expected)
+
+
+@pytest.mark.parametrize("bool_agg_func", ["any", "all"])
+@pytest.mark.parametrize("dtype", ["Int64", "Float64", "boolean"])
+@pytest.mark.parametrize("skipna", [True, False])
+def test_masked_bool_aggs_skipna(bool_agg_func, dtype, skipna, frame_or_series):
+    # GH#40585
+    obj = frame_or_series([pd.NA, 1], dtype=dtype)
+    expected_res = True
+    if not skipna and bool_agg_func == "all":
+        expected_res = pd.NA
+    expected = frame_or_series([expected_res], index=[1], dtype="boolean")
+
+    result = obj.groupby([1, 1]).agg(bool_agg_func, skipna=skipna)
+    tm.assert_equal(result, expected)

pandas/tests/reductions/test_reductions.py

@@ -949,14 +949,17 @@ def test_all_any_boolean(self):
         assert s3.all(skipna=False)
         assert not s4.any(skipna=False)
 
-        # Check level TODO(GH-33449) result should also be boolean
         s = Series(
             [False, False, True, True, False, True],
             index=[0, 0, 1, 1, 2, 2],
             dtype="boolean",
         )
-        tm.assert_series_equal(s.all(level=0), Series([False, True, False]))
-        tm.assert_series_equal(s.any(level=0), Series([False, True, True]))
+        tm.assert_series_equal(
+            s.all(level=0), Series([False, True, False], dtype="boolean")
+        )
+        tm.assert_series_equal(
+            s.any(level=0), Series([False, True, True], dtype="boolean")
+        )
 
     def test_any_axis1_bool_only(self):
         # GH#32432