diff --git a/pandas/tools/pivot.py b/pandas/tools/pivot.py index 89fe9463282b6..a4a175fb75716 100644 --- a/pandas/tools/pivot.py +++ b/pandas/tools/pivot.py @@ -159,6 +159,20 @@ def _add_margins(table, data, values, rows, cols, aggfunc): grand_margin = _compute_grand_margin(data, values, aggfunc) + # categorical index or columns will fail below when 'All' is added + # here we'll convert all categorical indices to object + def convert_categorical(ind): + _convert = lambda ind: (ind.astype('object') + if ind.dtype.name == 'category' else ind) + if isinstance(ind, MultiIndex): + return ind.set_levels([_convert(lev) for lev in ind.levels]) + else: + return _convert(ind) + + table.index = convert_categorical(table.index) + if hasattr(table, 'columns'): + table.columns = convert_categorical(table.columns) + if not values and isinstance(table, Series): # If there are no values and the table is a series, then there is only # one column in the data. Compute grand margin and return it. diff --git a/pandas/tools/tests/test_pivot.py b/pandas/tools/tests/test_pivot.py index 34789a3c52cb7..106e0fa7a259a 100644 --- a/pandas/tools/tests/test_pivot.py +++ b/pandas/tools/tests/test_pivot.py @@ -719,6 +719,20 @@ def test_crosstab_dropna(self): ('two', 'dull'), ('two', 'shiny')]) assert_equal(res.columns.values, m.values) + def test_categorical_margins(self): + # GH 10989 + data = pd.DataFrame({'x': np.arange(8), + 'y': np.arange(8) // 4, + 'z': np.arange(8) % 2}) + data.y = data.y.astype('category') + data.z = data.z.astype('category') + table = data.pivot_table('x', 'y', 'z', margins=True) + assert_equal(table.values, [[1, 2, 1.5], + [5, 6, 5.5], + [3, 4, 3.5]]) + + + if __name__ == '__main__': import nose nose.runmodule(argv=[__file__, '-vvs', '-x', '--pdb', '--pdb-failure'],