BUG: Pyarrow timestamp support for map() function

doc/source/whatsnew/v3.0.0.rst

@@ -853,6 +853,7 @@ Other
 - Bug in :meth:`Series.diff` allowing non-integer values for the ``periods`` argument. (:issue:`56607`)
 - Bug in :meth:`Series.dt` methods in :class:`ArrowDtype` that were returning incorrect values. (:issue:`57355`)
 - Bug in :meth:`Series.isin` raising ``TypeError`` when series is large (>10**6) and ``values`` contains NA (:issue:`60678`)
+- Bug in :meth:`Series.map` where mapping with a ``dict`` failed to match keys when the Series used ``timestamp[ns][pyarrow]`` dtype. (:issue:`61231`)
 - Bug in :meth:`Series.mode` where an exception was raised when taking the mode with nullable types with no null values in the series. (:issue:`58926`)
 - Bug in :meth:`Series.rank` that doesn't preserve missing values for nullable integers when ``na_option='keep'``. (:issue:`56976`)
 - Bug in :meth:`Series.replace` and :meth:`DataFrame.replace` inconsistently replacing matching instances when ``regex=True`` and missing values are present. (:issue:`56599`)

pandas/core/arrays/arrow/array.py

@@ -1481,7 +1481,7 @@ def to_numpy(
         return result
 
     def map(self, mapper, na_action: Literal["ignore"] | None = None):
-        if is_numeric_dtype(self.dtype):
+        if is_numeric_dtype(self.dtype) or self.dtype.kind in "mM":
             return map_array(self.to_numpy(), mapper, na_action=na_action)
         else:
             return super().map(mapper, na_action)

pandas/tests/extension/test_arrow.py

@@ -46,6 +46,7 @@
     pa_version_under20p0,
 )
 
+from pandas.core.dtypes.common import is_timedelta64_dtype
 from pandas.core.dtypes.dtypes import (
     ArrowDtype,
     CategoricalDtypeType,
@@ -281,9 +282,34 @@ 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)
-            tm.assert_numpy_array_equal(result, expected)
+            mapped = data_missing.map(lambda x: x, na_action=na_action)
+            result = pd.Series(mapped)
+            expected = pd.Series(data_missing.to_numpy())
+
+            orig_dtype = expected.dtype
+
+            if result.dtype == "float64" and (
+                is_datetime64_any_dtype(orig_dtype)
+                or is_timedelta64_dtype(orig_dtype)
+                or isinstance(orig_dtype, pd.DatetimeTZDtype)
+            ):
+                result = result.astype(orig_dtype)
+
+            if isinstance(orig_dtype, pd.DatetimeTZDtype):
+                pass
+            elif is_datetime64_any_dtype(orig_dtype):
+                result = result.astype("datetime64[ns]").astype("int64")
+                expected = expected.astype("datetime64[ns]").astype("int64")
+                result = pd.Series(result)
+                expected = pd.Series(expected)
+            elif is_timedelta64_dtype(orig_dtype):
+                result = result.astype("timedelta64[ns]")
+                expected = expected.astype("timedelta64[ns]")
+
+            tm.assert_series_equal(
+                result, expected, check_dtype=False, check_exact=False
+            )
+
         else:
             result = data_missing.map(lambda x: x, na_action=na_action)
             if data_missing.dtype == "float32[pyarrow]":

pandas/tests/series/methods/test_map.py

@@ -606,6 +606,17 @@ def test_map_kwargs():
     tm.assert_series_equal(result, expected)
 
 
+def test_map_arrow_timestamp_dict():
+    # GH 61231
+    pytest.importorskip("pyarrow")
+
+    ser = Series(date_range("2023-01-01", periods=3)).astype("timestamp[ns][pyarrow]")
+    mapper = {ts: i for i, ts in enumerate(ser)}
+    result = ser.map(mapper)
+    expected = Series([0, 1, 2], dtype="int64")
+    tm.assert_series_equal(result, expected)
+
+
 def test_map_arg_as_kwarg():
     with tm.assert_produces_warning(
         FutureWarning, match="`arg` has been renamed to `func`"