@@ -350,13 +350,13 @@ def _check_output(result, values_col, index=['A', 'B'],
350350 # no rows
351351 rtable = self .data .pivot_table (columns = ['AA' , 'BB' ], margins = True ,
352352 aggfunc = np .mean )
353- tm . assertIsInstance ( rtable , Series )
354-
355- table = self . data . pivot_table ( index = [ 'AA' , 'BB' ], margins = True ,
356- aggfunc = 'mean' )
357- for item in [ 'DD ' , 'EE' , 'FF' ]:
358- totals = table . loc [( 'All' , '' ), item ]
359- self . assertEqual ( totals , self . data [ item ]. mean () )
353+ expected = self . data . groupby ([ 'AA' , 'BB' ]). mean ( )
354+ expected . loc [( 'All' , '' ), :] = self . data [[ 'DD' , 'EE' , 'FF' ]]. mean ()
355+ expected = ( expected . stack ()
356+ . unstack ([ 'BB' , 'AA' ] )
357+ . stack ([ 'AA ' , 'BB' ])
358+ . to_frame ())
359+ tm . assert_frame_equal ( expected , rtable )
360360
361361 # issue number #8349: pivot_table with margins and dictionary aggfunc
362362 data = [
@@ -485,8 +485,11 @@ def test_margins_no_values_no_cols(self):
485485 # Regression test on pivot table: no values or cols passed.
486486 result = self .data [['A' , 'B' ]].pivot_table (
487487 index = ['A' , 'B' ], aggfunc = len , margins = True )
488- result_list = result .tolist ()
489- self .assertEqual (sum (result_list [:- 1 ]), result_list [- 1 ])
488+ expected = self .data [['A' , 'B' ]].groupby (['A' , 'B' ]).apply (len )
489+ expected .loc [('All' , '' )] = expected .sum ()
490+ expected = expected .to_frame ()
491+
492+ tm .assert_frame_equal (result , expected )
490493
491494 def test_margins_no_values_two_rows (self ):
492495 # Regression test on pivot table: no values passed but rows are a
@@ -854,6 +857,39 @@ def test_categorical_margins(self):
854857 table = data .pivot_table ('x' , 'y' , 'z' , margins = True )
855858 tm .assert_frame_equal (table , expected )
856859
860+ def test_always_return_dataframe (self ):
861+ # GH 4386
862+ df = DataFrame ({'col1' : [3 , 4 , 5 ],
863+ 'col2' : ['C' , 'D' , 'E' ],
864+ 'col3' : [1 , 3 , 9 ]})
865+ result = df .pivot_table ('col1' , index = ['col3' , 'col2' ], aggfunc = np .sum )
866+ m = MultiIndex .from_arrays ([[1 , 3 , 9 ],
867+ ['C' , 'D' , 'E' ]],
868+ names = ['col3' , 'col2' ])
869+ expected = DataFrame ([3 , 4 , 5 ],
870+ index = m , columns = ['col1' ])
871+
872+ tm .assert_frame_equal (result , expected )
873+
874+ result = df .pivot_table (
875+ 'col1' , index = 'col3' , columns = 'col2' , aggfunc = np .sum
876+ )
877+ expected = DataFrame ([[3 , np .NaN , np .NaN ],
878+ [np .NaN , 4 , np .NaN ],
879+ [np .NaN , np .NaN , 5 ]],
880+ index = Index ([1 , 3 , 9 ], name = 'col3' ),
881+ columns = Index (['C' , 'D' , 'E' ], name = 'col2' ))
882+
883+ tm .assert_frame_equal (result , expected )
884+
885+ result = df .pivot_table ('col1' , index = 'col3' , aggfunc = [np .sum ])
886+ m = MultiIndex .from_arrays ([['sum' ],
887+ ['col1' ]])
888+ expected = DataFrame ([3 , 4 , 5 ],
889+ index = Index ([1 , 3 , 9 ], name = 'col3' ),
890+ columns = m )
891+ tm .assert_frame_equal (result , expected )
892+
857893
858894class TestCrosstab (tm .TestCase ):
859895
0 commit comments