fix(DataFrame.unstack): fix bug when indexes contains nan

Fix bux when indexes contains `nan` and is not sorting would raise an
`IndexError` or `ValueError`.

Author: Alvaro-Kothe

pandas/core/reshape/reshape.py

@@ -128,8 +128,11 @@ def __init__(
 
         self.level = self.index._get_level_number(level)
 
-        # when index includes `nan`, need to lift levels/strides by 1
-        self.lift = 1 if -1 in self.index.codes[self.level] else 0
+        # `nan` values have code `-1`, when sorting, we lift to assign them
+        # at index 0
+        self.has_nan = -1 in self.index.codes[self.level]
+        should_lift = self.has_nan and self.sort
+        self.lift = 1 if should_lift else 0
 
         # Note: the "pop" below alters these in-place.
         self.new_index_levels = list(self.index.levels)
@@ -138,8 +141,22 @@ def __init__(
         self.removed_name = self.new_index_names.pop(self.level)
         self.removed_level = self.new_index_levels.pop(self.level)
         self.removed_level_full = index.levels[self.level]
+        self.unique_nan_index: int = -1
         if not self.sort:
             unique_codes = unique(self.index.codes[self.level])
+            if self.has_nan:
+                # drop nan codes, because they are not represented in level
+                nan_mask = unique_codes == -1
+
+                if TYPE_CHECKING:
+                    # make explicit that nan_mask is an array
+                    # to remove this pyright diagnostic:
+                    # The method "__invert__" in class "bool" is deprecated
+                    nan_mask = cast(ArrayLike, nan_mask)
+
+                unique_codes = unique_codes[~nan_mask]
+                self.unique_nan_index = np.flatnonzero(nan_mask)[0]
+
             self.removed_level = self.removed_level.take(unique_codes)
             self.removed_level_full = self.removed_level_full.take(unique_codes)
 
@@ -210,7 +227,7 @@ def _make_selectors(self) -> None:
         ngroups = len(obs_ids)
 
         comp_index = ensure_platform_int(comp_index)
-        stride = self.index.levshape[self.level] + self.lift
+        stride = self.index.levshape[self.level] + self.has_nan
         self.full_shape = ngroups, stride
 
         selector = self.sorted_labels[-1] + stride * comp_index + self.lift
@@ -362,13 +379,13 @@ def get_new_values(self, values, fill_value=None):
 
     def get_new_columns(self, value_columns: Index | None):
         if value_columns is None:
-            if self.lift == 0:
+            if not self.has_nan:
                 return self.removed_level._rename(name=self.removed_name)
 
             lev = self.removed_level.insert(0, item=self.removed_level._na_value)
             return lev.rename(self.removed_name)
 
-        stride = len(self.removed_level) + self.lift
+        stride = len(self.removed_level) + self.has_nan
         width = len(value_columns)
         propagator = np.repeat(np.arange(width), stride)
 
@@ -401,12 +418,18 @@ def _repeater(self) -> np.ndarray:
         if len(self.removed_level_full) != len(self.removed_level):
             # In this case, we remap the new codes to the original level:
             repeater = self.removed_level_full.get_indexer(self.removed_level)
-            if self.lift:
+            if self.has_nan:
+                # insert nan index at first position
                 repeater = np.insert(repeater, 0, -1)
         else:
             # Otherwise, we just use each level item exactly once:
-            stride = len(self.removed_level) + self.lift
+            stride = len(self.removed_level) + self.has_nan
             repeater = np.arange(stride) - self.lift
+            if self.has_nan and self.lift == 0:
+                # assign -1 where should be nan according to the unique values.
+                repeater[self.unique_nan_index] = -1
+                # compensate for the removed index level
+                repeater[self.unique_nan_index + 1 :] -= 1
 
         return repeater
 

pandas/tests/frame/test_stack_unstack.py

@@ -1386,6 +1386,46 @@ def test_unstack_sort_false(frame_or_series, dtype):
     tm.assert_frame_equal(result, expected)
 
 
+@pytest.mark.parametrize(
+    "levels2, expected_columns, expected_data",
+    [
+        (
+            Index([None, 1, 2, 3]),
+            [("value", np.nan), ("value", 1.0), ("value", 2.0), ("value", 3.0)],
+            [[0, 4], [1, 5], [2, 6], [3, 7]],
+        ),
+        (
+            Index([1, None, 2, 3]),
+            [("value", 1.0), ("value", np.nan), ("value", 2.0), ("value", 3.0)],
+            [[0, 4], [1, 5], [2, 6], [3, 7]],
+        ),
+        (
+            Index([1, 2, None, 3]),
+            [("value", 1.0), ("value", 2.0), ("value", np.nan), ("value", 3.0)],
+            [[0, 4], [1, 5], [2, 6], [3, 7]],
+        ),
+        (
+            Index([1, 2, 3, None]),
+            [("value", 1.0), ("value", 2.0), ("value", 3.0), ("value", np.nan)],
+            [[0, 4], [1, 5], [2, 6], [3, 7]],
+        ),
+    ],
+    ids=["nan=first", "nan=second", "nan=third", "nan=last"],
+)
+def test_unstack_sort_false_nan(levels2, expected_columns, expected_data):
+    # GH#61221
+    levels1 = ["b", "a"]
+    index = MultiIndex.from_product([levels1, levels2], names=["level1", "level2"])
+    df = DataFrame({"value": [0, 1, 2, 3, 4, 5, 6, 7]}, index=index)
+    result = df.unstack(level="level2", sort=False)
+    expected = DataFrame(
+        dict(zip(expected_columns, expected_data)),
+        index=Index(["b", "a"], name="level1"),
+        columns=MultiIndex.from_tuples(expected_columns, names=[None, "level2"]),
+    )
+    tm.assert_frame_equal(result, expected)
+
+
 def test_unstack_fill_frame_object():
     # GH12815 Test unstacking with object.
     data = Series(["a", "b", "c", "a"], dtype="object")