From 2b04d9f4b0016210b321e2a20abc9275d8fa8569 Mon Sep 17 00:00:00 2001 From: Jake VanderPlas Date: Fri, 4 Sep 2015 12:57:19 -0700 Subject: [PATCH 1/2] BUG: quick fix for #10989 --- pandas/tools/pivot.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/pandas/tools/pivot.py b/pandas/tools/pivot.py index 89fe9463282b6..4f500c818c06e 100644 --- a/pandas/tools/pivot.py +++ b/pandas/tools/pivot.py @@ -159,6 +159,19 @@ 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) + 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. From 74cac0e4060e9984fbd39f23bf8fe50186228368 Mon Sep 17 00:00:00 2001 From: Jake VanderPlas Date: Fri, 4 Sep 2015 13:23:31 -0700 Subject: [PATCH 2/2] TST: add test case from Issue #10989 --- pandas/tools/pivot.py | 3 ++- pandas/tools/tests/test_pivot.py | 14 ++++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/pandas/tools/pivot.py b/pandas/tools/pivot.py index 4f500c818c06e..a4a175fb75716 100644 --- a/pandas/tools/pivot.py +++ b/pandas/tools/pivot.py @@ -170,7 +170,8 @@ def convert_categorical(ind): return _convert(ind) table.index = convert_categorical(table.index) - table.columns = convert_categorical(table.columns) + 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 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'],