ENH: pass multiple values columns to pivot_table, GH #381

Author: wesm

pandas/tools/pivot.py

@@ -51,10 +51,22 @@ def pivot_table(data, values=None, rows=None, cols=None, aggfunc=np.mean,
     cols = _convert_by(cols)
 
     keys = rows + cols
+
+    values_passed = values is not None
+    if values_passed:
+        if isinstance(values, (list, tuple)):
+            values_multi = True
+        else:
+            values_multi = False
+            values = [values]
+
+    if values_passed:
+        data = data[keys + values]
+
     grouped = data.groupby(keys)
 
-    if values is not None:
-        grouped = grouped[values]
+    if values_passed and not values_multi:
+        grouped = grouped[values[0]]
 
     agged = grouped.agg(aggfunc)
 

pandas/tools/tests/test_pivot.py

@@ -19,7 +19,8 @@ def setUp(self):
                                       'dull', 'shiny', 'shiny', 'dull',
                                       'shiny', 'shiny', 'shiny'],
                                'D' : np.random.randn(11),
-                               'E' : np.random.randn(11)})
+                               'E' : np.random.randn(11),
+                               'F' : np.random.randn(11)})
 
     def test_pivot_table(self):
         rows = ['A', 'B']
@@ -46,6 +47,13 @@ def test_pivot_table_multiple(self):
         expected = self.data.groupby(rows + [cols]).agg(np.mean).unstack()
         assert_frame_equal(table, expected)
 
+    def test_pivot_multi_values(self):
+        result = pivot_table(self.data, values=['D', 'E'],
+                             rows='A', cols=['B', 'C'])
+        expected = pivot_table(self.data.drop(['F'], axis=1),
+                               rows='A', cols=['B', 'C'])
+        assert_frame_equal(result, expected)
+
 if __name__ == '__main__':
     import nose
     nose.runmodule(argv=[__file__,'-vvs','-x','--pdb', '--pdb-failure'],