@@ -350,13 +350,13 @@ def _check_output(result, values_col, index=['A', 'B'],
350
350
# no rows
351
351
rtable = self .data .pivot_table (columns = ['AA' , 'BB' ], margins = True ,
352
352
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 )
360
360
361
361
# issue number #8349: pivot_table with margins and dictionary aggfunc
362
362
data = [
@@ -485,8 +485,11 @@ def test_margins_no_values_no_cols(self):
485
485
# Regression test on pivot table: no values or cols passed.
486
486
result = self .data [['A' , 'B' ]].pivot_table (
487
487
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 )
490
493
491
494
def test_margins_no_values_two_rows (self ):
492
495
# Regression test on pivot table: no values passed but rows are a
@@ -854,6 +857,39 @@ def test_categorical_margins(self):
854
857
table = data .pivot_table ('x' , 'y' , 'z' , margins = True )
855
858
tm .assert_frame_equal (table , expected )
856
859
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
+
857
893
858
894
class TestCrosstab (tm .TestCase ):
859
895
0 commit comments