Skip to content

REF: move ExtensionIndex.map to be part of DatetimeLikeArrayMixin.map #51934

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 1 commit into from
Mar 14, 2023
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
25 changes: 19 additions & 6 deletions pandas/core/arrays/datetimelike.py
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,7 @@
from pandas.core.algorithms import (
checked_add_with_arr,
isin,
map_array,
unique1d,
)
from pandas.core.array_algos import datetimelike_accumulations
Expand Down Expand Up @@ -754,14 +755,26 @@ def map(self, mapper, na_action=None):
if na_action is not None:
raise NotImplementedError

# TODO(GH-23179): Add ExtensionArray.map
# Need to figure out if we want ExtensionArray.map first.
# If so, then we can refactor IndexOpsMixin._map_values to
# a standalone function and call from here..
# Else, just rewrite _map_infer_values to do the right thing.
from pandas import Index

return Index(self).map(mapper).array
idx = Index(self)
try:
result = mapper(idx)

# Try to use this result if we can
if isinstance(result, np.ndarray):
result = Index(result)

if not isinstance(result, Index):
raise TypeError("The map function must return an Index object")
except Exception:
result = map_array(self, mapper)
result = Index(result)

if isinstance(result, ABCMultiIndex):
return result.to_numpy()
else:
return result.array

def isin(self, values) -> npt.NDArray[np.bool_]:
"""
Expand Down
6 changes: 2 additions & 4 deletions pandas/core/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -881,10 +881,8 @@ def _map_values(self, mapper, na_action=None):
"""
arr = extract_array(self, extract_numpy=True, extract_range=True)

if is_extension_array_dtype(arr.dtype):
# Item "IndexOpsMixin" of "Union[IndexOpsMixin, ExtensionArray,
# ndarray[Any, Any]]" has no attribute "map"
return arr.map(mapper, na_action=na_action) # type: ignore[union-attr]
if isinstance(arr, ExtensionArray):
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

DatetimeArray/DatetimeIndex etc. have numpy dtypes, so using is_extension_array_dtype isn't enough, we need to check for ExtensionArray to reach the .map method of datetime-likes.

return arr.map(mapper, na_action=na_action)

# Argument 1 to "map_array" has incompatible type
# "Union[IndexOpsMixin, ExtensionArray, ndarray[Any, Any]]";
Expand Down
26 changes: 3 additions & 23 deletions pandas/core/indexes/extension.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,18 +9,15 @@
TypeVar,
)

import numpy as np

from pandas.util._decorators import (
cache_readonly,
doc,
)
from pandas.util._decorators import cache_readonly

from pandas.core.dtypes.generic import ABCDataFrame

from pandas.core.indexes.base import Index

if TYPE_CHECKING:
import numpy as np

from pandas._typing import (
ArrayLike,
npt,
Expand Down Expand Up @@ -154,23 +151,6 @@ def _validate_fill_value(self, value):
"""
return self._data._validate_setitem_value(value)

@doc(Index.map)
def map(self, mapper, na_action=None):
# Try to run function on index first, and then on elements of index
# Especially important for group-by functionality
try:
result = mapper(self)

# Try to use this result if we can
if isinstance(result, np.ndarray):
result = Index(result)

if not isinstance(result, Index):
raise TypeError("The map function must return an Index object")
return result
except Exception:
return self.astype(object).map(mapper)

@cache_readonly
def _isnan(self) -> npt.NDArray[np.bool_]:
# error: Incompatible return value type (got "ExtensionArray", expected
Expand Down