ENH: handle dict return values and vbench, close #823

Author: wesm

pandas/core/groupby.py

@@ -1213,7 +1213,11 @@ def _get_index():
                 index = Index(keys, name=self.grouper.names[0])
             return index
 
-        if isinstance(values[0], Series):
+        if isinstance(values[0], dict):
+            # # GH #823
+            return DataFrame(values, index=keys).stack()
+
+        if isinstance(values[0], (Series, dict)):
             return self._concat_objects(keys, values,
                                         not_indexed_same=not_indexed_same)
         elif isinstance(values[0], DataFrame):

pandas/tests/test_groupby.py

@@ -1701,6 +1701,19 @@ def test_multifunc_sum_bug(self):
         result = grouped.agg({'fl':'sum',2:'size'})
         self.assert_(result['fl'].dtype == np.float64)
 
+    def test_handle_dict_return_value(self):
+        def f(group):
+            return {'min': group.min(), 'max': group.max()}
+
+        def g(group):
+            return Series({'min': group.min(), 'max': group.max()})
+
+        result = self.df.groupby('A')['C'].apply(f)
+        expected = self.df.groupby('A')['C'].apply(g)
+
+        self.assert_(isinstance(result, Series))
+        assert_series_equal(result, expected)
+
 def _check_groupby(df, result, keys, field, f=lambda x: x.sum()):
     tups = map(tuple, df[keys].values)
     tups = com._asarray_tuplesafe(tups)

vb_suite/groupby.py

@@ -144,3 +144,14 @@ def f():
 groupby_pivot_table = Benchmark(stmt, setup, start_date=datetime(2011, 12, 15))
 
 
+#----------------------------------------------------------------------
+# dict return values
+
+setup = common_setup + """
+labels = np.arange(1000).repeat(10)
+data = Series(randn(len(labels)))
+f = lambda x: {'first': x.values[0], 'last': x.values[-1]}
+"""
+
+groupby_apply_dict_return = Benchmark('data.groupby(labels).apply(f)',
+                                      setup, start_date=datetime(2011, 12, 15))