REF: ExtensionIndex.searchsorted -> IndexOpsMixin.searchsorted (#43653)

Author: jbrockmendel

pandas/core/algorithms.py

@@ -96,6 +96,7 @@
     )
     from pandas.core.arrays import (
         DatetimeArray,
+        ExtensionArray,
         TimedeltaArray,
     )
 
@@ -1535,7 +1536,7 @@ def take(
 
 def searchsorted(
     arr: ArrayLike,
-    value: NumpyValueArrayLike,
+    value: NumpyValueArrayLike | ExtensionArray,
     side: Literal["left", "right"] = "left",
     sorter: NumpySorter = None,
 ) -> npt.NDArray[np.intp] | np.intp:
@@ -1616,7 +1617,9 @@ def searchsorted(
         # and `value` is a pd.Timestamp, we may need to convert value
         arr = ensure_wrapped_if_datetimelike(arr)
 
-    return arr.searchsorted(value, side=side, sorter=sorter)
+    # Argument 1 to "searchsorted" of "ndarray" has incompatible type
+    # "Union[NumpyValueArrayLike, ExtensionArray]"; expected "NumpyValueArrayLike"
+    return arr.searchsorted(value, side=side, sorter=sorter)  # type: ignore[arg-type]
 
 
 # ---- #

pandas/core/base.py

@@ -14,6 +14,7 @@
     TypeVar,
     cast,
     final,
+    overload,
 )
 
 import numpy as np
@@ -1230,14 +1231,50 @@ def factorize(self, sort: bool = False, na_sentinel: int | None = -1):
         0  # wrong result, correct would be 1
         """
 
+    # This overload is needed so that the call to searchsorted in
+    # pandas.core.resample.TimeGrouper._get_period_bins picks the correct result
+
+    @overload
+    # The following ignore is also present in numpy/__init__.pyi
+    # Possibly a mypy bug??
+    # error: Overloaded function signatures 1 and 2 overlap with incompatible
+    # return types  [misc]
+    def searchsorted(  # type: ignore[misc]
+        self,
+        value: npt._ScalarLike_co,
+        side: Literal["left", "right"] = "left",
+        sorter: NumpySorter = None,
+    ) -> np.intp:
+        ...
+
+    @overload
+    def searchsorted(
+        self,
+        value: npt.ArrayLike | ExtensionArray,
+        side: Literal["left", "right"] = "left",
+        sorter: NumpySorter = None,
+    ) -> npt.NDArray[np.intp]:
+        ...
+
     @doc(_shared_docs["searchsorted"], klass="Index")
     def searchsorted(
         self,
-        value: NumpyValueArrayLike,
+        value: NumpyValueArrayLike | ExtensionArray,
         side: Literal["left", "right"] = "left",
         sorter: NumpySorter = None,
     ) -> npt.NDArray[np.intp] | np.intp:
-        return algorithms.searchsorted(self._values, value, side=side, sorter=sorter)
+
+        values = self._values
+        if not isinstance(values, np.ndarray):
+            # Going through EA.searchsorted directly improves performance GH#38083
+            return values.searchsorted(value, side=side, sorter=sorter)
+
+        return algorithms.searchsorted(
+            values,
+            value,
+            side=side,
+            sorter=sorter,
+        )
 
     def drop_duplicates(self, keep="first"):
         duplicated = self._duplicated(keep=keep)

pandas/core/indexes/extension.py

@@ -4,11 +4,8 @@
 from __future__ import annotations
 
 from typing import (
-    TYPE_CHECKING,
     Hashable,
-    Literal,
     TypeVar,
-    overload,
 )
 
 import numpy as np
@@ -38,17 +35,9 @@
     TimedeltaArray,
 )
 from pandas.core.arrays._mixins import NDArrayBackedExtensionArray
-from pandas.core.arrays.base import ExtensionArray
 from pandas.core.indexers import deprecate_ndim_indexing
 from pandas.core.indexes.base import Index
 
-if TYPE_CHECKING:
-
-    from pandas._typing import (
-        NumpySorter,
-        NumpyValueArrayLike,
-    )
-
 _T = TypeVar("_T", bound="NDArrayBackedExtensionIndex")
 
 
@@ -207,40 +196,6 @@ def __getitem__(self, key):
         deprecate_ndim_indexing(result)
         return result
 
-    # This overload is needed so that the call to searchsorted in
-    # pandas.core.resample.TimeGrouper._get_period_bins picks the correct result
-
-    @overload
-    # The following ignore is also present in numpy/__init__.pyi
-    # Possibly a mypy bug??
-    # error: Overloaded function signatures 1 and 2 overlap with incompatible
-    # return types  [misc]
-    def searchsorted(  # type: ignore[misc]
-        self,
-        value: npt._ScalarLike_co,
-        side: Literal["left", "right"] = "left",
-        sorter: NumpySorter = None,
-    ) -> np.intp:
-        ...
-
-    @overload
-    def searchsorted(
-        self,
-        value: npt.ArrayLike | ExtensionArray,
-        side: Literal["left", "right"] = "left",
-        sorter: NumpySorter = None,
-    ) -> npt.NDArray[np.intp]:
-        ...
-
-    def searchsorted(
-        self,
-        value: NumpyValueArrayLike | ExtensionArray,
-        side: Literal["left", "right"] = "left",
-        sorter: NumpySorter = None,
-    ) -> npt.NDArray[np.intp] | np.intp:
-        # overriding IndexOpsMixin improves performance GH#38083
-        return self._data.searchsorted(value, side=side, sorter=sorter)
-
     # ---------------------------------------------------------------------
 
     def delete(self, loc):

pandas/core/series.py

@@ -2790,13 +2790,14 @@ def __rmatmul__(self, other):
         return self.dot(np.transpose(other))
 
     @doc(base.IndexOpsMixin.searchsorted, klass="Series")
-    def searchsorted(
+    # Signature of "searchsorted" incompatible with supertype "IndexOpsMixin"
+    def searchsorted(  # type: ignore[override]
         self,
-        value: NumpyValueArrayLike,
+        value: NumpyValueArrayLike | ExtensionArray,
         side: Literal["left", "right"] = "left",
         sorter: NumpySorter = None,
     ) -> npt.NDArray[np.intp] | np.intp:
-        return algorithms.searchsorted(self._values, value, side=side, sorter=sorter)
+        return base.IndexOpsMixin.searchsorted(self, value, side=side, sorter=sorter)
 
     # -------------------------------------------------------------------
     # Combination