ENH: first cut at margins for pivot_table. testing still needed, #114

Author: wesm

pandas/core/index.py

@@ -1202,6 +1202,26 @@ def swaplevel(self, i, j):
         return MultiIndex(levels=new_levels, labels=new_labels,
                           names=new_names)
 
+    def reorder_levels(self, order):
+        """
+        Rearrange levels using input order. May not drop or duplicate levels
+
+        Parameters
+        ----------
+        """
+        try:
+            assert(set(order) == set(range(self.nlevels)))
+        except AssertionError:
+            raise Exception('New order must be permutation of range(%d)' %
+                            self.nlevels)
+
+        new_levels = [self.levels[i] for i in order]
+        new_labels = [self.labels[i] for i in order]
+        new_names = [self.names[i] for i in order]
+
+        return MultiIndex(levels=new_levels, labels=new_labels,
+                          names=new_names)
+
     def __getslice__(self, i, j):
         return self.__getitem__(slice(i, j))
 

pandas/tools/pivot.py

@@ -2,7 +2,7 @@
 import numpy as np
 
 def pivot_table(data, values=None, rows=None, cols=None, aggfunc=np.mean,
-                fill_value=None):
+                fill_value=None, margins=False):
     """
     Create a spreadsheet-style pivot table as a DataFrame. The levels in the
     pivot table will be stored in MultiIndex objects (hierarchical indexes) on
@@ -19,6 +19,8 @@ def pivot_table(data, values=None, rows=None, cols=None, aggfunc=np.mean,
     aggfunc : function, default numpy.mean
     fill_value : scalar, default None
         Value to replace missing values with
+    margins : boolean, default False
+        Add all row / columns (e.g. for subtotal / grand totals)
 
     Examples
     --------
@@ -59,15 +61,14 @@ def pivot_table(data, values=None, rows=None, cols=None, aggfunc=np.mean,
         else:
             values_multi = False
             values = [values]
+    else:
+        values = list(data.columns.drop(keys))
 
     if values_passed:
         data = data[keys + values]
 
     grouped = data.groupby(keys)
 
-    if values_passed and not values_multi:
-        grouped = grouped[values[0]]
-
     agged = grouped.agg(aggfunc)
 
     table = agged
@@ -77,10 +78,61 @@ def pivot_table(data, values=None, rows=None, cols=None, aggfunc=np.mean,
     if fill_value is not None:
         table = table.fillna(value=fill_value)
 
+    if margins:
+        table = _add_margins(table, data, values, rows=rows,
+                             cols=cols, aggfunc=aggfunc)
+
+    # discard the top level
+    if values_passed and not values_multi:
+        table = table[values[0]]
+
     return table
 
 DataFrame.pivot_table = pivot_table
 
+def _add_margins(table, data, values, rows=None, cols=None, aggfunc=np.mean):
+    if rows is not None:
+        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):
+            all_key = (key, 'All') + ('',) * (len(cols) - 1)
+            piece[all_key] = col_margin[key]
+            table_pieces.append(piece)
+            margin_keys.append(all_key)
+
+        result = table_pieces[0]
+        for piece in table_pieces[1:]:
+            result = result.join(piece)
+    else:
+        result = table
+        margin_keys = []
+
+    grand_margin = data[values].apply(aggfunc)
+
+    if cols is not None:
+        row_margin = data[cols + values].groupby(cols).agg(aggfunc)
+        row_margin = row_margin.stack()
+
+        # slight hack
+        new_order = [len(cols)] + range(len(cols))
+        row_margin.index = row_margin.index.reorder_levels(new_order)
+
+        key = ('All',) + ('',) * (len(rows) - 1)
+
+        row_margin = row_margin.reindex(result.columns)
+        # populate grand margin
+        for k in margin_keys:
+            row_margin[k] = grand_margin[k[0]]
+
+        margin_dummy = DataFrame(row_margin, columns=[key]).T
+        result = result.append(margin_dummy)
+
+    return result
+
 def _convert_by(by):
     if by is None:
         by = []

pandas/tools/tests/test_pivot.py

@@ -27,6 +27,9 @@ def test_pivot_table(self):
         cols=  'C'
         table = pivot_table(self.data, values='D', rows=rows, cols=cols)
 
+        table2 = self.data.pivot_table(values='D', rows=rows, cols=cols)
+        assert_frame_equal(table, table2)
+
         # this works
         pivot_table(self.data, values='D', rows=rows)
 
@@ -57,6 +60,9 @@ def test_pivot_multi_values(self):
                                rows='A', cols=['B', 'C'], fill_value=0)
         assert_frame_equal(result, expected)
 
+    def test_margins(self):
+        pass
+
 if __name__ == '__main__':
     import nose
     nose.runmodule(argv=[__file__,'-vvs','-x','--pdb', '--pdb-failure'],