REF: avoid special-casing in SelectN (#45956)

Author: jbrockmendel

pandas/core/algorithms.py

@@ -64,6 +64,7 @@
 )
 from pandas.core.dtypes.concat import concat_compat
 from pandas.core.dtypes.dtypes import (
+    BaseMaskedDtype,
     ExtensionDtype,
     PandasDtype,
 )
@@ -103,6 +104,7 @@
         Series,
     )
     from pandas.core.arrays import (
+        BaseMaskedArray,
         DatetimeArray,
         ExtensionArray,
         TimedeltaArray,
@@ -142,6 +144,15 @@ def _ensure_data(values: ArrayLike) -> np.ndarray:
     if is_object_dtype(values.dtype):
         return ensure_object(np.asarray(values))
 
+    elif isinstance(values.dtype, BaseMaskedDtype):
+        # i.e. BooleanArray, FloatingArray, IntegerArray
+        values = cast("BaseMaskedArray", values)
+        if not values._hasna:
+            # No pd.NAs -> We can avoid an object-dtype cast (and copy) GH#41816
+            #  recurse to avoid re-implementing logic for eg bool->uint8
+            return _ensure_data(values._data)
+        return np.asarray(values)
+
     elif is_bool_dtype(values.dtype):
         if isinstance(values, np.ndarray):
             # i.e. actually dtype == np.dtype("bool")
@@ -1188,18 +1199,6 @@ def compute(self, method: str) -> Series:
         dropped = self.obj.dropna()
         nan_index = self.obj.drop(dropped.index)
 
-        if is_extension_array_dtype(dropped.dtype):
-            # GH#41816 bc we have dropped NAs above, MaskedArrays can use the
-            #  numpy logic.
-            from pandas.core.arrays import BaseMaskedArray
-
-            arr = dropped._values
-            if isinstance(arr, BaseMaskedArray):
-                ser = type(dropped)(arr._data, index=dropped.index, name=dropped.name)
-
-                result = type(self)(ser, n=self.n, keep=self.keep).compute(method)
-                return result.astype(arr.dtype)
-
         # slow method
         if n >= len(self.obj):
             ascending = method == "nsmallest"

pandas/core/arrays/masked.py

@@ -27,13 +27,9 @@
     SequenceIndexer,
     Shape,
     npt,
-    type_t,
 )
 from pandas.errors import AbstractMethodError
-from pandas.util._decorators import (
-    cache_readonly,
-    doc,
-)
+from pandas.util._decorators import doc
 from pandas.util._validators import validate_fillna_kwargs
 
 from pandas.core.dtypes.astype import astype_nansafe
@@ -51,6 +47,7 @@
     is_string_dtype,
     pandas_dtype,
 )
+from pandas.core.dtypes.dtypes import BaseMaskedDtype
 from pandas.core.dtypes.inference import is_array_like
 from pandas.core.dtypes.missing import (
     array_equivalent,
@@ -91,43 +88,6 @@
 BaseMaskedArrayT = TypeVar("BaseMaskedArrayT", bound="BaseMaskedArray")
 
 
-class BaseMaskedDtype(ExtensionDtype):
-    """
-    Base class for dtypes for BaseMaskedArray subclasses.
-    """
-
-    name: str
-    base = None
-    type: type
-
-    na_value = libmissing.NA
-
-    @cache_readonly
-    def numpy_dtype(self) -> np.dtype:
-        """Return an instance of our numpy dtype"""
-        return np.dtype(self.type)
-
-    @cache_readonly
-    def kind(self) -> str:
-        return self.numpy_dtype.kind
-
-    @cache_readonly
-    def itemsize(self) -> int:
-        """Return the number of bytes in this dtype"""
-        return self.numpy_dtype.itemsize
-
-    @classmethod
-    def construct_array_type(cls) -> type_t[BaseMaskedArray]:
-        """
-        Return the array type associated with this dtype.
-
-        Returns
-        -------
-        type
-        """
-        raise NotImplementedError
-
-
 class BaseMaskedArray(OpsMixin, ExtensionArray):
     """
     Base class for masked arrays (which use _data and _mask to store the data).

pandas/core/dtypes/dtypes.py

@@ -14,6 +14,7 @@
 import numpy as np
 import pytz
 
+from pandas._libs import missing as libmissing
 from pandas._libs.interval import Interval
 from pandas._libs.properties import cache_readonly
 from pandas._libs.tslibs import (
@@ -57,6 +58,7 @@
         Index,
     )
     from pandas.core.arrays import (
+        BaseMaskedArray,
         DatetimeArray,
         IntervalArray,
         PandasArray,
@@ -1376,3 +1378,40 @@ def itemsize(self) -> int:
         The element size of this data-type object.
         """
         return self._dtype.itemsize
+
+
+class BaseMaskedDtype(ExtensionDtype):
+    """
+    Base class for dtypes for BaseMaskedArray subclasses.
+    """
+
+    name: str
+    base = None
+    type: type
+
+    na_value = libmissing.NA
+
+    @cache_readonly
+    def numpy_dtype(self) -> np.dtype:
+        """Return an instance of our numpy dtype"""
+        return np.dtype(self.type)
+
+    @cache_readonly
+    def kind(self) -> str:
+        return self.numpy_dtype.kind
+
+    @cache_readonly
+    def itemsize(self) -> int:
+        """Return the number of bytes in this dtype"""
+        return self.numpy_dtype.itemsize
+
+    @classmethod
+    def construct_array_type(cls) -> type_t[BaseMaskedArray]:
+        """
+        Return the array type associated with this dtype.
+
+        Returns
+        -------
+        type
+        """
+        raise NotImplementedError