Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Sep 29, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions pandas/core/algorithms.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,7 @@
)
from pandas.core.arrays import (
DatetimeArray,
ExtensionArray,
TimedeltaArray,
)

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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]


# ---- #
Expand Down
41 changes: 39 additions & 2 deletions pandas/core/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
TypeVar,
cast,
final,
overload,
)

import numpy as np
Expand Down Expand Up @@ -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)
Expand Down
45 changes: 0 additions & 45 deletions pandas/core/indexes/extension.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,8 @@
from __future__ import annotations

from typing import (
TYPE_CHECKING,
Hashable,
Literal,
TypeVar,
overload,
)

import numpy as np
Expand Down Expand Up @@ -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")


Expand Down Expand Up @@ -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 _get_engine_target(self) -> np.ndarray:
Expand Down
7 changes: 4 additions & 3 deletions pandas/core/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down