Repository: pandas-dev/pandas
Pull request #56668 (2023-12-28 21:44:11)

Backport PR #56654 on branch 2.2.x (BUG: assert_series_equal not properly respecting check-dtype)
Backport PR #56654: BUG: assert_series_equal not properly respecting check-dtype


Patch:
diff --git a/doc/source/whatsnew/v2.2.0.rst b/doc/source/whatsnew/v2.2.0.rst
index 1da18cd9be..129f5cedb8 100644
--- a/doc/source/whatsnew/v2.2.0.rst
+++ b/doc/source/whatsnew/v2.2.0.rst
@@ -865,6 +865,7 @@ Groupby/resample/rolling
 - Bug in :meth:`DataFrame.resample` when resampling on a :class:`ArrowDtype` of ``pyarrow.timestamp`` or ``pyarrow.duration`` type (:issue:`55989`)
 - Bug in :meth:`DataFrame.resample` where bin edges were not correct for :class:`~pandas.tseries.offsets.BusinessDay` (:issue:`55281`)
 - Bug in :meth:`DataFrame.resample` where bin edges were not correct for :class:`~pandas.tseries.offsets.MonthBegin` (:issue:`55271`)
+- Bug in :meth:`DataFrame.rolling` and :meth:`Series.rolling` where either the ``index`` or ``on`` column was :class:`ArrowDtype` with ``pyarrow.timestamp`` type (:issue:`55849`)
 
 Reshaping
 ^^^^^^^^^
diff --git a/pandas/core/arrays/datetimelike.py b/pandas/core/arrays/datetimelike.py
index 11a0c7bf18..a0e0a1434e 100644
--- a/pandas/core/arrays/datetimelike.py
+++ b/pandas/core/arrays/datetimelike.py
@@ -92,6 +92,7 @@ from pandas.core.dtypes.common import (
     pandas_dtype,
 )
 from pandas.core.dtypes.dtypes import (
+    ArrowDtype,
     CategoricalDtype,
     DatetimeTZDtype,
     ExtensionDtype,
@@ -2531,7 +2532,7 @@ def _validate_inferred_freq(
     return freq
 
 
-def dtype_to_unit(dtype: DatetimeTZDtype | np.dtype) -> str:
+def dtype_to_unit(dtype: DatetimeTZDtype | np.dtype | ArrowDtype) -> str:
     """
     Return the unit str corresponding to the dtype's resolution.
 
@@ -2546,4 +2547,8 @@ def dtype_to_unit(dtype: DatetimeTZDtype | np.dtype) -> str:
     """
     if isinstance(dtype, DatetimeTZDtype):
         return dtype.unit
+    elif isinstance(dtype, ArrowDtype):
+        if dtype.kind not in "mM":
+            raise ValueError(f"{dtype=} does not have a resolution.")
+        return dtype.pyarrow_dtype.unit
     return np.datetime_data(dtype)[0]
diff --git a/pandas/core/window/rolling.py b/pandas/core/window/rolling.py
index e78bd258c1..68cec16ec9 100644
--- a/pandas/core/window/rolling.py
+++ b/pandas/core/window/rolling.py
@@ -14,7 +14,6 @@ from typing import (
     Any,
     Callable,
     Literal,
-    cast,
 )
 
 import numpy as np
@@ -39,6 +38,7 @@ from pandas.core.dtypes.common import (
     is_numeric_dtype,
     needs_i8_conversion,
 )
+from pandas.core.dtypes.dtypes import ArrowDtype
 from pandas.core.dtypes.generic import (
     ABCDataFrame,
     ABCSeries,
@@ -104,6 +104,7 @@ if TYPE_CHECKING:
         NDFrameT,
         QuantileInterpolation,
         WindowingRankType,
+        npt,
     )
 
     from pandas import (
@@ -404,11 +405,12 @@ class BaseWindow(SelectionMixin):
                 result[name] = extra_col
 
     @property
-    def _index_array(self):
+    def _index_array(self) -> npt.NDArray[np.int64] | None:
         # TODO: why do we get here with e.g. MultiIndex?
-        if needs_i8_conversion(self._on.dtype):
-            idx = cast("PeriodIndex | DatetimeIndex | TimedeltaIndex", self._on)
-            return idx.asi8
+        if isinstance(self._on, (PeriodIndex, DatetimeIndex, TimedeltaIndex)):
+            return self._on.asi8
+        elif isinstance(self._on.dtype, ArrowDtype) and self._on.dtype.kind in "mM":
+            return self._on.to_numpy(dtype=np.int64)
         return None
 
     def _resolve_output(self, out: DataFrame, obj: DataFrame) -> DataFrame:
@@ -439,7 +441,7 @@ class BaseWindow(SelectionMixin):
         self, homogeneous_func: Callable[..., ArrayLike], name: str | None = None
     ) -> Series:
         """
-        Series version of _apply_blockwise
+        Series version of _apply_columnwise
         """
         obj = self._create_data(self._selected_obj)
 
@@ -455,7 +457,7 @@ class BaseWindow(SelectionMixin):
         index = self._slice_axis_for_step(obj.index, result)
         return obj._constructor(result, index=index, name=obj.name)
 
-    def _apply_blockwise(
+    def _apply_columnwise(
         self,
         homogeneous_func: Callable[..., ArrayLike],
         name: str,
@@ -614,7 +616,7 @@ class BaseWindow(SelectionMixin):
             return result
 
         if self.method == "single":
-            return self._apply_blockwise(homogeneous_func, name, numeric_only)
+            return self._apply_columnwise(homogeneous_func, name, numeric_only)
         else:
             return self._apply_tablewise(homogeneous_func, name, numeric_only)
 
@@ -1232,7 +1234,9 @@ class Window(BaseWindow):
 
             return result
 
-        return self._apply_blockwise(homogeneous_func, name, numeric_only)[:: self.step]
+        return self._apply_columnwise(homogeneous_func, name, numeric_only)[
+            :: self.step
+        ]
 
     @doc(
         _shared_docs["aggregate"],
@@ -1868,6 +1872,7 @@ class Rolling(RollingAndExpandingMixin):
         if (
             self.obj.empty
             or isinstance(self._on, (DatetimeIndex, TimedeltaIndex, PeriodIndex))
+            or (isinstance(self._on.dtype, ArrowDtype) and self._on.dtype.kind in "mM")
         ) and isinstance(self.window, (str, BaseOffset, timedelta)):
             self._validate_datetimelike_monotonic()
 


Test patch:
diff --git a/pandas/_testing/asserters.py b/pandas/_testing/asserters.py
index 800b037075..d0f38c8586 100644
--- a/pandas/_testing/asserters.py
+++ b/pandas/_testing/asserters.py
@@ -949,9 +949,15 @@ def assert_series_equal(
                 obj=str(obj),
             )
         else:
+            # convert both to NumPy if not, check_dtype would raise earlier
+            lv, rv = left_values, right_values
+            if isinstance(left_values, ExtensionArray):
+                lv = left_values.to_numpy()
+            if isinstance(right_values, ExtensionArray):
+                rv = right_values.to_numpy()
             assert_numpy_array_equal(
-                left_values,
-                right_values,
+                lv,
+                rv,
                 check_dtype=check_dtype,
                 obj=str(obj),
                 index_values=left.index,
diff --git a/pandas/tests/extension/test_numpy.py b/pandas/tests/extension/test_numpy.py
index aaf49f53ba..e38144f4c6 100644
--- a/pandas/tests/extension/test_numpy.py
+++ b/pandas/tests/extension/test_numpy.py
@@ -421,16 +421,6 @@ class TestNumpyExtensionArray(base.ExtensionTests):
     def test_EA_types(self, engine, data, request):
         super().test_EA_types(engine, data, request)
 
-    @pytest.mark.xfail(reason="Expect NumpyEA, get np.ndarray")
-    def test_compare_array(self, data, comparison_op):
-        super().test_compare_array(data, comparison_op)
-
-    def test_compare_scalar(self, data, comparison_op, request):
-        if data.dtype.kind == "f" or comparison_op.__name__ in ["eq", "ne"]:
-            mark = pytest.mark.xfail(reason="Expect NumpyEA, get np.ndarray")
-            request.applymarker(mark)
-        super().test_compare_scalar(data, comparison_op)
-
 
 class Test2DCompat(base.NDArrayBacked2DTests):
     pass
diff --git a/pandas/tests/util/test_assert_frame_equal.py b/pandas/tests/util/test_assert_frame_equal.py
index a074898f60..79132591b1 100644
--- a/pandas/tests/util/test_assert_frame_equal.py
+++ b/pandas/tests/util/test_assert_frame_equal.py
@@ -211,10 +211,7 @@ def test_assert_frame_equal_extension_dtype_mismatch():
         "\\[right\\]: int[32|64]"
     )
 
-    # TODO: this shouldn't raise (or should raise a better error message)
-    # https://github.com/pandas-dev/pandas/issues/56131
-    with pytest.raises(AssertionError, match="classes are different"):
-        tm.assert_frame_equal(left, right, check_dtype=False)
+    tm.assert_frame_equal(left, right, check_dtype=False)
 
     with pytest.raises(AssertionError, match=msg):
         tm.assert_frame_equal(left, right, check_dtype=True)
@@ -246,7 +243,6 @@ def test_assert_frame_equal_ignore_extension_dtype_mismatch():
     tm.assert_frame_equal(left, right, check_dtype=False)
 
 
-@pytest.mark.xfail(reason="https://github.com/pandas-dev/pandas/issues/56131")
 def test_assert_frame_equal_ignore_extension_dtype_mismatch_cross_class():
     # https://github.com/pandas-dev/pandas/issues/35715
     left = DataFrame({"a": [1, 2, 3]}, dtype="Int64")
@@ -300,9 +296,7 @@ def test_frame_equal_mixed_dtypes(frame_or_series, any_numeric_ea_dtype, indexer
     dtypes = (any_numeric_ea_dtype, "int64")
     obj1 = frame_or_series([1, 2], dtype=dtypes[indexer[0]])
     obj2 = frame_or_series([1, 2], dtype=dtypes[indexer[1]])
-    msg = r'(Series|DataFrame.iloc\[:, 0\] \(column name="0"\) classes) are different'
-    with pytest.raises(AssertionError, match=msg):
-        tm.assert_equal(obj1, obj2, check_exact=True, check_dtype=False)
+    tm.assert_equal(obj1, obj2, check_exact=True, check_dtype=False)
 
 
 def test_assert_frame_equal_check_like_different_indexes():
diff --git a/pandas/tests/util/test_assert_series_equal.py b/pandas/tests/util/test_assert_series_equal.py
index f722f619bc..c4ffc19729 100644
--- a/pandas/tests/util/test_assert_series_equal.py
+++ b/pandas/tests/util/test_assert_series_equal.py
@@ -290,10 +290,7 @@ Attribute "dtype" are different
 \\[left\\]:  Int64
 \\[right\\]: int[32|64]"""
 
-    # TODO: this shouldn't raise (or should raise a better error message)
-    # https://github.com/pandas-dev/pandas/issues/56131
-    with pytest.raises(AssertionError, match="Series classes are different"):
-        tm.assert_series_equal(left, right, check_dtype=False)
+    tm.assert_series_equal(left, right, check_dtype=False)
 
     with pytest.raises(AssertionError, match=msg):
         tm.assert_series_equal(left, right, check_dtype=True)
@@ -372,7 +369,6 @@ def test_assert_series_equal_ignore_extension_dtype_mismatch():
     tm.assert_series_equal(left, right, check_dtype=False)
 
 
-@pytest.mark.xfail(reason="https://github.com/pandas-dev/pandas/issues/56131")
 def test_assert_series_equal_ignore_extension_dtype_mismatch_cross_class():
     # https://github.com/pandas-dev/pandas/issues/35715
     left = Series([1, 2, 3], dtype="Int64")
@@ -456,3 +452,13 @@ def test_large_unequal_ints(dtype):
     right = Series([1577840521123543], dtype=dtype)
     with pytest.raises(AssertionError, match="Series are different"):
         tm.assert_series_equal(left, right)
+
+
+@pytest.mark.parametrize("dtype", [None, object])
+@pytest.mark.parametrize("check_exact", [True, False])
+@pytest.mark.parametrize("val", [3, 3.5])
+def test_ea_and_numpy_no_dtype_check(val, check_exact, dtype):
+    # GH#56651
+    left = Series([1, 2, val], dtype=dtype)
+    right = Series(pd.array([1, 2, val]))
+    tm.assert_series_equal(left, right, check_dtype=False, check_exact=check_exact)
diff --git a/pandas/tests/window/test_timeseries_window.py b/pandas/tests/window/test_timeseries_window.py
index c99fc8a8eb..bd0fadeb3e 100644
--- a/pandas/tests/window/test_timeseries_window.py
+++ b/pandas/tests/window/test_timeseries_window.py
@@ -1,9 +1,12 @@
 import numpy as np
 import pytest
 
+import pandas.util._test_decorators as td
+
 from pandas import (
     DataFrame,
     DatetimeIndex,
+    Index,
     MultiIndex,
     NaT,
     Series,
@@ -697,3 +700,16 @@ def test_nat_axis_error(msg, axis):
     with pytest.raises(ValueError, match=f"{msg} values must not have NaT"):
         with tm.assert_produces_warning(FutureWarning, match=warn_msg):
             df.rolling("D", axis=axis).mean()
+
+
+@td.skip_if_no("pyarrow")
+def test_arrow_datetime_axis():
+    # GH 55849
+    expected = Series(
+        np.arange(5, dtype=np.float64),
+        index=Index(
+            date_range("2020-01-01", periods=5), dtype="timestamp[ns][pyarrow]"
+        ),
+    )
+    result = expected.rolling("1D").sum()
+    tm.assert_series_equal(result, expected)