[BUG] DataFrame.stack with MultiIndex columns does not preserve order

[amanlai]

Describe the bug

A call to .stack() on a DataFrame with MultiIndex columns by some level results in a DataFrame whose columns are sorted lexicographically (does not preserve the original order).

The following constructor:

df = cudf.DataFrame({('b', 3): [1, 2], ('a', 2): [3, 4]})

creates a DataFrame that looks like

   b  a
   3  2
0  1  3
1  2  4

Stacking it using df.stack() produces

       a    b
0 2  NaN  3.0
  3  1.0  NaN
1 2  NaN  4.0
  3  2.0  NaN

Since the original column order was 'b' and 'a', the expected output is

       b    a
0 2  NaN  3.0
  3  1.0  NaN
1 2  NaN  4.0
  3  2.0  NaN

This is reproducible on Google Colab using the Rapids installer.

On a related note, I believe stack() in pandas behave like this in older versions (e.g. 1.5) but it preserves the original order in newer versions (e.g. 2.2).

[shwina]

Thanks for reporting, @amanlai! cc: @galipremsagar any chance this is fixed in the upcoming 2.2 compatibility changes?

[galipremsagar]

Yes, it is fixed in 24.04:

In [1]: import cudf


In [2]: df = cudf.DataFrame({('b', 3): [1, 2], ('a', 2): [3, 4]})

In [3]: 

In [3]: df
Out[3]: 
   b  a
   3  2
0  1  3
1  2  4

In [4]: df.stack()
/nvme/0/pgali/envs/cudfdev/lib/python3.10/site-packages/cudf/core/dataframe.py:6848: FutureWarning: The previous implementation of stack is deprecated and will be removed in a future version of cudf. Specify future_stack=True to adopt the new implementation and silence this warning.
  warnings.warn(
Out[4]: 
        b     a
0 2  <NA>     3
  3     1  <NA>
1 2  <NA>     4
  3     2  <NA>

In [5]: df.to_pandas().stack()
<ipython-input-5-2234f1c7afa9>:1: FutureWarning: The previous implementation of stack is deprecated and will be removed in a future version of pandas. See the What's New notes for pandas 2.1.0 for details. Specify future_stack=True to adopt the new implementation and silence this warning.
  df.to_pandas().stack()
Out[5]: 
       b    a
0 2  NaN  3.0
  3  1.0  NaN
1 2  NaN  4.0
  3  2.0  NaN

[shwina]

Fantastic! @amanlai our nightly packages should contain these fixes if you're interested in trying them out! I'm closing this issue for now, but please feel free to reopen if you require further help!