Compute mask usage inside algo

pandas-dev · jreback · Apr 21, 2021 · Mar 26, 2021 · Mar 26, 2021 · Mar 26, 2021
commit 0009dfd367d6bfb74381a72a3c3e4314b39edfcc
diff --git a/pandas/_libs/groupby.pyx b/pandas/_libs/groupby.pyx
@@ -1236,7 +1236,6 @@ def group_cummin_max(groupby_t[:, ::1] out,
                      const intp_t[:] labels,
                      int ngroups,
                      bint is_datetimelike,
-                     bint use_mask,
                      bint compute_max):
     """
     Cumulative minimum/maximum of columns of `values`, in row groups `labels`.
@@ -1256,10 +1255,6 @@ def group_cummin_max(groupby_t[:, ::1] out,
         Number of groups, larger than all entries of `labels`.
     is_datetimelike : bool
         True if `values` contains datetime-like entries.
-    use_mask : bool
-        True if the mask should be used (otherwise we continue
-        as if it is not a masked algorithm). Avoids the cost
-        of checking for a completely zeroed mask.
    
8000
 compute_max : bool
         True if cumulative maximum should be computed, False
         if cumulative minimum should be computed
@@ -1273,7 +1268,9 @@ def group_cummin_max(groupby_t[:, ::1] out,
         groupby_t val, mval
         ndarray[groupby_t, ndim=2] accum
         intp_t lab
-        bint val_is_nan
+        bint val_is_nan, use_mask
+
+    use_mask = mask is not None
 
     N, K = (<object>values).shape
     accum = np.empty((ngroups, K), dtype=np.asarray(values).dtype)
@@ -1331,8 +1328,7 @@ def group_cummin(groupby_t[:, ::1] out,
                  uint8_t[:, ::1] mask,
                  const intp_t[:] labels,
                  int ngroups,
-                 bint is_datetimelike,
-                 bint use_mask):
+                 bint is_datetimelike):
     """See group_cummin_max.__doc__"""
     group_cummin_max(
         out,
@@ -1341,7 +1337,6 @@ def group_cummin(groupby_t[:, ::1] out,
         labels,
         ngroups,
         is_datetimelike,
-        use_mask,
         compute_max=False
     )
 
@@ -1353,8 +1348,7 @@ def group_cummax(groupby_t[:, ::1] out,
                  uint8_t[:, ::1] mask,
                  const intp_t[:] labels,
                  int ngroups,
-                 bint is_datetimelike,
-                 bint use_mask):
+                 bint is_datetimelike):
     """See group_cummin_max.__doc__"""
        group_cummin_max(
         out,
@@ -1363,6 +1357,5 @@ def group_cummax(groupby_t[:, ::1] out,
         labels,
         ngroups,
         is_datetimelike,
-        use_mask,
         compute_max=True
     )
diff --git a/pandas/core/groupby/ops.py b/pandas/core/groupby/ops.py
@@ -652,6 +652,7 @@ def _cython_operation(
         # can we do this operation with our cython functions
         # if not raise NotImplementedError
         self._disallow_invalid_ops(dtype, how, is_numeric)
+
         func_uses_mask = cython_function_uses_mask(how)
         if is_extension_array_dtype(dtype):
             if isinstance(values, BaseMaskedArray) and func_uses_mask:
@@ -718,10 +719,6 @@ def _cython_operation(
             out_shape = (ngroups,) + values.shape[1:]
 
         func, values = self._get_cython_func_and_vals(kind, how, values, is_numeric)
-        use_mask = mask is not None
-        if func_uses_mask:
-            if mask is None:
-                mask = np.zeros_like(values, dtype=np.uint8, order="C")
 
         if how == "rank":
             out_dtype = "float"
@@ -740,10 +737,10 @@ def _cython_operation(
         elif kind == "transform":
             # TODO: min_count
             result = self._transform(
-                result, values, func, is_datetimelike, use_mask, mask, **kwargs
+                result, values, func, is_datetimelike, mask, func_uses_mask, **kwargs
             )
 
-        if not use_mask and is_integer_dtype(result.dtype) and not is_datetimelike:
+        if mask is None and is_integer_dtype(result.dtype) and not is_datetimelike:
             result_mask = result == iNaT
             if result_mask.any():
                 result = result.astype("float64")
@@ -784,21 +781,20 @@ def _transform(
         values: np.ndarray,
         transform_func,
         is_datetimelike: bool,
-        use_mask: bool,
         mask: np.ndarray | None,
+        needs_mask: bool = False,
         **kwargs,
     ) -> np.ndarray:
 
         comp_ids, _, ngroups = self.group_info
-        if mask is not None:
+        if needs_mask:
             transform_func(
                 result,
                 values,
                 mask,
                 comp_ids,
                 ngroups,
                 is_datetimelike,
-                use_mask,
                 **kwargs,
             )
         else: