ENH: can pass list of functions to pivot_table

Author: wesm

pandas/__init__.py

@@ -29,5 +29,5 @@
 from pandas.io.pytables import HDFStore
 from pandas.util.testing import debug
 
-from pandas.tools.pivot import pivot_table
 from pandas.tools.merge import merge, concat
+from pandas.tools.pivot import pivot_table

pandas/tools/merge.py

@@ -934,25 +934,31 @@ def _make_concat_multiindex(indexes, keys, levels=None, names=None):
 
     new_index = indexes[0]
     n = len(new_index)
+    kpieces = len(indexes)
 
     # also copies
-    names = names + [indexes[0].name]
+    new_names = list(names)
+    new_levels = list(levels)
 
-    new_levels = levels
+    # construct labels
+    new_labels = []
 
     # do something a bit more speedy
-    new_levels.append(new_index)
-
-    # construct labels
-    labels = []
 
-    for hlevel, level in zip(zipped, new_levels[:-1]):
+    for hlevel, level in zip(zipped, levels):
         mapped = level.get_indexer(hlevel)
-        labels.append(np.repeat(mapped, n))
+        new_labels.append(np.repeat(mapped, n))
+
+    if isinstance(new_index, MultiIndex):
+        new_levels.extend(new_index.levels)
+        new_labels.extend([np.tile(lab, kpieces) for lab in new_index.labels])
+        new_names.extend(new_index.names)
+    else:
+        new_levels.append(new_index)
+        new_names.append(new_index.name)
+        new_labels.append(np.tile(np.arange(n), kpieces))
 
-    # last labels for the new level
-    labels.append(np.tile(np.arange(n), len(indexes)))
-    return MultiIndex(levels=new_levels, labels=labels, names=names)
+    return MultiIndex(levels=new_levels, labels=new_labels, names=new_names)
 
 def _get_consensus_names(indexes):
     consensus_name = indexes[0].names

pandas/tools/pivot.py

@@ -1,4 +1,7 @@
+# pylint: disable=E1103
+
 from pandas import Series, DataFrame
+from pandas.tools.merge import concat
 import numpy as np
 
 def pivot_table(data, values=None, rows=None, cols=None, aggfunc=np.mean,
@@ -16,7 +19,10 @@ def pivot_table(data, values=None, rows=None, cols=None, aggfunc=np.mean,
         Columns to group on the x-axis of the pivot table
     cols : list
         Columns to group on the x-axis of the pivot table
-    aggfunc : function, default numpy.mean
+    aggfunc : function, default numpy.mean, or list of functions
+        If list of functions passed, the resulting pivot table will have
+        hierarchical columns whose top level are the function names (inferred
+        from the function objects themselves)
     fill_value : scalar, default None
         Value to replace missing values with
     margins : boolean, default False
@@ -52,6 +58,17 @@ def pivot_table(data, values=None, rows=None, cols=None, aggfunc=np.mean,
     rows = _convert_by(rows)
     cols = _convert_by(cols)
 
+    if isinstance(aggfunc, list):
+        pieces = []
+        keys = []
+        for func in aggfunc:
+            table = pivot_table(data, values=values, rows=rows, cols=cols,
+                                fill_value=fill_value, aggfunc=func,
+                                margins=margins)
+            pieces.append(table)
+            keys.append(func.__name__)
+        return concat(pieces, keys=keys, axis=1)
+
     keys = rows + cols
 
     values_passed = values is not None
@@ -68,7 +85,6 @@ def pivot_table(data, values=None, rows=None, cols=None, aggfunc=np.mean,
         data = data[keys + values]
 
     grouped = data.groupby(keys)
-
     agged = grouped.agg(aggfunc)
 
     table = agged
@@ -95,7 +111,6 @@ def _add_margins(table, data, values, rows=None, cols=None, aggfunc=np.mean):
         col_margin = data[rows + values].groupby(rows).agg(aggfunc)
 
         # need to "interleave" the margins
-
         table_pieces = []
         margin_keys = []
         for key, piece in table.groupby(level=0, axis=1):
@@ -104,9 +119,7 @@ def _add_margins(table, data, values, rows=None, cols=None, aggfunc=np.mean):
             table_pieces.append(piece)
             margin_keys.append(all_key)
 
-        result = table_pieces[0]
-        for piece in table_pieces[1:]:
-            result = result.join(piece)
+        result = concat(table_pieces, axis=1)
     else:
         result = table
         margin_keys = table.columns

pandas/tools/tests/test_pivot.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from pandas import DataFrame
+from pandas import DataFrame, concat
 from pandas.tools.pivot import pivot_table
 import pandas.util.testing as tm
 
@@ -60,6 +60,27 @@ def test_pivot_multi_values(self):
                                rows='A', cols=['B', 'C'], fill_value=0)
         tm.assert_frame_equal(result, expected)
 
+    def test_pivot_multi_functions(self):
+        f = lambda func: pivot_table(self.data, values=['D', 'E'],
+                                     rows=['A', 'B'], cols='C',
+                                     aggfunc=func)
+        result = f([np.mean, np.std])
+        means = f(np.mean)
+        stds = f(np.std)
+        expected = concat([means, stds], keys=['mean', 'std'], axis=1)
+        tm.assert_frame_equal(result, expected)
+
+        # margins not supported??
+        f = lambda func: pivot_table(self.data, values=['D', 'E'],
+                                     rows=['A', 'B'], cols='C',
+                                     aggfunc=func, margins=True)
+        result = f([np.mean, np.std])
+        means = f(np.mean)
+        stds = f(np.std)
+        expected = concat([means, stds], keys=['mean', 'std'], axis=1)
+        tm.assert_frame_equal(result, expected)
+
+
     def test_margins(self):
         def _check_output(res, col, rows=['A', 'B'], cols=['C']):
             cmarg = res['All'][:-1]