BUG: Fix groupby.apply (#28662)

Author: dsaxton

doc/source/whatsnew/v1.0.0.rst

@@ -915,7 +915,7 @@ Plotting
 Groupby/resample/rolling
 ^^^^^^^^^^^^^^^^^^^^^^^^
 
--
+- Bug in :meth:`DataFrame.groupby.apply` only showing output from a single group when function returns an :class:`Index` (:issue:`28652`)
 - Bug in :meth:`DataFrame.groupby` with multiple groups where an ``IndexError`` would be raised if any group contained all NA values (:issue:`20519`)
 - Bug in :meth:`pandas.core.resample.Resampler.size` and :meth:`pandas.core.resample.Resampler.count` returning wrong dtype when used with an empty series or dataframe (:issue:`28427`)
 - Bug in :meth:`DataFrame.rolling` not allowing for rolling over datetimes when ``axis=1`` (:issue:`28192`)

pandas/_libs/reduction.pyx

@@ -1,3 +1,4 @@
+from copy import copy
 from distutils.version import LooseVersion
 
 from cython import Py_ssize_t
@@ -15,7 +16,7 @@ from numpy cimport (ndarray,
 cnp.import_array()
 
 cimport pandas._libs.util as util
-from pandas._libs.lib import maybe_convert_objects
+from pandas._libs.lib import maybe_convert_objects, is_scalar
 
 
 cdef _check_result_array(object obj, Py_ssize_t cnt):
@@ -492,14 +493,19 @@ def apply_frame_axis0(object frame, object f, object names,
             # Need to infer if low level index slider will cause segfaults
             require_slow_apply = i == 0 and piece is chunk
             try:
-                if piece.index is chunk.index:
-                    piece = piece.copy(deep='all')
-                else:
+                if piece.index is not chunk.index:
                     mutated = True
             except AttributeError:
                 # `piece` might not have an index, could be e.g. an int
                 pass
 
+            if not is_scalar(piece):
+                # Need to copy data to avoid appending references
+                if hasattr(piece, "copy"):
+                    piece = piece.copy(deep="all")
+                else:
+                    piece = copy(piece)
+
             results.append(piece)
 
             # If the data was modified inplace we need to

pandas/tests/groupby/test_apply.py

@@ -686,6 +686,17 @@ def test_apply_with_mixed_types():
     tm.assert_frame_equal(result, expected)
 
 
+def test_func_returns_object():
+    # GH 28652
+    df = DataFrame({"a": [1, 2]}, index=pd.Int64Index([1, 2]))
+    result = df.groupby("a").apply(lambda g: g.index)
+    expected = Series(
+        [pd.Int64Index([1]), pd.Int64Index([2])], index=pd.Int64Index([1, 2], name="a")
+    )
+
+    tm.assert_series_equal(result, expected)
+
+
 @pytest.mark.parametrize(
     "group_column_dtlike",
     [datetime.today(), datetime.today().date(), datetime.today().time()],