Address review comments

pandas-dev · jreback · Apr 13, 2021 · Apr 6, 2021 · Apr 6, 2021 · Apr 6, 2021
commit c9b9d5f25820523898970be1a746f752b8e0dc4b
diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py
@@ -77,11 +77,11 @@ class providing the base-class of operations.
 from pandas.core import nanops
 import pandas.core.algorithms as algorithms
 from pandas.core.arrays import (
+    BaseMaskedArray,
+    BooleanArray,
     Categorical,
     ExtensionArray,
 )
-from pandas.core.arrays.boolean import BooleanArray
-from pandas.core.arrays.masked import BaseMaskedArray
 from pandas.core.base import (
     DataError,
     PandasObject,
@@ -1444,9 +1444,9 @@ def result_to_bool(
             needs_mask=True,
             pre_processing=objs_to_bool,
             post_processing=result_to_bool,
+            use_nullable_input_arg=True,
             val_test=val_test,
             skipna=skipna,
-            masked=False,
         )
 
     @final
@@ -2630,6 +2630,7 @@ def _get_cythonized_result(
         result_is_index: bool = False,
         pre_processing=None,
         post_processing=None,
+        use_nullable_input_arg: bool = False,
         **kwargs,
     ):
         """
@@ -2677,6 +2678,9 @@ def _get_cythonized_result(
             second argument, i.e. the signature should be
             (ndarray, Type). Optionally, a third argument can be "masked", to
             allow for processing specific to nullable values
+        use_nullable_input_arg : bool, default False
+            If True, pass an argument "masked" to the cython function specifying
+            whether or not the input should be treated as masked
         **kwargs : dict
             Extra arguments to be passed back to Cython funcs
 
@@ -2705,7 +2709,6 @@ def _get_cythonized_result(
         post_processing_accepts_masked = post_processing is not None and (
             "masked" in inspect.signature(post_processing).parameters
         )
-        kwargs_accepts_masked = "masked" in kwargs
 
         error_msg = ""
         for idx, obj in enumerate(self._iterate_slices()):
@@ -2757,8 +2760,8 @@ def _get_cythonized_result(
             if needs_ngroups:
                 func = partial(func, ngroups)
 
-            if kwargs_accepts_masked:
-                kwargs["masked"] = is_nullable
+            if use_nullable_input_arg:
+                func = partial(func, masked=is_nullable)
 
             func(**kwargs)  # Call func to modify indexer values in place
 

diff --git a/pandas/tests/groupby/test_any_all.py b/pandas/tests/groupby/test_any_all.py
@@ -74,8 +74,9 @@ def test_bool_aggs_dup_column_labels(bool_agg_func):
 @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",
+    # expected_data indexed as [[skipna=False/any, skipna=False/all],
+    #                           [skipna=True/any, skipna=True/all]]
+    "data,expected_data",
     [
         ([False, False, False], [[False, False], [False, False]]),
         ([True, True, True], [[True, True], [True, True]]),
@@ -85,16 +86,22 @@ def test_bool_aggs_dup_column_labels(bool_agg_func):
         ([True, pd.NA, False], [[True, False], [True, False]]),
     ],
 )
-def test_masked_kleene_logic(bool_agg_func, data, expected, skipna):
+def test_masked_kleene_logic(bool_agg_func, data, expected_data, skipna):
     # GH#37506
     df = DataFrame(data, dtype="boolean")
     expected = DataFrame(
-        [expected[skipna][bool_agg_func == "all"]], dtype="boolean", index=[1]
+        [expected_data[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)
 
+    # The expected result we compared to should match aggregating on the whole
+    # series
+    result = getattr(df[0], bool_agg_func)(skipna=skipna)
+    expected = expected_data[skipna][bool_agg_func == "all"]
+    assert (result is pd.NA and expected is pd.NA) or result == expected
+
 
 @pytest.mark.parametrize(
     "dtype1,dtype2,exp_col1,exp_col2",