|
6 | 6 |
|
7 | 7 | from warnings import catch_warnings
|
8 | 8 | from pandas import (date_range, Timestamp,
|
9 |
| - Index, MultiIndex, DataFrame, Series) |
| 9 | + Index, MultiIndex, DataFrame, Series, CategoricalIndex) |
10 | 10 | from pandas.util.testing import (assert_panel_equal, assert_frame_equal,
|
11 | 11 | assert_series_equal, assert_almost_equal)
|
12 | 12 | from pandas.compat import lrange, long
|
@@ -251,6 +251,29 @@ def test_groupby_levels_and_columns(self):
|
251 | 251 | by_columns.columns = pd.Index(by_columns.columns, dtype=np.int64)
|
252 | 252 | tm.assert_frame_equal(by_levels, by_columns)
|
253 | 253 |
|
| 254 | + def test_groupby_categorical_index_and_columns(self): |
| 255 | + # GH18432 |
| 256 | + columns = ['A', 'B', 'A', 'B'] |
| 257 | + categories = ['B', 'A'] |
| 258 | + data = np.ones((5, 4), int) |
| 259 | + cat_columns = CategoricalIndex(columns, |
| 260 | + categories=categories, |
| 261 | + ordered=True) |
| 262 | + df = DataFrame(data=data, columns=cat_columns) |
| 263 | + result = df.groupby(axis=1, level=0).sum() |
| 264 | + expected_data = 2 * np.ones((5, 2), int) |
| 265 | + expected_columns = CategoricalIndex(categories, |
| 266 | + categories=categories, |
| 267 | + ordered=True) |
| 268 | + expected = DataFrame(data=expected_data, columns=expected_columns) |
| 269 | + assert_frame_equal(result, expected) |
| 270 | + |
| 271 | + # test transposed version |
| 272 | + df = DataFrame(data.T, index=cat_columns) |
| 273 | + result = df.groupby(axis=0, level=0).sum() |
| 274 | + expected = DataFrame(data=expected_data.T, index=expected_columns) |
| 275 | + assert_frame_equal(result, expected) |
| 276 | + |
254 | 277 | def test_grouper_getting_correct_binner(self):
|
255 | 278 |
|
256 | 279 | # GH 10063
|
|
0 commit comments