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()],