@@ -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 , DataFrame )
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 .values .flatten ().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
0 commit comments