Test for Group By - Apply Key Error

Author: brunoacpcosta

pandas/tests/groupby/test_apply.py

@@ -1177,3 +1177,34 @@ def test_apply_empty_string_nan_coerce_bug():
         index=MultiIndex.from_tuples([(1, ""), (2, "")], names=["a", "b"]),
     )
     tm.assert_frame_equal(result, expected)
+
+
+def test_apply_index_key_error_bug():
+    # GH 44310
+    # Tests for a key error while applying in a grouped DF
+
+    result = DataFrame(
+        {
+            "a": ["aa", "a2", "a3"],
+            "b": [1, 2, 3],
+        },
+        index=Index([1.0, 2.0, 3.0]),
+    )
+
+    expected = DataFrame(
+        {
+            "a": ["aa", "a2", "a3"],
+            "b": [1, 2, 3],
+        },
+        index=Index([1, 2, 3]),
+    )
+    try:
+        result = result.groupby("a").apply(
+            lambda df: Series([df["b"].mean()], index=["b_mean"])
+        )
+        expected = expected.groupby("a").apply(
+            lambda df: Series([df["b"].mean()], index=["b_mean"])
+        )
+        tm.assert_frame_equal(result, expected)
+    except KeyError:
+        assert False