Avoid _maybe_mask_result

rhshadrach · Oct 24, 2023 · 75638a5 · 75638a5
1 parent 5c416fe
commit 75638a5
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/pandas/core/arrays/masked.py b/pandas/core/arrays/masked.py
@@ -1506,9 +1506,13 @@ def _groupby_op(
             arity = op._cython_arity.get(op.how, 1)
             result_mask = np.tile(result_mask, (arity, 1)).T
 
-        # res_values should already have the correct dtype, we just need to
-        #  wrap in a MaskedArray
-        return self._maybe_mask_result(res_values, result_mask)
+        if op.how in ["idxmin", "idxmax"]:
+            # Result values are indexes to take, keep as ndarray
+            return res_values
+        else:
+            # res_values should already have the correct dtype, we just need to
+            #  wrap in a MaskedArray
+            return self._maybe_mask_result(res_values, result_mask)
 
 
 def transpose_homogeneous_masked_arrays(

diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py
@@ -5845,8 +5845,6 @@ def _wrap_idxmax_idxmin(self, res: NDFrameT) -> NDFrameT:
             if isinstance(index, MultiIndex):
                 index = index.to_flat_index()
             values = res._values
-            if isinstance(values, ExtensionArray):
-                values = values.to_numpy(copy=False)
             na_value = na_value_for_dtype(index.dtype, compat=False)
             if isinstance(res, Series):
                 # mypy: expression has type "Series", variable has type "NDFrameT"