TEST: Add test for pandas-dev#22541 to ensure that error does not occur again

phofl · patrick · commit 5a381c35503a · 2020-03-28T13:43:47.000+01:00
diff --git a/pandas/tests/groupby/test_apply.py b/pandas/tests/groupby/test_apply.py
@@ -866,3 +866,16 @@ def fct(group):
         [[1.0, 2.0], [3.0], [np.nan]], index=pd.Index(["a", "b", "none"], name="A")
     )
     tm.assert_series_equal(result, expected)
+
+
+def test_apply_function_index_return():
+    # GH: 22541
+    df = pd.DataFrame([1, 2, 2, 2, 1, 2, 3, 1, 3, 1], columns=["id"])
+    result = df.groupby("id").apply(lambda gr: gr.index)
+    result_2 = df.groupby("id").apply(lambda gr: gr.index + 1 - 1)
+    expected = pd.Series(
+        [pd.Index([0, 4, 7, 9]), pd.Index([1, 2, 3, 5]), pd.Index([6, 8])],
+        index=pd.Index([1, 2, 3], name="id"),
+    )
+    tm.assert_series_equal(result, expected)
+    tm.assert_series_equal(result, result_2)
