TST: groupby.nth with MultiIndex Interval (#52968)

pandas-dev · Apr 29, 2023 · a3c9a5c · a3c9a5c
1 parent 9f5b44c
commit a3c9a5c
Showing 1 changed file with 23 additions and 0 deletions.
diff --git a/pandas/tests/groupby/test_nth.py b/pandas/tests/groupby/test_nth.py
@@ -767,6 +767,29 @@ def test_groupby_nth_with_column_axis():
     tm.assert_frame_equal(result, expected)
 
 
+def test_groupby_nth_interval():
+    # GH#24205
+    idx_result = MultiIndex(
+        [
+            pd.CategoricalIndex([pd.Interval(0, 1), pd.Interval(1, 2)]),
+            pd.CategoricalIndex([pd.Interval(0, 10), pd.Interval(10, 20)]),
+        ],
+        [[0, 0, 0, 1, 1], [0, 1, 1, 0, -1]],
+    )
+    df_result = DataFrame({"col": range(len(idx_result))}, index=idx_result)
+    result = df_result.groupby(level=[0, 1], observed=False).nth(0)
+    val_expected = [0, 1, 3]
+    idx_expected = MultiIndex(
+        [
+            pd.CategoricalIndex([pd.Interval(0, 1), pd.Interval(1, 2)]),
+            pd.CategoricalIndex([pd.Interval(0, 10), pd.Interval(10, 20)]),
+        ],
+        [[0, 0, 1], [0, 1, 0]],
+    )
+    expected = DataFrame(val_expected, index=idx_expected, columns=["col"])
+    tm.assert_frame_equal(result, expected)
+
+
 @pytest.mark.parametrize(
     "start, stop, expected_values, expected_columns",
     [