From 856d71435394a059af22b9e036f42ac28d103ed8 Mon Sep 17 00:00:00 2001 From: jbrockmendel Date: Sun, 8 Nov 2020 16:38:07 -0800 Subject: [PATCH] Backport PR #37657 on branch 1.1.x: BUG: unpickling modifies Block.ndim --- doc/source/whatsnew/v1.1.5.rst | 1 + pandas/core/internals/managers.py | 9 ++++++--- pandas/tests/io/test_pickle.py | 12 ++++++++++++ 3 files changed, 19 insertions(+), 3 deletions(-) diff --git a/doc/source/whatsnew/v1.1.5.rst b/doc/source/whatsnew/v1.1.5.rst index a122154904996..e0fa68e3b9f80 100644 --- a/doc/source/whatsnew/v1.1.5.rst +++ b/doc/source/whatsnew/v1.1.5.rst @@ -24,6 +24,7 @@ Fixed regressions Bug fixes ~~~~~~~~~ - Bug in metadata propagation for ``groupby`` iterator (:issue:`37343`) +- Bug in indexing on a :class:`Series` with ``CategoricalDtype`` after unpickling (:issue:`37631`) - .. --------------------------------------------------------------------------- diff --git a/pandas/core/internals/managers.py b/pandas/core/internals/managers.py index c9ac9cb0f140a..4c52343d08513 100644 --- a/pandas/core/internals/managers.py +++ b/pandas/core/internals/managers.py @@ -284,14 +284,17 @@ def __getstate__(self): return axes_array, block_values, block_items, extra_state def __setstate__(self, state): - def unpickle_block(values, mgr_locs): - return make_block(values, placement=mgr_locs) + def unpickle_block(values, mgr_locs, ndim: int): + # TODO(EA2D): ndim would be unnecessary with 2D EAs + return make_block(values, placement=mgr_locs, ndim=ndim) if isinstance(state, tuple) and len(state) >= 4 and "0.14.1" in state[3]: state = state[3]["0.14.1"] self.axes = [ensure_index(ax) for ax in state["axes"]] + ndim = len(self.axes) self.blocks = tuple( - unpickle_block(b["values"], b["mgr_locs"]) for b in state["blocks"] + unpickle_block(b["values"], b["mgr_locs"], ndim=ndim) + for b in state["blocks"] ) else: raise NotImplementedError("pre-0.14.1 pickles are no longer supported") diff --git a/pandas/tests/io/test_pickle.py b/pandas/tests/io/test_pickle.py index e4d43db7834e3..376091c62619b 100644 --- a/pandas/tests/io/test_pickle.py +++ b/pandas/tests/io/test_pickle.py @@ -477,3 +477,15 @@ def test_read_pickle_with_subclass(): tm.assert_series_equal(result[0], expected[0]) assert isinstance(result[1], MyTz) + + +def test_pickle_preserves_block_ndim(): + # GH#37631 + ser = pd.Series(list("abc")).astype("category").iloc[[0]] + res = tm.round_trip_pickle(ser) + + assert res._mgr.blocks[0].ndim == 1 + assert res._mgr.blocks[0].shape == (1,) + + # GH#37631 OP issue was about indexing, underlying problem was pickle + tm.assert_series_equal(res[[True]], ser)