BUG: preserve index names in GroupBy.apply concat step, GH #481

Author: wesm

RELEASE.rst

@@ -112,6 +112,10 @@ pandas 0.6.1
   - Handle differently-indexed Series input to DataFrame constructor (GH #475)
   - Omit nuisance columns in multi-groupby with Python function
   - Buglet in handling of single grouping in general apply
+  - Handle type inference properly when passing list of lists or tuples to
+    DataFrame constructor (GH #484)
+  - Preserve Index / MultiIndex names in GroupBy.apply concatenation step (GH
+    #481)
 
 Thanks
 ------

pandas/core/groupby.py

@@ -1188,6 +1188,8 @@ def _concat_frames_hierarchical(frames, keys, groupings, axis=0):
         return DataFrame(new_values, index=new_index, columns=new_columns)
 
 def _make_concat_multiindex(indexes, keys, groupings):
+    names = [ping.name for ping in groupings]
+
     if not _all_indexes_same(indexes):
         label_list = []
 
@@ -1213,13 +1215,22 @@ def _make_concat_multiindex(indexes, keys, groupings):
         else:
             label_list.append(concat_index.values)
 
-        return MultiIndex.from_arrays(label_list)
+        consensus_name = indexes[0].names
+        for index in indexes[1:]:
+            if index.names != consensus_name:
+                consensus_name = [None] * index.nlevels
+                break
+        names.extend(consensus_name)
+
+        return MultiIndex.from_arrays(label_list, names=names)
 
     new_index = indexes[0]
     n = len(new_index)
 
+    names.append(indexes[0].name)
+
     # do something a bit more speedy
-    levels = [ping.group_index for ping in  groupings]
+    levels = [ping.group_index for ping in groupings]
     levels.append(new_index)
 
     # construct labels
@@ -1237,7 +1248,7 @@ def _make_concat_multiindex(indexes, keys, groupings):
 
     # last labels for the new level
     labels.append(np.tile(np.arange(n), len(indexes)))
-    return MultiIndex(levels=levels, labels=labels)
+    return MultiIndex(levels=levels, labels=labels, names=names)
 
 def _all_indexes_same(indexes):
     first = indexes[0]

pandas/tests/test_groupby.py

@@ -62,6 +62,19 @@ def setUp(self):
         self.mframe = DataFrame(np.random.randn(10, 3), index=index,
                                 columns=['A', 'B', 'C'])
 
+        self.three_group = DataFrame({'A' : ['foo', 'foo', 'foo', 'foo',
+                                             'bar', 'bar', 'bar', 'bar',
+                                             'foo', 'foo', 'foo'],
+                                      'B' : ['one', 'one', 'one', 'two',
+                                             'one', 'one', 'one', 'two',
+                                             'two', 'two', 'one'],
+                                      'C' : ['dull', 'dull', 'shiny', 'dull',
+                                             'dull', 'shiny', 'shiny', 'dull',
+                                             'shiny', 'shiny', 'shiny'],
+                                      'D' : np.random.randn(11),
+                                      'E' : np.random.randn(11),
+                                      'F' : np.random.randn(11)})
+
     def test_basic(self):
         data = Series(np.arange(9) // 3, index=np.arange(9))
 
@@ -686,25 +699,46 @@ def test_omit_nuisance(self):
         result = self.assertRaises(TypeError, grouped.agg, np.sum)
 
     def test_omit_nuisance_python_multiple(self):
-        data = DataFrame({'A' : ['foo', 'foo', 'foo', 'foo',
-                                 'bar', 'bar', 'bar', 'bar',
-                                 'foo', 'foo', 'foo'],
-                          'B' : ['one', 'one', 'one', 'two',
-                                 'one', 'one', 'one', 'two',
-                                 'two', 'two', 'one'],
-                          'C' : ['dull', 'dull', 'shiny', 'dull',
-                                 'dull', 'shiny', 'shiny', 'dull',
-                                 'shiny', 'shiny', 'shiny'],
-                          'D' : np.random.randn(11),
-                          'E' : np.random.randn(11),
-                          'F' : np.random.randn(11)})
-
-        grouped = data.groupby(['A', 'B'])
+        grouped = self.three_group.groupby(['A', 'B'])
 
         agged = grouped.agg(np.mean)
         exp = grouped.mean()
         assert_frame_equal(agged, exp)
 
+    def test_apply_concat_preserve_names(self):
+        grouped = self.three_group.groupby(['A', 'B'])
+
+        def desc(group):
+            result = group.describe()
+            result.index.name = 'stat'
+            return result
+
+        def desc2(group):
+            result = group.describe()
+            result.index.name = 'stat'
+            result = result[:len(group)]
+            # weirdo
+            return result
+
+        def desc3(group):
+            result = group.describe()
+
+            # names are different
+            result.index.name = 'stat_%d' % len(group)
+
+            result = result[:len(group)]
+            # weirdo
+            return result
+
+        result = grouped.apply(desc)
+        self.assertEquals(result.index.names, ['A', 'B', 'stat'])
+
+        result2 = grouped.apply(desc2)
+        self.assertEquals(result2.index.names, ['A', 'B', 'stat'])
+
+        result3 = grouped.apply(desc3)
+        self.assertEquals(result3.index.names, ['A', 'B', None])
+
     def test_nonsense_func(self):
         df = DataFrame([0])
         self.assertRaises(Exception, df.groupby, lambda x: x + 'foo')