BUG: fix GroupBy.apply bug, GH #237

Author: wesm

pandas/core/groupby.py

@@ -1069,22 +1069,35 @@ def transform(self, func, *args, **kwargs):
                               axis=self.axis)
 
 def _concat_frames(frames, index, columns=None, axis=0):
-    if axis == 0:
-        all_index = [np.asarray(x.index) for x in frames]
-        new_index = Index(np.concatenate(all_index))
+    if len(frames) == 1:
+        return frames[0]
 
+    if axis == 0:
+        new_index = _concat_indexes([x.index for x in frames])
         if columns is None:
             new_columns = frames[0].columns
         else:
             new_columns = columns
     else:
-        all_columns = [np.asarray(x.columns) for x in frames]
-        new_columns = Index(np.concatenate(all_columns))
+        new_columns = _concat_indexes([x.columns for x in frames])
         new_index = index
 
-    new_values = np.concatenate([x.values for x in frames], axis=axis)
-    result = DataFrame(new_values, index=new_index, columns=new_columns)
-    return result.reindex(index=index, columns=columns)
+    if frames[0]._is_mixed_type:
+        new_data = {}
+        for col in new_columns:
+            new_data[col] = np.concatenate([x[col].values for x in frames])
+        return DataFrame(new_data, index=new_index, columns=new_columns)
+    else:
+        new_values = np.concatenate([x.values for x in frames], axis=axis)
+        result = DataFrame(new_values, index=new_index, columns=new_columns)
+        return result.reindex(index=index, columns=columns)
+
+def _concat_indexes(indexes):
+    if len(indexes) == 1:
+        new_index = indexes[0]
+    else:
+        new_index = indexes[0].append(indexes[1:])
+    return new_index
 
 def _concat_frames_hierarchical(frames, keys, groupings, axis=0):
     if axis == 0:
@@ -1096,8 +1109,14 @@ def _concat_frames_hierarchical(frames, keys, groupings, axis=0):
         new_columns = _make_concat_multiindex(all_columns, keys, groupings)
         new_index = frames[0].index
 
-    new_values = np.concatenate([x.values for x in frames], axis=axis)
-    return DataFrame(new_values, index=new_index, columns=new_columns)
+    if frames[0]._is_mixed_type:
+        new_data = {}
+        for col in new_columns:
+            new_data[col] = np.concatenate([x[col].values for x in frames])
+        return DataFrame(new_data, index=new_index, columns=new_columns)
+    else:
+        new_values = np.concatenate([x.values for x in frames], axis=axis)
+        return DataFrame(new_values, index=new_index, columns=new_columns)
 
 def _make_concat_multiindex(indexes, keys, groupings):
     if not _all_indexes_same(indexes):

pandas/tests/test_groupby.py

@@ -872,7 +872,7 @@ def test_grouping_ndarray(self):
         expected = self.df.groupby('A').sum()
         assert_frame_equal(result, expected)
 
-    def test_apply_example(self):
+    def test_apply_typecast_fail(self):
         df = DataFrame({'d' : [1.,1.,1.,2.,2.,2.],
                         'c' : np.tile(['a','b','c'], 2),
                         'v' : np.arange(1., 7.)})
@@ -889,6 +889,24 @@ def f(group):
 
         assert_frame_equal(result, expected)
 
+    def test_apply_multiindex_fail(self):
+        index = MultiIndex.from_arrays([[0, 0, 0, 1, 1, 1],
+                                        [1, 2, 3, 1, 2, 3]])
+        df = DataFrame({'d' : [1.,1.,1.,2.,2.,2.],
+                        'c' : np.tile(['a','b','c'], 2),
+                        'v' : np.arange(1., 7.)}, index=index)
+
+        def f(group):
+            v = group['v']
+            group['v2'] = (v - v.min()) / (v.max() - v.min())
+            return group
+
+        result = df.groupby('d').apply(f)
+
+        expected = df.copy()
+        expected['v2'] = np.tile([0., 0.5, 1], 2)
+
+        assert_frame_equal(result, expected)
 
 class TestPanelGroupBy(unittest.TestCase):