moved logic from algorithms.py to ArrowExtensionArray.map

Author: arthurlw

pandas/core/algorithms.py

@@ -24,7 +24,6 @@
     lib,
 )
 from pandas._libs.missing import NA
-from pandas._libs.tslibs.timestamps import Timestamp
 from pandas._typing import (
     AnyArrayLike,
     ArrayLike,
@@ -60,7 +59,6 @@
 )
 from pandas.core.dtypes.concat import concat_compat
 from pandas.core.dtypes.dtypes import (
-    ArrowDtype,
     BaseMaskedDtype,
     CategoricalDtype,
     ExtensionDtype,
@@ -1693,14 +1691,6 @@ def map_array(
         if na_action == "ignore":
             mapper = mapper[mapper.index.notna()]
 
-        if isinstance(arr.dtype, ArrowDtype) and arr.dtype.name.startswith("timestamp"):
-            try:
-                # Convert elements to pandas.Timestamp (or datetime64[ns])
-                arr = arr.astype("datetime64[ns]")
-            except Exception:
-                # fallback: safe, slow path
-                arr = np.array([Timestamp(x.as_py()) for x in arr])
-
         # Since values were input this means we came from either
         # a dict or a series and mapper should be an index
         indexer = mapper.index.get_indexer(arr)

pandas/core/arrays/arrow/array.py

@@ -1483,6 +1483,14 @@ def to_numpy(
     def map(self, mapper, na_action: Literal["ignore"] | None = None):
         if is_numeric_dtype(self.dtype):
             return map_array(self.to_numpy(), mapper, na_action=na_action)
+        if self.dtype.name.startswith("timestamp"):
+            try:
+                # Convert elements to pandas.Timestamp (or datetime64[ns])
+                self = self.astype("datetime64[ns]")
+            except Exception:
+                # fallback: safe, slow path
+                self = np.array([Timestamp(x.as_py()) for x in self])
+            return map_array(self, mapper, na_action=na_action)
         else:
             return super().map(mapper, na_action)