result index instead of levels[0]

Author: MarcoGorelli

pandas/core/groupby/ops.py

@@ -140,18 +140,16 @@ def _get_grouper(self):
         return self.groupings[0].grouper
 
     def _get_group_keys(self):
-        # if len(self.groupings) == 1:
-        #     return self.levels[0]
-        # else:
-        comp_ids, _, ngroups = self.group_info
-
-        # provide "flattened" iterator for multi-group setting
-        flattened_iterator = get_flattened_iterator(
-            comp_ids, ngroups, self.levels, self.codes
-        )
         if len(self.groupings) == 1:
-            return Index([i[0] for i in flattened_iterator], name=self.levels[0].name)
-        return flattened_iterator
+            return self.result_index
+        else:
+            comp_ids, _, ngroups = self.group_info
+
+            # provide "flattened" iterator for multi-group setting
+            flattened_iterator = get_flattened_iterator(
+                comp_ids, ngroups, self.levels, self.codes
+            )
+            return flattened_iterator
 
     def apply(self, f, data: FrameOrSeries, axis: int = 0):
         # breakpoint()