Add tests

Author: Kei

pandas/tests/groupby/test_apply.py

@@ -1583,3 +1583,110 @@ def f_4(grp):
     e.loc["Pony"] = np.nan
     e.name = None
     tm.assert_series_equal(result, e)
+
+
+# GH58291
+def test_apply_frame_not_as_index_returns_single_index():
+    df = DataFrame(
+        [
+            ["group_a", 0],
+            ["group_a", 2],
+            ["group_b", 1],
+            ["group_b", 3],
+            ["group_b", 5],
+        ],
+        columns=["group", "value"],
+    )
+    gb = df.groupby("group", as_index=False)[["group", "value"]]
+
+    def up_to_two_rows(df: DataFrame) -> DataFrame:
+        return df.head(2)
+
+    result = gb.apply(up_to_two_rows)
+
+    expected = DataFrame(
+        [["group_a", 0], ["group_a", 2], ["group_b", 1], ["group_b", 3]],
+        columns=["group", "value"],
+        index=range(4),
+    )
+    tm.assert_frame_equal(result, expected)
+
+
+# GH58291
+def test_apply_column_groupby_frame_not_as_index_returns_single_index(df):
+    gb = df.groupby("A", as_index=False)[["A", "C", "D"]]
+    result = gb.apply(lambda x: x.min())
+    print(str(result))
+
+    expected = DataFrame(
+        [["bar", -2.441467, -2.441467], ["foo", -0.413064, -0.413064]],
+        columns=["A", "C", "D"],
+        index=range(2),
+    )
+    tm.assert_frame_equal(result, expected)
+
+
+# GH58291
+def test_apply_non_column_groupby_frame_not_as_index_returns_single_index(tsframe):
+    gb = tsframe.groupby(lambda x: x.month, as_index=False)
+    result = gb.apply(lambda x: x.mean())
+    print(str(result))
+
+    expected = DataFrame(
+        [
+            [1, 0.115464, 0.260960, 0.187824, -0.411523],
+            [2, 0.047104, -0.183591, 0.330640, -0.207427],
+        ],
+        columns=["index", "A", "B", "C", "D"],
+        index=range(2),
+    )
+    tm.assert_frame_equal(result, expected)
+
+    result = gb.apply(lambda x: x.mean() + x.max())
+    print(str(result))
+
+    expected = DataFrame(
+        [
+            [1, 2.490351, 2.062381, 2.244527, 0.912824],
+            [2, 1.918738, 2.273746, 1.759184, 0.700975],
+        ],
+        columns=["index", "A", "B", "C", "D"],
+        index=range(2),
+    )
+    tm.assert_frame_equal(result, expected)
+
+    result = gb.apply(lambda x: 1)
+    print(str(result))
+
+    expected = DataFrame([[1, 1], [2, 1]], columns=["index", None], index=range(2))
+    tm.assert_frame_equal(result, expected)
+
+    result = gb.apply(lambda x: x.quantile([0.2, 0.38]))
+    print(str(result))
+
+    expected = DataFrame(
+        [
+            [-0.742155, -0.607186, -0.325423, -1.254187],
+            [-0.212477, -0.034416, -0.117452, -0.561536],
+            [-0.842537, -0.920719, -0.404496, -0.729036],
+            [-0.411161, -0.830757, 0.321981, -0.444410],
+        ],
+        columns=["A", "B", "C", "D"],
+        index=[0.20, 0.38] * 2,
+    )
+    tm.assert_frame_equal(result, expected)
+
+    result = gb.apply(DataFrame.quantile, [0.2, 0.38])
+    print(str(result))
+
+    expected = DataFrame(
+        [
+            [-0.742155, -0.607186, -0.325423, -1.254187],
+            [-0.212477, -0.034416, -0.117452, -0.561536],
+            [-0.842537, -0.920719, -0.404496, -0.729036],
+            [-0.411161, -0.830757, 0.321981, -0.444410],
+        ],
+        columns=["A", "B", "C", "D"],
+        index=[0.20, 0.38] * 2,
+    )
+    tm.assert_frame_equal(result, expected)