diff --git a/pandas/core/categorical.py b/pandas/core/categorical.py index afae11163b0dca..019aa10d7df3b0 100644 --- a/pandas/core/categorical.py +++ b/pandas/core/categorical.py @@ -127,6 +127,8 @@ def maybe_to_categorical(array): """ coerce to a categorical if a series is given """ if isinstance(array, (ABCSeries, ABCCategoricalIndex)): return array._values + elif isinstance(array, np.ndarray): + return Categorical(array) return array diff --git a/pandas/core/internals.py b/pandas/core/internals.py index 4416a9736d1d9b..7e8d186c80decc 100644 --- a/pandas/core/internals.py +++ b/pandas/core/internals.py @@ -4082,7 +4082,7 @@ def unstack(self, unstacker): new_columns = dummy.get_new_columns() new_index = dummy.get_new_index() new_blocks = [] - mask_blocks = np.zeros_like(new_columns, dtype=bool) + mask_columns = np.zeros_like(new_columns, dtype=bool) for blk in self.blocks: bunstacker = unstacker( @@ -4090,12 +4090,11 @@ def unstack(self, unstacker): new_items = bunstacker.get_new_columns() new_values, mask = bunstacker.get_new_values() new_placement = new_columns.get_indexer(new_items) - mask_blocks[new_placement] = mask.any(0) + mask_columns[new_placement] = mask.any(0) new_blocks.extend(blk._unstack(new_values.T, new_placement)) - new_blocks = [b for b, keep in zip(new_blocks, mask_blocks) if keep] - new_axes = [new_columns[mask_blocks], new_index] - bm = BlockManager(new_blocks, new_axes) + bm = BlockManager(new_blocks, [new_columns, new_index]) + bm = bm.take(mask_columns.nonzero()[0], axis=0) return bm