BUG: Fix pandas-dev#61221: Exception with unstack(sort=False) and NA in index.

Author: Gabe Small

pandas/core/reshape/reshape.py

@@ -134,6 +134,11 @@ def __init__(
         self.removed_level_full = index.levels[self.level]
         if not self.sort:
             unique_codes = unique(self.index.codes[self.level])
+            # Bug Fix GH 61221
+            # The -1 in the unsorted unique codes causes for doubling and an eventual ValueError
+            # saving the NA location to be used in the repeater
+            self.na = np.where(unique_codes == -1)[0][0] if -1 in unique_codes else None
+            unique_codes = unique_codes[unique_codes != -1]  
             self.removed_level = self.removed_level.take(unique_codes)
             self.removed_level_full = self.removed_level_full.take(unique_codes)
 
@@ -381,11 +386,22 @@ def _repeater(self) -> np.ndarray:
             # 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:
-                repeater = np.insert(repeater, 0, -1)
+                if not self.sort and self.na:
+                    repeater = np.insert(repeater, self.na, -1)
+                else:
+                    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
+            if self.sort or not self.na:
+                repeater = np.arange(stride) - self.lift
+            else :
+                #move the -1 to the position at self.na
+                repeater = np.arange(stride)
+                if(self.na):
+                    repeater[self.na] = -1
+                    if(self.na + 1) < len(repeater):
+                        repeater[self.na + 1:] -= 1
 
         return repeater
 
@@ -1049,7 +1065,7 @@ def stack_reshape(
             else:
                 data.columns = default_index(len(data.columns))
         buf.append(data)
-
+ 
     if len(buf) > 0 and not frame.empty:
         result = concat(buf, ignore_index=True)
     else:

pandas/tests/frame/test_stack_unstack.py

@@ -1605,6 +1605,84 @@ def test_stack_sort_false(future_stack):
     tm.assert_frame_equal(result, expected)
 
 
+def assert_na_safe_equal(left, right):
+    """Compare DataFrames ignoring NA type differences"""
+    left = left.rename(columns={pd.NA: np.nan}, level=1)
+    right = right.rename(columns={pd.NA: np.nan}, level=1)
+    tm.assert_frame_equal(left, right, check_dtype=False)
+    
+def test_unstack_sort_false_na():
+    # GH 61221
+    levels1 = ['b','a']
+    levels2 = pd.Index([1, 2, 3, pd.NA], dtype=pd.Int64Dtype())
+    index = pd.MultiIndex.from_product([levels1, levels2], names=['level1', 'level2'])
+    df = pd.DataFrame(dict(value=range(len(index))), index=index)
+    result = df.unstack(level='level2', sort=False)
+    expected = pd.DataFrame(
+     {
+         ('value', 1): [0, 4],
+         ('value', 2): [1, 5],
+         ('value', 3): [2, 6],
+         ('value', pd.Int64Dtype().na_value): [3, 7]  
+     },
+     index=pd.Index(['b', 'a'], name='level1'),
+     columns=pd.MultiIndex.from_tuples([
+         ('value', 1), ('value', 2), ('value', 3), ('value', pd.Int64Dtype().na_value)
+     ], names=[None, 'level2'])
+ )
+    assert_na_safe_equal(result, expected)
+    levels2 = pd.Index([pd.NA, 1, 2, 3], dtype=pd.Int64Dtype())
+    index = pd.MultiIndex.from_product([levels1, levels2], names=['level1', 'level2'])
+    df = pd.DataFrame(dict(value=range(len(index))), index=index)
+    result = df.unstack(level='level2', sort=False)
+    expected = pd.DataFrame(
+     {
+         ('value', pd.Int64Dtype().na_value): [0, 4],
+         ('value', 1): [1, 5],
+         ('value', 2): [2, 6],
+         ('value', 3): [3, 7]  # Use actual pd.NA object
+     },
+     index=pd.Index(['b', 'a'], name='level1'),
+     columns=pd.MultiIndex.from_tuples([
+         ('value', pd.Int64Dtype().na_value), ('value', 1), ('value', 2), ('value', 3)
+     ], names=[None, 'level2'])
+ )
+    assert_na_safe_equal(result, expected)
+    levels2 = pd.Index([ 1, pd.NA, 2, 3], dtype=pd.Int64Dtype())
+    index = pd.MultiIndex.from_product([levels1, levels2], names=['level1', 'level2'])
+    df = pd.DataFrame(dict(value=range(len(index))), index=index)
+    result = df.unstack(level='level2', sort=False)
+    expected = pd.DataFrame(
+     {
+         ('value', 1): [0, 4],
+         ('value', pd.Int64Dtype().na_value): [1, 5],
+         ('value', 2): [2, 6],
+         ('value', 3): [3, 7]  # Use actual pd.NA object
+     },
+     index=pd.Index(['b', 'a'], name='level1'),
+     columns=pd.MultiIndex.from_tuples([
+         ('value', 1), ('value', pd.Int64Dtype().na_value), ('value', 2), ('value', 3)
+     ], names=[None, 'level2'])
+ )
+    assert_na_safe_equal(result, expected)
+    levels2 = pd.Index([3, pd.NA, 1, 2], dtype=pd.Int64Dtype())
+    index = pd.MultiIndex.from_product([levels1, levels2], names=['level1', 'level2'])
+    df = pd.DataFrame(dict(value=range(len(index))), index=index)
+    result = df.unstack(level='level2', sort=False)
+    expected = pd.DataFrame(
+     {
+         ('value', 3): [0, 4],
+         ('value', pd.Int64Dtype().na_value): [1, 5],
+         ('value', 1): [2, 6],
+         ('value', 2): [3, 7]  # Use actual pd.NA object
+     },
+     index=pd.Index(['b', 'a'], name='level1'),
+     columns=pd.MultiIndex.from_tuples([
+         ('value', 3), ('value', pd.Int64Dtype().na_value), ('value', 1), ('value', 2)
+     ], names=[None, 'level2'])
+ )
+    assert_na_safe_equal(result, expected)
+
 @pytest.mark.filterwarnings("ignore:The previous implementation of stack is deprecated")
 def test_stack_sort_false_multi_level(future_stack):
     # GH 15105