Backport PR #35875 on branch 1.1.x (COMPAT: Ensure rolling indexers return intp during take operations) (#37471)

Co-authored-by: Matthew Roeschke <[email protected]>

Author: meeseeksmachine

pandas/core/window/indexers.py

@@ -7,6 +7,8 @@
 from pandas._libs.window.indexers import calculate_variable_window_bounds
 from pandas.util._decorators import Appender
 
+from pandas.core.dtypes.common import ensure_platform_int
+
 from pandas.tseries.offsets import Nano
 
 get_window_bounds_doc = """
@@ -296,9 +298,9 @@ def get_window_bounds(
         start_arrays = []
         end_arrays = []
         window_indicies_start = 0
-        for key, indicies in self.groupby_indicies.items():
+        for key, indices in self.groupby_indicies.items():
             if self.index_array is not None:
-                index_array = self.index_array.take(indicies)
+                index_array = self.index_array.take(ensure_platform_int(indices))
             else:
                 index_array = self.index_array
             indexer = self.rolling_indexer(
@@ -307,22 +309,22 @@ def get_window_bounds(
                 **self.indexer_kwargs,
             )
             start, end = indexer.get_window_bounds(
-                len(indicies), min_periods, center, closed
+                len(indices), min_periods, center, closed
             )
             start = start.astype(np.int64)
             end = end.astype(np.int64)
             # Cannot use groupby_indicies as they might not be monotonic with the object
             # we're rolling over
             window_indicies = np.arange(
-                window_indicies_start, window_indicies_start + len(indicies),
+                window_indicies_start, window_indicies_start + len(indices),
             )
-            window_indicies_start += len(indicies)
+            window_indicies_start += len(indices)
             # Extend as we'll be slicing window like [start, end)
             window_indicies = np.append(
                 window_indicies, [window_indicies[-1] + 1]
             ).astype(np.int64)
-            start_arrays.append(window_indicies.take(start))
-            end_arrays.append(window_indicies.take(end))
+            start_arrays.append(window_indicies.take(ensure_platform_int(start)))
+            end_arrays.append(window_indicies.take(ensure_platform_int(end)))
         start = np.concatenate(start_arrays)
         end = np.concatenate(end_arrays)
         # GH 35552: Need to adjust start and end based on the nans appended to values

pandas/tests/resample/test_resampler_grouper.py

@@ -6,7 +6,7 @@
 from pandas.util._test_decorators import async_mark
 
 import pandas as pd
-from pandas import DataFrame, Series, Timestamp, compat
+from pandas import DataFrame, Series, Timestamp
 import pandas._testing as tm
 from pandas.core.indexes.datetimes import date_range
 
@@ -317,7 +317,6 @@ def test_resample_groupby_with_label():
     tm.assert_frame_equal(result, expected)
 
 
-@pytest.mark.xfail(not compat.IS64, reason="GH-35148")
 def test_consistency_with_window():
 
     # consistent return values with window

pandas/tests/window/test_api.py

@@ -6,7 +6,7 @@
 import pandas.util._test_decorators as td
 
 import pandas as pd
-from pandas import DataFrame, Index, Series, Timestamp, compat, concat
+from pandas import DataFrame, Index, Series, Timestamp, concat
 import pandas._testing as tm
 from pandas.core.base import SpecificationError
 
@@ -277,7 +277,7 @@ def test_preserve_metadata():
 @pytest.mark.parametrize(
     "func,window_size,expected_vals",
     [
-        pytest.param(
+        (
             "rolling",
             2,
             [
@@ -289,7 +289,6 @@ def test_preserve_metadata():
                 [35.0, 40.0, 60.0, 40.0],
                 [60.0, 80.0, 85.0, 80],
             ],
-            marks=pytest.mark.xfail(not compat.IS64, reason="GH-35294"),
         ),
         (
             "expanding",

pandas/tests/window/test_apply.py

@@ -4,7 +4,7 @@
 from pandas.errors import NumbaUtilError
 import pandas.util._test_decorators as td
 
-from pandas import DataFrame, Index, MultiIndex, Series, Timestamp, compat, date_range
+from pandas import DataFrame, Index, MultiIndex, Series, Timestamp, date_range
 import pandas._testing as tm
 
 
@@ -142,7 +142,6 @@ def test_invalid_kwargs_nopython():
 
 
 @pytest.mark.parametrize("args_kwargs", [[None, {"par": 10}], [(10,), None]])
-@pytest.mark.xfail(not compat.IS64, reason="GH-35294")
 def test_rolling_apply_args_kwargs(args_kwargs):
     # GH 33433
     def foo(x, par):

pandas/tests/window/test_grouper.py

@@ -2,7 +2,7 @@
 import pytest
 
 import pandas as pd
-from pandas import DataFrame, Series, compat
+from pandas import DataFrame, Series
 import pandas._testing as tm
 from pandas.core.groupby.groupby import get_groupby
 
@@ -23,7 +23,6 @@ def test_mutated(self):
         g = get_groupby(self.frame, by="A", mutated=True)
         assert g.mutated
 
-    @pytest.mark.xfail(not compat.IS64, reason="GH-35294")
     def test_getitem(self):
         g = self.frame.groupby("A")
         g_mutated = get_groupby(self.frame, by="A", mutated=True)
@@ -56,7 +55,6 @@ def test_getitem_multiple(self):
         result = r.B.count()
         tm.assert_series_equal(result, expected)
 
-    @pytest.mark.xfail(not compat.IS64, reason="GH-35294")
     def test_rolling(self):
         g = self.frame.groupby("A")
         r = g.rolling(window=4)
@@ -74,7 +72,6 @@ def test_rolling(self):
     @pytest.mark.parametrize(
         "interpolation", ["linear", "lower", "higher", "midpoint", "nearest"]
     )
-    @pytest.mark.xfail(not compat.IS64, reason="GH-35294")
     def test_rolling_quantile(self, interpolation):
         g = self.frame.groupby("A")
         r = g.rolling(window=4)
@@ -105,7 +102,6 @@ def func(x):
             expected = g.apply(func)
             tm.assert_series_equal(result, expected)
 
-    @pytest.mark.xfail(not compat.IS64, reason="GH-35294")
     def test_rolling_apply(self, raw):
         g = self.frame.groupby("A")
         r = g.rolling(window=4)
@@ -115,7 +111,6 @@ def test_rolling_apply(self, raw):
         expected = g.apply(lambda x: x.rolling(4).apply(lambda y: y.sum(), raw=raw))
         tm.assert_frame_equal(result, expected)
 
-    @pytest.mark.xfail(not compat.IS64, reason="GH-35294")
     def test_rolling_apply_mutability(self):
         # GH 14013
         df = pd.DataFrame({"A": ["foo"] * 3 + ["bar"] * 3, "B": [1] * 6})
@@ -197,7 +192,6 @@ def test_expanding_apply(self, raw):
         tm.assert_frame_equal(result, expected)
 
     @pytest.mark.parametrize("expected_value,raw_value", [[1.0, True], [0.0, False]])
-    @pytest.mark.xfail(not compat.IS64, reason="GH-35294")
     def test_groupby_rolling(self, expected_value, raw_value):
         # GH 31754
 
@@ -215,7 +209,6 @@ def foo(x):
         )
         tm.assert_series_equal(result, expected)
 
-    @pytest.mark.xfail(not compat.IS64, reason="GH-35294")
     def test_groupby_rolling_center_center(self):
         # GH 35552
         series = Series(range(1, 6))
@@ -281,7 +274,6 @@ def test_groupby_rolling_center_center(self):
         )
         tm.assert_frame_equal(result, expected)
 
-    @pytest.mark.xfail(not compat.IS64, reason="GH-35294")
     def test_groupby_subselect_rolling(self):
         # GH 35486
         df = DataFrame(
@@ -307,7 +299,6 @@ def test_groupby_subselect_rolling(self):
         )
         tm.assert_series_equal(result, expected)
 
-    @pytest.mark.xfail(not compat.IS64, reason="GH-35294")
     def test_groupby_rolling_custom_indexer(self):
         # GH 35557
         class SimpleIndexer(pd.api.indexers.BaseIndexer):
@@ -331,7 +322,6 @@ def get_window_bounds(
         expected = df.groupby(df.index).rolling(window=3, min_periods=1).sum()
         tm.assert_frame_equal(result, expected)
 
-    @pytest.mark.xfail(not compat.IS64, reason="GH-35294")
     def test_groupby_rolling_subset_with_closed(self):
         # GH 35549
         df = pd.DataFrame(
@@ -356,7 +346,6 @@ def test_groupby_rolling_subset_with_closed(self):
         )
         tm.assert_series_equal(result, expected)
 
-    @pytest.mark.xfail(not compat.IS64, reason="GH-35294")
     def test_groupby_subset_rolling_subset_with_closed(self):
         # GH 35549
         df = pd.DataFrame(
@@ -384,7 +373,6 @@ def test_groupby_subset_rolling_subset_with_closed(self):
         )
         tm.assert_series_equal(result, expected)
 
-    @pytest.mark.xfail(not compat.IS64, reason="GH-35294")
     @pytest.mark.parametrize("func", ["max", "min"])
     def test_groupby_rolling_index_changed(self, func):
         # GH: #36018 nlevels of MultiIndex changed

pandas/tests/window/test_rolling.py

@@ -7,7 +7,7 @@
 import pandas.util._test_decorators as td
 
 import pandas as pd
-from pandas import DataFrame, Series, compat, date_range
+from pandas import DataFrame, Series, date_range
 import pandas._testing as tm
 from pandas.core.window import Rolling
 
@@ -150,7 +150,6 @@ def test_closed_one_entry(func):
 
 
 @pytest.mark.parametrize("func", ["min", "max"])
-@pytest.mark.xfail(not compat.IS64, reason="GH-35294")
 def test_closed_one_entry_groupby(func):
     # GH24718
     ser = pd.DataFrame(
@@ -683,7 +682,6 @@ def test_iter_rolling_datetime(expected, expected_index, window):
         ),
     ],
 )
-@pytest.mark.xfail(not compat.IS64, reason="GH-35294")
 def test_rolling_positional_argument(grouping, _index, raw):
     # GH 34605
 

pandas/tests/window/test_timeseries_window.py

@@ -7,7 +7,6 @@
     MultiIndex,
     Series,
     Timestamp,
-    compat,
     date_range,
     to_datetime,
 )
@@ -657,7 +656,6 @@ def agg_by_day(x):
 
             tm.assert_frame_equal(result, expected)
 
-    @pytest.mark.xfail(not compat.IS64, reason="GH-35294")
     def test_groupby_monotonic(self):
 
         # GH 15130
@@ -687,7 +685,6 @@ def test_groupby_monotonic(self):
         result = df.groupby("name").rolling("180D", on="date")["amount"].sum()
         tm.assert_series_equal(result, expected)
 
-    @pytest.mark.xfail(not compat.IS64, reason="GH-35294")
     def test_non_monotonic(self):
         # GH 13966 (similar to #15130, closed by #15175)