Skip to content

Commit 87e554d

Browse files
authored
BUG: RollingGroupby when groupby key is in the index (#37661)
1 parent 71a3375 commit 87e554d

File tree

3 files changed

+59
-8
lines changed

3 files changed

+59
-8
lines changed

doc/source/whatsnew/v1.1.5.rst

+1
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ Bug fixes
2525
~~~~~~~~~
2626
- Bug in metadata propagation for ``groupby`` iterator (:issue:`37343`)
2727
- Bug in indexing on a :class:`Series` with ``CategoricalDtype`` after unpickling (:issue:`37631`)
28+
- Bug in :class:`RollingGroupby` with the resulting :class:`MultiIndex` when grouping by a label that is in the index (:issue:`37641`)
2829
-
2930

3031
.. ---------------------------------------------------------------------------

pandas/core/window/rolling.py

+27-7
Original file line numberDiff line numberDiff line change
@@ -762,19 +762,39 @@ def _apply(
762762
use_numba_cache,
763763
**kwargs,
764764
)
765-
# Compose MultiIndex result from grouping levels then rolling level
766-
# Aggregate the MultiIndex data as tuples then the level names
767-
grouped_object_index = self.obj.index
768-
grouped_index_name = [*grouped_object_index.names]
769-
groupby_keys = [grouping.name for grouping in self._groupby.grouper._groupings]
770-
result_index_names = groupby_keys + grouped_index_name
765+
# Reconstruct the resulting MultiIndex from tuples
766+
# 1st set of levels = group by labels
767+
# 2nd set of levels = original index
768+
# Ignore 2nd set of levels if a group by label include an index level
769+
result_index_names = [
770+
grouping.name for grouping in self._groupby.grouper._groupings
771+
]
772+
grouped_object_index = None
773+
774+
column_keys = [
775+
key
776+
for key in result_index_names
777+
if key not in self.obj.index.names or key is None
778+
]
779+
780+
if len(column_keys) == len(result_index_names):
781+
grouped_object_index = self.obj.index
782+
grouped_index_name = [*grouped_object_index.names]
783+
result_index_names += grouped_index_name
784+
else:
785+
# Our result will have still kept the column in the result
786+
result = result.drop(columns=column_keys, errors="ignore")
771787

772788
result_index_data = []
773789
for key, values in self._groupby.grouper.indices.items():
774790
for value in values:
775791
data = [
776792
*com.maybe_make_list(key),
777-
*com.maybe_make_list(grouped_object_index[value]),
793+
*com.maybe_make_list(
794+
grouped_object_index[value]
795+
if grouped_object_index is not None
796+
else []
797+
),
778798
]
779799
result_index_data.append(tuple(data))
780800

pandas/tests/window/test_grouper.py

+31-1
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
import pytest
33

44
import pandas as pd
5-
from pandas import DataFrame, Series
5+
from pandas import DataFrame, MultiIndex, Series
66
import pandas._testing as tm
77
from pandas.core.groupby.groupby import get_groupby
88

@@ -601,3 +601,33 @@ def test_groupby_rolling_nans_in_index(self, rollings, key):
601601
df = df.set_index("a")
602602
with pytest.raises(ValueError, match=f"{key} must be monotonic"):
603603
df.groupby("c").rolling("60min", **rollings)
604+
605+
def test_groupby_rolling_group_keys(self):
606+
# GH 37641
607+
arrays = [["val1", "val1", "val2"], ["val1", "val1", "val2"]]
608+
index = MultiIndex.from_arrays(arrays, names=("idx1", "idx2"))
609+
610+
s = Series([1, 2, 3], index=index)
611+
result = s.groupby(["idx1", "idx2"], group_keys=False).rolling(1).mean()
612+
expected = Series(
613+
[1.0, 2.0, 3.0],
614+
index=MultiIndex.from_tuples(
615+
[("val1", "val1"), ("val1", "val1"), ("val2", "val2")],
616+
names=["idx1", "idx2"],
617+
),
618+
)
619+
tm.assert_series_equal(result, expected)
620+
621+
def test_groupby_rolling_index_level_and_column_label(self):
622+
arrays = [["val1", "val1", "val2"], ["val1", "val1", "val2"]]
623+
index = MultiIndex.from_arrays(arrays, names=("idx1", "idx2"))
624+
625+
df = DataFrame({"A": [1, 1, 2], "B": range(3)}, index=index)
626+
result = df.groupby(["idx1", "A"]).rolling(1).mean()
627+
expected = DataFrame(
628+
{"B": [0.0, 1.0, 2.0]},
629+
index=MultiIndex.from_tuples(
630+
[("val1", 1), ("val1", 1), ("val2", 2)], names=["idx1", "A"]
631+
),
632+
)
633+
tm.assert_frame_equal(result, expected)

0 commit comments

Comments
 (0)