@@ -350,13 +350,14 @@ 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 )
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 ())#this still returns a series
354
359
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 ())
360
+ tm .assert_frame_equal (expected , rtable )
360
361
361
362
# issue number #8349: pivot_table with margins and dictionary aggfunc
362
363
data = [
@@ -485,8 +486,11 @@ def test_margins_no_values_no_cols(self):
485
486
# Regression test on pivot table: no values or cols passed.
486
487
result = self .data [['A' , 'B' ]].pivot_table (
487
488
index = ['A' , 'B' ], aggfunc = len , margins = True )
488
- result_list = result .tolist ()
489
- self .assertEqual (sum (result_list [:- 1 ]), result_list [- 1 ])
489
+ expected = self .data [['A' , 'B' ]].groupby (['A' , 'B' ]).apply (len )
490
+ expected .loc [('All' , '' )] = expected .sum ()
491
+ expected = expected .to_frame ()
492
+
493
+ tm .assert_frame_equal (result , expected ) #need to return dataframe
490
494
491
495
def test_margins_no_values_two_rows (self ):
492
496
# Regression test on pivot table: no values passed but rows are a
@@ -854,6 +858,39 @@ def test_categorical_margins(self):
854
858
table = data .pivot_table ('x' , 'y' , 'z' , margins = True )
855
859
tm .assert_frame_equal (table , expected )
856
860
861
+ def test_always_return_dataframe (self ):
862
+ # GH 4386
863
+ df = DataFrame ({'col1' : [3 , 4 , 5 ],
864
+ 'col2' : ['C' , 'D' , 'E' ],
865
+ 'col3' : [1 , 3 , 9 ]})
866
+ result = df .pivot_table ('col1' , index = ['col3' , 'col2' ], aggfunc = np .sum )
867
+ m = MultiIndex .from_arrays ([[1 , 3 , 9 ],
868
+ ['C' , 'D' , 'E' ]],
869
+ names = ['col3' , 'col2' ])
870
+ expected = DataFrame ([3 , 4 , 5 ],
871
+ index = m , columns = ['col1' ])
872
+
873
+ tm .assert_frame_equal (result , expected )
874
+
875
+ result = df .pivot_table (
876
+ 'col1' , index = 'col3' , columns = 'col2' , aggfunc = np .sum
877
+ )
878
+ expected = DataFrame ([[3 , np .NaN , np .NaN ],
879
+ [np .NaN , 4 , np .NaN ],
880
+ [np .NaN , np .NaN , 5 ]],
881
+ index = Index ([1 , 3 , 9 ], name = 'col3' ),
882
+ columns = Index (['C' , 'D' , 'E' ], name = 'col2' ))
883
+
884
+ tm .assert_frame_equal (result , expected )
885
+
886
+ result = df .pivot_table ('col1' , index = 'col3' , aggfunc = [np .sum ])
887
+ m = MultiIndex .from_arrays ([['sum' ],
888
+ ['col1' ]])
889
+ expected = DataFrame ([3 , 4 , 5 ],
890
+ index = Index ([1 , 3 , 9 ], name = 'col3' ),
891
+ columns = m )
892
+ tm .assert_frame_equal (result , expected )
893
+
857
894
858
895
class TestCrosstab (tm .TestCase ):
859
896
0 commit comments