groupy apply: Ensure same index is returned for slow and fast path

Author: fjetter

doc/source/whatsnew/v1.0.1.rst

@@ -24,6 +24,7 @@ Deprecations
 Bug fixes
 ~~~~~~~~~
 - Bug in :meth:`GroupBy.apply` was raising ``TypeError`` if called with function which returned a non-pandas non-scalar object (e.g. a list) (:issue:`31441`)
+- Bug in :meth:`GroupBy.apply` where the output index type was depending on internals (:issue:`31612`)
 
 Categorical
 ^^^^^^^^^^^

pandas/_libs/reduction.pyx

@@ -493,7 +493,7 @@ 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 not chunk.index:
+ if not piece.index.equals(chunk.index):
  mutated = True
  except AttributeError:
  # `piece` might not have an index, could be e.g. an int

pandas/tests/groupby/test_apply.py

@@ -851,3 +851,20 @@ def test_apply_function_returns_non_pandas_non_scalar(function, expected_values)
  result = df.groupby("groups").apply(function)
  expected = pd.Series(expected_values, index=pd.Index(["A", "B"], name="groups"))
  tm.assert_series_equal(result, expected)
+
+
+def test_apply_fast_slow_identical():
+
+ df = DataFrame({"A": [0, 0, 1], "b": range(3)})
+
+ def slow(group):
+ # slow apply because of check `result is input`, c.f. https://github.com/pandas-dev/pandas/blob/44782c0809e296a8e57b7f77d963e999c7e0f4a7/pandas/_libs/reduction.pyx#L494
+ return group
+
+ def fast(group):
+ return group.copy()
+
+ fast_df = df.groupby("A").apply(fast)
+ slow_df = df.groupby("A").apply(slow)
+
+ tm.assert_frame_equal(fast_df, slow_df)