PERF: unstack (pandas-dev#43025)

Author: jbrockmendel

pandas/_libs/internals.pyi

@@ -44,6 +44,7 @@ class BlockPlacement:
     def __len__(self) -> int: ...
     def delete(self, loc) -> BlockPlacement: ...
     def append(self, others: list[BlockPlacement]) -> BlockPlacement: ...
+    def tile_for_unstack(self, factor: int) -> np.ndarray: ...
 
 class SharedBlock:
     _mgr_locs: BlockPlacement

pandas/_libs/internals.pyx

@@ -210,6 +210,27 @@ cdef class BlockPlacement:
 
         return self._as_slice
 
+    def tile_for_unstack(self, factor: int) -> np.ndarray:
+        """
+        Find the new mgr_locs for the un-stacked version of a Block.
+        """
+        cdef:
+            slice slc = self._ensure_has_slice()
+            slice new_slice
+            ndarray new_placement
+
+        if slc is not None and slc.step == 1:
+            new_slc = slice(slc.start * factor, slc.stop * factor, 1)
+            new_placement = np.arange(new_slc.start, new_slc.stop, dtype=np.intp)
+        else:
+            # Note: test_pivot_table_empty_aggfunc gets here with `slc is not None`
+            mapped = [
+                np.arange(x * factor, (x + 1) * factor, dtype=np.intp)
+                for x in self
+            ]
+            new_placement = np.concatenate(mapped)
+        return new_placement
+
 
 cdef slice slice_canonize(slice s):
     """

pandas/core/internals/managers.py

@@ -1363,10 +1363,16 @@ def unstack(self, unstacker, fill_value) -> BlockManager:
         new_blocks: list[Block] = []
         columns_mask: list[np.ndarray] = []
 
+        if len(self.items) == 0:
+            factor = 1
+        else:
+            fac = len(new_columns) / len(self.items)
+            assert fac == int(fac)
+            factor = int(fac)
+
         for blk in self.blocks:
-            blk_cols = self.items[blk.mgr_locs.indexer]
-            new_items = unstacker.get_new_columns(blk_cols)
-            new_placement = new_columns.get_indexer(new_items)
+            mgr_locs = blk.mgr_locs
+            new_placement = mgr_locs.tile_for_unstack(factor)
 
             blocks, mask = blk._unstack(
                 unstacker,

pandas/core/reshape/reshape.py

@@ -287,7 +287,7 @@ def get_new_values(self, values, fill_value=None):
 
         return new_values, new_mask
 
-    def get_new_columns(self, value_columns):
+    def get_new_columns(self, value_columns: Index | None):
         if value_columns is None:
             if self.lift == 0:
                 return self.removed_level._rename(name=self.removed_name)
@@ -308,6 +308,16 @@ def get_new_columns(self, value_columns):
             new_names = [value_columns.name, self.removed_name]
             new_codes = [propagator]
 
+        repeater = self._repeater
+
+        # The entire level is then just a repetition of the single chunk:
+        new_codes.append(np.tile(repeater, width))
+        return MultiIndex(
+            levels=new_levels, codes=new_codes, names=new_names, verify_integrity=False
+        )
+
+    @cache_readonly
+    def _repeater(self) -> np.ndarray:
         # The two indices differ only if the unstacked level had unused items:
         if len(self.removed_level_full) != len(self.removed_level):
             # In this case, we remap the new codes to the original level:
@@ -316,13 +326,10 @@ def get_new_columns(self, value_columns):
                 repeater = np.insert(repeater, 0, -1)
         else:
             # Otherwise, we just use each level item exactly once:
+            stride = len(self.removed_level) + self.lift
             repeater = np.arange(stride) - self.lift
 
-        # The entire level is then just a repetition of the single chunk:
-        new_codes.append(np.tile(repeater, width))
-        return MultiIndex(
-            levels=new_levels, codes=new_codes, names=new_names, verify_integrity=False
-        )
+        return repeater
 
     @cache_readonly
     def new_index(self):

pandas/tests/frame/test_stack_unstack.py

@@ -1285,6 +1285,21 @@ def test_stack_positional_level_duplicate_column_names():
     tm.assert_frame_equal(result, expected)
 
 
+def test_unstack_non_slice_like_blocks(using_array_manager):
+    # Case where the mgr_locs of a DataFrame's underlying blocks are not slice-like
+
+    mi = MultiIndex.from_product([range(5), ["A", "B", "C"]])
+    df = DataFrame(np.random.randn(15, 4), index=mi)
+    df[1] = df[1].astype(np.int64)
+    if not using_array_manager:
+        assert any(not x.mgr_locs.is_slice_like for x in df._mgr.blocks)
+
+    res = df.unstack()
+
+    expected = pd.concat([df[n].unstack() for n in range(4)], keys=range(4), axis=1)
+    tm.assert_frame_equal(res, expected)
+
+
 class TestStackUnstackMultiLevel:
     def test_unstack(self, multiindex_year_month_day_dataframe_random_data):
         # just check that it works for now