BUG: Pyarrow timestamp support for map() function #61236
Conversation
pandas/core/algorithms.py
Outdated
| try: | ||
| # Convert elements to pandas.Timestamp (or datetime64[ns]) | ||
| arr = arr.astype("datetime64[ns]") | ||
| except Exception: |
There was a problem hiding this comment.
This is the wrong place to fix this. This should be fixed in ArrowExtensionArray.map
pandas/core/arrays/arrow/array.py
Outdated
| @@ -1483,6 +1483,8 @@ 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) | |||
| elif self.dtype == "timestamp[ns][pyarrow]": | |||
| return map_array(self.to_numpy(dtype=object), mapper, na_action=na_action) | |||
There was a problem hiding this comment.
Can you avoid the type cast to object?
There was a problem hiding this comment.
I tried using datetime64[ns] instead of object, but some tests expect Python objects (pd.Timestamp, ) and do not pass. I think keeping object helps preserve that expected behavior. Let me know if you'd prefer adjusting the test instead.
There was a problem hiding this comment.
I think the failing test would need adjustment (we get a better result when we don't return object)
|
|
||
| def test_map_arrow_timestamp_dict(): | ||
| # GH 61231 | ||
| pytest.importorskip("pyarrow", minversion="10.0.1") |
There was a problem hiding this comment.
| pytest.importorskip("pyarrow", minversion="10.0.1") | |
| pytest.importorskip("pyarrow") |
| @@ -1483,6 +1483,10 @@ 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) | |||
| elif self.dtype == "timestamp[ns][pyarrow]": | |||
There was a problem hiding this comment.
Instead of adding an elif you can modify the existing if statement as if is_numeric_dtype(self.dtype) or self.dtype.kind in "mM":
Added type annotations to new arguments/methods/functions.doc/source/whatsnew/vX.X.X.rstfile if fixing a bug or adding a new feature.