@@ -801,6 +801,26 @@ def test_pivot_table_margins_name_with_aggfunc_list(self):
801
801
expected = pd .DataFrame (table .values , index = ix , columns = cols )
802
802
tm .assert_frame_equal (table , expected )
803
803
804
+ def test_categorical_margins (self ):
805
+ # GH 10989
806
+ df = pd .DataFrame ({'x' : np .arange (8 ),
807
+ 'y' : np .arange (8 ) // 4 ,
808
+ 'z' : np .arange (8 ) % 2 })
809
+
810
+ expected = pd .DataFrame ([[1.0 , 2.0 , 1.5 ], [5 , 6 , 5.5 ], [3 , 4 , 3.5 ]])
811
+ expected .index = Index ([0 , 1 , 'All' ], name = 'y' )
812
+ expected .columns = Index ([0 , 1 , 'All' ], name = 'z' )
813
+
814
+ data = df .copy ()
815
+ table = data .pivot_table ('x' , 'y' , 'z' , margins = True )
816
+ tm .assert_frame_equal (table , expected )
817
+
818
+ data = df .copy ()
819
+ data .y = data .y .astype ('category' )
820
+ data .z = data .z .astype ('category' )
821
+ table = data .pivot_table ('x' , 'y' , 'z' , margins = True )
822
+ tm .assert_frame_equal (table , expected )
823
+
804
824
805
825
class TestCrosstab (tm .TestCase ):
806
826
@@ -919,26 +939,6 @@ def test_crosstab_dropna(self):
919
939
names = ['b' , 'c' ])
920
940
tm .assert_index_equal (res .columns , m )
921
941
922
- def test_categorical_margins (self ):
923
- # GH 10989
924
- df = pd .DataFrame ({'x' : np .arange (8 ),
925
- 'y' : np .arange (8 ) // 4 ,
926
- 'z' : np .arange (8 ) % 2 })
927
-
928
- expected = pd .DataFrame ([[1.0 , 2.0 , 1.5 ], [5 , 6 , 5.5 ], [3 , 4 , 3.5 ]])
929
- expected .index = Index ([0 , 1 , 'All' ], name = 'y' )
930
- expected .columns = Index ([0 , 1 , 'All' ], name = 'z' )
931
-
932
- data = df .copy ()
933
- table = data .pivot_table ('x' , 'y' , 'z' , margins = True )
934
- tm .assert_frame_equal (table , expected )
935
-
936
- data = df .copy ()
937
- data .y = data .y .astype ('category' )
938
- data .z = data .z .astype ('category' )
939
- table = data .pivot_table ('x' , 'y' , 'z' , margins = True )
940
- tm .assert_frame_equal (table , expected )
941
-
942
942
def test_crosstab_no_overlap (self ):
943
943
# GS 10291
944
944
0 commit comments