Force mask usage

Author: mzeitlin11

pandas/_libs/groupby.pyx

@@ -1256,8 +1256,7 @@ cdef group_cummin_max(groupby_t[:, ::1] out,
     values : np.ndarray[groupby_t, ndim=2]
         Values to take cummin/max of.
     mask : array[uint8_t] or None
-        If not None, indices represent missing values,
-        otherwise the mask will not be used
+        Indices representing missing values,
     labels : np.ndarray[np.intp]
         Labels to group by.
     ngroups : int
@@ -1277,9 +1276,7 @@ cdef group_cummin_max(groupby_t[:, ::1] out,
         groupby_t val, mval
         groupby_t[:, ::1] accum
         intp_t lab
-        bint val_is_nan, use_mask
-
-    use_mask = mask is not None
+        bint val_is_nan
 
     N, K = (<object>values).shape
     accum = np.empty((ngroups, K), dtype=np.asarray(values).dtype, order='C')
@@ -1293,33 +1290,12 @@ cdef group_cummin_max(groupby_t[:, ::1] out,
     with nogil:
         for i in range(N):
             lab = labels[i]
-
             if lab < 0:
                 continue
-            for j in range(K):
-                val_is_nan = False
-
-                if use_mask:
-                    if mask[i, j]:
-
-                        # `out` does not need to be set since it
-                        # will be masked anyway
-                        val_is_nan = True
-                    else:
-
-                        # If using the mask, we can avoid grabbing the
-                        # value unless necessary
-                        val = values[i, j]
 
-                # Otherwise, `out` must be set accordingly if the
-                # value is missing
-                else:
+            for j in range(K):
+                if not mask[i, j]:
                     val = values[i, j]
-                    if _treat_as_na(val, is_datetimelike):
-                        val_is_nan = True
-                        out[i, j] = val
-
-                if not val_is_nan:
                     mval = accum[lab, j]
                     if compute_max:
                         if val > mval:

pandas/core/groupby/ops.py

@@ -575,7 +575,14 @@ def _is_builtin_func(self, arg):
 
     @final
     def _ea_wrap_cython_operation(
-        self, kind: str, values, how: str, axis: int, min_count: int = -1, **kwargs
+        self,
+        kind: str,
+        values,
+        how: str,
+        axis: int,
+        min_count: int = -1,
+        mask: np.ndarray | None = None,
+        **kwargs,
     ) -> ArrayLike:
         """
         If we have an ExtensionArray, unwrap, call _cython_operation, and
@@ -589,7 +596,7 @@ def _ea_wrap_cython_operation(
             #  operate on the tz-naive equivalents
             values = values.view("M8[ns]")
             res_values = self._cython_operation(
-                kind, values, how, axis, min_count, **kwargs
+                kind, values, how, axis, min_count, mask=mask, **kwargs
             )
             if how in ["rank"]:
                 # preserve float64 dtype
@@ -603,7 +610,7 @@ def _ea_wrap_cython_operation(
             # IntegerArray or BooleanArray
             values = ensure_int_or_float(values)
             res_values = self._cython_operation(
-                kind, values, how, axis, min_count, **kwargs
+                kind, values, how, axis, min_count, mask=mask, **kwargs
             )
             dtype = maybe_cast_result_dtype(orig_values.dtype, how)
             if isinstance(dtype, ExtensionDtype):
@@ -616,7 +623,7 @@ def _ea_wrap_cython_operation(
             # FloatingArray
             values = values.to_numpy(values.dtype.numpy_dtype, na_value=np.nan)
             res_values = self._cython_operation(
-                kind, values, how, axis, min_count, **kwargs
+                kind, values, how, axis, min_count, mask=mask, **kwargs
             )
             result = type(orig_values)._from_sequence(res_values)
             return result
@@ -632,7 +639,8 @@ def _masked_ea_wrap_cython_operation(
         values: BaseMaskedArray,
         how: str,
         axis: int,
-        min_count: int = -1,
+        min_count: int,
+        mask: np.ndarray,
         **kwargs,
     ) -> BaseMaskedArray:
         """
@@ -641,9 +649,6 @@ def _masked_ea_wrap_cython_operation(
         """
         orig_values = values
 
-        # isna just directly returns self._mask, so copy here to prevent
-        # modifying the original
-        mask = isna(values).copy()
         arr = values._data
 
         if is_integer_dtype(values.dtype) or is_bool_dtype(values.dtype):
@@ -658,7 +663,7 @@ def _masked_ea_wrap_cython_operation(
         cls = dtype.construct_array_type()
 
         return cls(
-            res_values.astype(dtype.type, copy=False), mask.astype(bool, copy=False)
+            res_values.astype(dtype.type, copy=False), mask.astype(bool, copy=True)
         )
 
     @final
@@ -695,14 +700,20 @@ def _cython_operation(
         cy_op.disallow_invalid_ops(dtype, is_numeric)
 
         func_uses_mask = cy_op.uses_mask()
+
+        # Only compute the mask if we haven't yet
+        if func_uses_mask and mask is None:
+            mask = isna(values)
+
         if is_extension_array_dtype(dtype):
             if isinstance(values, BaseMaskedArray) and func_uses_mask:
+                assert mask is not None
                 return self._masked_ea_wrap_cython_operation(
-                    kind, values, how, axis, min_count, **kwargs
+                    kind, values, how, axis, min_count, mask=mask, **kwargs
                 )
             else:
                 return self._ea_wrap_cython_operation(
-                    kind, values, how, axis, min_count, **kwargs
+                    kind, values, how, axis, min_count, mask=mask, **kwargs
                 )
 
         elif values.ndim == 1: