From b2f53b53d40076e0695b08487272338e8398c680 Mon Sep 17 00:00:00 2001 From: Matthew Roeschke <10647082+mroeschke@users.noreply.github.com> Date: Tue, 13 Jun 2023 11:57:29 -0700 Subject: [PATCH 1/2] BUG: Indexing a timestamp ArrowDtype Index --- doc/source/whatsnew/v2.0.3.rst | 2 +- pandas/core/indexes/base.py | 3 ++- pandas/tests/indexing/test_datetime.py | 18 ++++++++++++++++++ 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/doc/source/whatsnew/v2.0.3.rst b/doc/source/whatsnew/v2.0.3.rst index 3da469c2e1fe6..0e32cdc95dbbc 100644 --- a/doc/source/whatsnew/v2.0.3.rst +++ b/doc/source/whatsnew/v2.0.3.rst @@ -24,7 +24,7 @@ Bug fixes - Bug in :func:`RangeIndex.union` when using ``sort=True`` with another :class:`RangeIndex` (:issue:`53490`) - Bug in :func:`read_csv` when defining ``dtype`` with ``bool[pyarrow]`` for the ``"c"`` and ``"python"`` engines (:issue:`53390`) - Bug in :meth:`Series.str.split` and :meth:`Series.str.rsplit` with ``expand=True`` for :class:`ArrowDtype` with ``pyarrow.string`` (:issue:`53532`) -- +- Bug when indexing a :class:`DataFrame` or :class:`Series` with an :class:`Index` with a timestamp :class:`ArrowDtype` would raise an ``AttributeError`` (:issue:`53644`) .. --------------------------------------------------------------------------- .. _whatsnew_203.other: diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index 2c6c6dd0e6ed1..b4a3500ab257d 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -113,6 +113,7 @@ ) from pandas.core.dtypes.concat import concat_compat from pandas.core.dtypes.dtypes import ( + ArrowDtype, CategoricalDtype, DatetimeTZDtype, ExtensionDtype, @@ -6037,7 +6038,7 @@ def _get_indexer_strict(self, key, axis_name: str_t) -> tuple[Index, np.ndarray] if isinstance(key, Index): # GH 42790 - Preserve name from an Index keyarr.name = key.name - if keyarr.dtype.kind in "mM": + if keyarr.dtype.kind in "mM" and not isinstance(keyarr.dtype, ArrowDtype): # DTI/TDI.take can infer a freq in some cases when we dont want one if isinstance(key, list) or ( isinstance(key, type(self)) diff --git a/pandas/tests/indexing/test_datetime.py b/pandas/tests/indexing/test_datetime.py index 15e1fae77d65b..6510612ba6f87 100644 --- a/pandas/tests/indexing/test_datetime.py +++ b/pandas/tests/indexing/test_datetime.py @@ -168,3 +168,21 @@ def test_getitem_str_slice_millisecond_resolution(self, frame_or_series): ], ) tm.assert_equal(result, expected) + + def test_getitem_pyarrow_index(self, frame_or_series): + # GH 53644 + pytest.importorskip("pyarrow") + obj = frame_or_series( + range(5), + index=date_range("2020", freq="D", periods=5).astype( + "timestamp[us][pyarrow]" + ), + ) + result = obj.loc[obj.index[:-3]] + expected = frame_or_series( + range(2), + index=date_range("2020", freq="D", periods=2).astype( + "timestamp[us][pyarrow]" + ), + ) + tm.assert_equal(result, expected) From 724bd8e4eff2356c2ccbc79a62566bf8b0407882 Mon Sep 17 00:00:00 2001 From: Matthew Roeschke <10647082+mroeschke@users.noreply.github.com> Date: Tue, 13 Jun 2023 15:55:49 -0700 Subject: [PATCH 2/2] more explicit check --- pandas/core/indexes/base.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index b4a3500ab257d..006381ce691fd 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -113,7 +113,6 @@ ) from pandas.core.dtypes.concat import concat_compat from pandas.core.dtypes.dtypes import ( - ArrowDtype, CategoricalDtype, DatetimeTZDtype, ExtensionDtype, @@ -6038,7 +6037,9 @@ def _get_indexer_strict(self, key, axis_name: str_t) -> tuple[Index, np.ndarray] if isinstance(key, Index): # GH 42790 - Preserve name from an Index keyarr.name = key.name - if keyarr.dtype.kind in "mM" and not isinstance(keyarr.dtype, ArrowDtype): + if lib.is_np_dtype(keyarr.dtype, "mM") or isinstance( + keyarr.dtype, DatetimeTZDtype + ): # DTI/TDI.take can infer a freq in some cases when we dont want one if isinstance(key, list) or ( isinstance(key, type(self))