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

Author: fjetter

doc/source/whatsnew/v1.0.1.rst

@@ -41,6 +41,7 @@ Deprecations
 
 Bug fixes
 ~~~~~~~~~
+- Bug in :meth:`core.groupby.DataFrameGroupBy.apply` where the result shape was incorrect when the output index was not identical to the input index (:issue:`31612`)
 
 **Datetimelike**
 

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,22 @@ 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():
+    # GH 31613
+
+    df = DataFrame({"A": [0, 0, 1], "b": range(3)})
+
+    # For simple index structures we check for fast/slow apply using
+    # an identity check on in/output
+    def slow(group):
+        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)