BUG: fix error when unstacking with sort=False when indexes contains NA

[Alvaro-Kothe]

• closes BUG: Exception with unstack(sort=False) and NA in index #61221 (Replace xxxx with the GitHub issue number)
• Tests added and passed if fixing a bug or adding a new feature
• All code checks passed.
• Added type annotations to new arguments/methods/functions.
• Added an entry in the latest doc/source/whatsnew/vX.X.X.rst file if fixing a bug or adding a new feature.

---

This pull requests makes the unstack method differentiate how it manages NA when not sorting. The main problem was that NA values when not sorting didn't have -1 value, resulting in IndexError on mask.put.

The main change is that it keeps track on where NA should be and assign it before creating the multiindex in the _repeater method.

[mroeschke]

Nit: Could you make the Index call in the body of the test since it's the same across all parameters?

[Alvaro-Kothe]

So the parametrized argument change would be

- Index([None, 1, 2, 3])
+ [None, 1, 2, 3]

or do you expect something different?

[mroeschke]

Yes correct.

[Alvaro-Kothe]

After applying the suggested change I had to alter the expected_columns values, the columns from .unstack() became integers instead of floats.

I also removed expected_data from the parametrization and put it into the test body, as it was the same for all tests.

[Alvaro-Kothe]

I am sorry, I ended up not calling Index in the test, that's why it is no longer float. But I can still reproduce the error that I was seeing in the issue with these new tests:

❯ git restore --source=main pandas/core/reshape/reshape.py

❯ pytest pandas/tests/frame/test_stack_unstack.py::test_unstack_sort_false_nan
...
FAILED pandas/tests/frame/test_stack_unstack.py::test_unstack_sort_false_nan[nan=first] - IndexError: index 8 is out of bounds for axis 0 with size 8
FAILED pandas/tests/frame/test_stack_unstack.py::test_unstack_sort_false_nan[nan=second] - IndexError: index 8 is out of bounds for axis 0 with size 8
FAILED pandas/tests/frame/test_stack_unstack.py::test_unstack_sort_false_nan[nan=third] - IndexError: index 8 is out of bounds for axis 0 with size 8
FAILED pandas/tests/frame/test_stack_unstack.py::test_unstack_sort_false_nan[nan=last] - IndexError: index 8 is out of bounds for axis 0 with size 8

[Alvaro-Kothe]

For the sake of completeness, here is the result of the reproduction from #61221 with this patch:

import pandas as pd

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)
print(df)
#                value
# level1 level2
# b      1           0
#        2           1
#        3           2
#        <NA>        3
# a      1           4
#        2           5
#        3           6
#        <NA>        7

print(df.unstack(level='level2'))
#        value
# level2  <NA>  1  2  3
# level1
# a          7  4  5  6
# b          3  0  1  2

print(df.unstack(level='level2', sort=False))
#        value
# level2     1  2  3 <NA>
# level1
# b          0  1  2    3
# a          4  5  6    7

[mroeschke]

Great! Thanks @Alvaro-Kothe