diff --git a/pandas/core/arrays/arrow/array.py b/pandas/core/arrays/arrow/array.py index ad1d576bfec32..1bb738da73b84 100644 --- a/pandas/core/arrays/arrow/array.py +++ b/pandas/core/arrays/arrow/array.py @@ -1511,10 +1511,9 @@ def to_numpy( return result def map(self, mapper, na_action: Literal["ignore"] | None = None): - if is_numeric_dtype(self.dtype): + if self.dtype.kind in "mM" or is_numeric_dtype(self.dtype): return map_array(self.to_numpy(), mapper, na_action=na_action) - else: - return super().map(mapper, na_action) + return super().map(mapper, na_action) @doc(ExtensionArray.duplicated) def duplicated( diff --git a/pandas/tests/extension/test_arrow.py b/pandas/tests/extension/test_arrow.py index c3e1d33ec93df..52ab085e07dfd 100644 --- a/pandas/tests/extension/test_arrow.py +++ b/pandas/tests/extension/test_arrow.py @@ -64,6 +64,7 @@ is_string_dtype, is_unsigned_integer_dtype, ) +from pandas.core.algorithms import map_array from pandas.tests.extension import base pa = pytest.importorskip("pyarrow") @@ -278,8 +279,9 @@ def test_compare_scalar(self, data, comparison_op): @pytest.mark.parametrize("na_action", [None, "ignore"]) def test_map(self, data_missing, na_action): if data_missing.dtype.kind in "mM": - result = data_missing.map(lambda x: x, na_action=na_action) - expected = data_missing.to_numpy(dtype=object) + mapper = lambda x: x + result = data_missing.map(mapper, na_action=na_action) + expected = map_array(data_missing.to_numpy(), mapper, na_action=na_action) tm.assert_numpy_array_equal(result, expected) else: result = data_missing.map(lambda x: x, na_action=na_action) @@ -3579,3 +3581,16 @@ def test_timestamp_dtype_matches_to_datetime(): expected = pd.Series([ts], dtype=dtype1).convert_dtypes(dtype_backend="pyarrow") tm.assert_series_equal(result, expected) + + +@pytest.mark.parametrize("dtype", ["timestamp[ns][pyarrow]", "datetime64[ns]"]) +def test_map_timestamp(dtype): + # GH#61231 + date_range = pd.date_range("2018-01-01", "2018-01-07") + df = pd.DataFrame({"a": date_range}).astype({"a": dtype}) + date2pos = {date: i for i, date in enumerate(df["a"])} + + result = df["a"].map(date2pos) + expected = pd.Series(range(len(date_range)), name="a") + + tm.assert_series_equal(result, expected, check_dtype=False)