@@ -191,6 +191,45 @@ def test_groupby_categorical_index(self):
191
191
[0 , 1 , 2 , 3 ], levels , ordered = True ), name = 'cats' )
192
192
assert_frame_equal (result , expected )
193
193
194
+ def test_groupby_categorical_columns_index (self ):
195
+ # GH18432
196
+ s = np .random .RandomState (0 )
197
+ columns = ['A' , 'B' , 'A' , 'B' ]
198
+ categories = ['B' , 'A' ]
199
+ cat_columns = CategoricalIndex (columns ,
200
+ categories = categories ,
201
+ ordered = True )
202
+ data = s .rand (5 , len (columns ))
203
+ df = DataFrame (data , columns = cat_columns )
204
+ result = df .groupby (axis = 1 , level = 0 ).sum ()
205
+ df = DataFrame (data , columns = columns )
206
+ expected = df .groupby (axis = 1 , level = 0 ).sum ()
207
+ expected_cat_columns = CategoricalIndex (expected .columns ,
208
+ categories = categories ,
209
+ ordered = True )
210
+ expected .columns = expected_cat_columns
211
+ expected = expected .sort_index (axis = 1 )
212
+ assert_frame_equal (result , expected )
213
+
214
+ # test multi-index version
215
+ levels = [['a' , 'b' , 'c' , 'd' ],
216
+ Categorical (['A' , 'B' ], categories = categories , ordered = True )]
217
+ mixed_multi_index = MultiIndex .from_product (levels )
218
+ data = s .rand (5 , len (mixed_multi_index ))
219
+ df = DataFrame (data = data , columns = mixed_multi_index )
220
+ result = df .groupby (axis = 1 , level = 1 ).sum ()
221
+ vanilla_multi_index = MultiIndex .from_product ([
222
+ levels [0 ],
223
+ levels [1 ].get_values ()])
224
+ df = DataFrame (data = data , columns = vanilla_multi_index )
225
+ expected = df .groupby (axis = 1 , level = 1 ).sum ()
226
+ expected_cat_index = CategoricalIndex (expected .columns ,
227
+ categories = levels [1 ],
228
+ ordered = True )
229
+ expected .columns = expected_cat_index
230
+ expected = expected .sort_index (axis = 1 )
231
+ assert_frame_equal (result , expected )
232
+
194
233
def test_groupby_describe_categorical_columns (self ):
195
234
# GH 11558
196
235
cats = pd .CategoricalIndex (['qux' , 'foo' , 'baz' , 'bar' ],
0 commit comments