DEPR/BUG: DataFrame.stack including all null rows when stacking multiple levels

[rhshadrach]

Ref: #53094 (comment)

From the link above, DataFrame.stack with dropna=False will include column combinations that did not exist in the original DataFrame; the values are necessarily all null:

>>> df = pd.DataFrame(np.arange(6).reshape(2,3), columns=pd.MultiIndex.from_tuples([('A','x'), ('A','y'), ('B','z')], names=['Upper', 'Lower']))
>>> df
Upper  A     B
Lower  x  y  z
0      0  1  2
1      3  4  5

>>> df.stack(level=[0,1], dropna=False)
   Upper  Lower
0  A      x        0.0
          y        1.0
          z        NaN
   B      x        NaN
          y        NaN
          z        2.0
1  A      x        3.0
          y        4.0
          z        NaN
   B      x        NaN
          y        NaN
          z        5.0
dtype: float64

I do not think it would be expected for this operation to essentially invent combinations of the stacked columns that did not exist in the original DataFrame. I propose we deprecate this behavior, and only level combinations that occur in the original DataFrame should be included in the result.

The current implementation of DataFrame.stack does not appear easily modified to enforce this deprecation. The following function (which is a bit ugly and could use quite some refinement) implements stack that would be enforced upon deprecation. The implementation is only for the case of multiple levels (so the input must have a MultiIndex).

Implementation of new_stack

def new_stack(df, levels):
    stack_cols = (
        df.columns
        .droplevel([k for k in range(df.columns.nlevels) if k not in levels])
    )
    _, taker = np.unique(levels, return_inverse=True)
    if len(levels) > 1:
        # Arrange columns in the order we want to take them
        ordered_stack_cols = stack_cols.reorder_levels(taker)
    else:
        ordered_stack_cols = stack_cols

    buf = []
    for idx in stack_cols.unique():
        if not isinstance(idx, tuple):
            idx = (idx,)
        # Take the data from df corresponding to this idx value
        gen = iter(idx)
        column_indexer = tuple(
            next(gen) if k in levels else slice(None)
            for k in range(df.columns.nlevels)
        )
        data = df.loc[:, column_indexer]

        # When len(levels) == df.columns.nlevels, we're stacking all columns
        # and end up with a Series
        if len(levels) < df.columns.nlevels:
            data.columns = data.columns.droplevel(levels)

        buf.append(data)

    result = pd.concat(buf)

    # Construct the correct MultiIndex by combining the input's index and
    # stacked columns.
    if isinstance(df.index, MultiIndex):
        index_levels = [level.unique() for level in df.index.levels]
    else:
        index_levels = [
            df.index.get_level_values(k).values for k in range(df.index.nlevels)
        ]
    column_levels = [
        ordered_stack_cols.get_level_values(e).unique()
        for e in range(ordered_stack_cols.nlevels)
    ]
    if isinstance(df.index, MultiIndex):
        index_codes = np.tile(df.index.codes, (1, len(result) // len(df)))
    else:
        index_codes = np.tile(np.arange(len(df)), (1, len(result) // len(df)))
    index_codes = [e for e in index_codes]
    if isinstance(stack_cols, MultiIndex):
        column_codes = ordered_stack_cols.drop_duplicates().codes
    else:
        column_codes = [np.arange(stack_cols.nunique())]
    column_codes = [np.repeat(codes, len(df)) for codes in column_codes]
    index_names = df.index.names
    column_names = list(ordered_stack_cols.names)
    result.index = pd.MultiIndex(
        levels=index_levels + column_levels,
        codes=index_codes + column_codes,
        names=index_names + column_names,
    )

    # sort result, but faster than calling sort_index since we know the order we need
    len_df = len(df)
    n_uniques = len(ordered_stack_cols.unique())
    idxs = (
        np.tile(len_df * np.arange(n_uniques), len_df)
        + np.repeat(np.arange(len_df), n_uniques)
    )
    result = result.take(idxs)

    return result

It is not as quite more performant on two/three levels with NumPy dtypes

size = 100000
df = pd.DataFrame(
    np.arange(3 * size).reshape(size, 3), 
    columns=pd.MultiIndex.from_tuples(
        [('A','x', 'a'), ('A', 'y', 'a'), ('B','z', 'b')], 
        names=['l1', 'l2', 'l3'],
    ),
    dtype="int64",
)

%timeit new_stack(df, [0, 1])
%timeit df.stack(['l1', 'l2'], sort=False)
# 10.6 ms ± 77 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
# 22.1 ms ± 28.3 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit new_stack(df, [0, 1, 2])
%timeit df.stack(['l1', 'l2', 'l3'], sort=False)
# 5.69 ms ± 36.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
# 34.1 ms ± 315 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

But And is more performant on small data (for what it's worth); this is size=1000 with the same code above:

%timeit new_stack(df, [0, 1])
%timeit df.stack(['l1', 'l2'], sort=False)
1.56 ms ± 3.25 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
5.01 ms ± 12.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit new_stack(df, [0, 1, 2])
%timeit df.stack(['l1', 'l2', 'l3'], sort=False)
908 µs ± 2.86 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
5.25 ms ± 12.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

and it more efficient when working with nullable arrays; this is with size=100000 and dtype="Int64":

%timeit new_stack(df, [0, 1, 2])
%timeit df.stack(['l1', 'l2', 'l3'], sort=False)
5.9 ms ± 27.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
127 ms ± 961 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit new_stack(df, [0, 1, 2])
%timeit df.stack(['l1', 'l2', 'l3'], sort=False)
5.93 ms ± 212 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
131 ms ± 526 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In addition, it does not coerce dtypes unnecessarily as the current implementation does (ref: #51059, #17886). This implementation would close those two issues. E.g. this is from the top example (with dropna=True) that results in float64 today:

   Upper  Lower
0  A      x        0
          y        1
   B      z        2
1  A      x        3
          y        4
   B      z        5
dtype: int64

In particular, when operating on Int64 dtypes, the current implementation results in object dtype whereas this implementation results in Int64.

cc @jorisvandenbossche @jbrockmendel @mroeschke

[rhshadrach]

For the deprecation process here, I don't see any way other than introduce a new argument that defaults to being backward compatible, deprecating the default, and then removing the argument after deprecation during 3.x.

[jorisvandenbossche]

I was initially thinking we could keep an (opt-in) option for the current behaviour with dropna=False (using a full product-wise MultiIndex), if that would be faster, but based on your experiment, that is not necessarily the case.

I didn't look in detail at the implementations, but just a note that if you try new_stack on the following dataframe, it gives a wrong result:

df = pd.DataFrame(
    np.arange(12 * size).reshape(size, 12), dtype="int64",
    columns=pd.MultiIndex.from_product([["A", "B"], ["x", "y", "z"], ["a", "b"]], names=["l1", "l2", "l3"])
)

(the calculation of idxs for the final result.take(idx) might be wrong?)

I do not think it would be expected for this operation to essentially invent combinations of the stacked columns that did not exist in the original DataFrame. I propose we deprecate this behavior, and only level combinations that occur in the original DataFrame should be included in the result.

+1

Long term, if we don't have this behaviour of introducing NAs from the stacking, I would also say that there is no need at all for a dropna keyword? (and if so, I am wondering if we could (mis)use the keyword for the deprecation)

[rhshadrach]

(the calculation of idxs for the final result.take(idx) might be wrong?)

Indeed it was, and levels=[0, 1] vs levels=[1, 0] gave the same result. These are now fixed. With the change in implementation, I've updated times above. They are slightly longer, but I think the overall picture remains the same.

I've been trying to test it more thoroughly but issues with the current implementation of stack is getting in the way. E.g. with your example df, df.stack([0, 1], sort=False) raises. Similarly, I'm seeing this (with size=3) which I think is buggy behavior:

print(df.stack([0], sort=True))
# l2     x       y       z    
# l3     a   b   a   b   a   b
#   l1                        
# 0 A    0   1   2   3   4   5
#   B    6   7   8   9  10  11
# 1 A   12  13  14  15  16  17
#   B   18  19  20  21  22  23
# 2 A   24  25  26  27  28  29
#   B   30  31  32  33  34  35

print(df.stack([0], sort=False))
# l2     x       y       z       x       y       z    
# l3     a   b   a   b   a   b   a   b   a   b   a   b
#   l1                                                
# 0 A    0   1   2   3   4   5   0   1   2   3   4   5
#   B    6   7   8   9  10  11   6   7   8   9  10  11
# 1 A   12  13  14  15  16  17  12  13  14  15  16  17
#   B   18  19  20  21  22  23  18  19  20  21  22  23
# 2 A   24  25  26  27  28  29  24  25  26  27  28  29
#   B   30  31  32  33  34  35  30  31  32  33  34  35

In any case, the following runs now without raising:

import itertools as it

size = 1000
df = pd.DataFrame(
    np.arange(12 * size).reshape(size, 12), dtype="int64",
    columns=pd.MultiIndex.from_product([["A", "B"], ["x", "y", "z"], ["a", "b"]], names=["l1", "l2", "l3"])
)

for r in [1, 2, 3]:
    for levels in it.combinations([0, 1, 2], r=r):
        tm.assert_equal(new_stack(df, levels).sort_index(), df.stack(levels, sort=True))

The implementation for new_stack is for sort=False.

[rhshadrach]

Long term, if we don't have this behaviour of introducing NAs from the stacking, I would also say that there is no need at all for a dropna keyword? (and if so, I am wondering if we could (mis)use the keyword for the deprecation)

Yes - I think we can deprecate the keyword. We could use it for this deprecation by introducing dropna=[sentinel], maybe dropna=None? This would give the user the behavior detailed here, and is perhaps an okay name because this won't drop any NAs (nor introduce them artificially!). Then we can deprecate the current default in 2.x and deprecate the argument entirely in 3.x.

[jorisvandenbossche]

Similarly, I'm seeing this (with size=3) which I think is buggy behavior:

Yes, that sort=False case (duplicating the columns) is certainly buggy ..

---

FWIW, the dataframe that I used above (with MI columns using from_product) is an example where new_stack is more significantly slower:

size = 100000

df = pd.DataFrame(
    np.arange(12 * size).reshape(size, 12), dtype="int64",
    columns=pd.MultiIndex.from_product([["A", "B"], ["x", "y", "z"], ["a", "b"]], names=["l1", "l2", "l3"])
)

In [29]: %timeit df.stack([0, 1])
28.7 ms ± 1.16 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [30]: %timeit new_stack(df, [0, 1])
89.6 ms ± 4.02 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

From a quick profile, it seems there is a lot of overhead in the concat(buf) step in your new_stack for concatting the index. This is maybe something that could be avoided if we only assign the index (data.index = new_index) after concatting and can create the overall new_index in a smarter way (just some quick thoughts, didn't check in detail).

[jorisvandenbossche]

From a quick profile, it seems there is a lot of overhead in the concat(buf) step in your new_stack for concatting the index.

I also opened a separate issue based on a quick exploration into why this step is so surprisingly slow: #53574

[rhshadrach]

Thanks @jorisvandenbossche - I plan to look into your suggestion of enhancing perf here. But given the bug fixes and overall better behavior, not only with dtypes but also not creating unnecessary nulls, would you agree that the perf hit as-is is worth it?

[rhshadrach]

@jorisvandenbossche - thanks for the suggestion! I've updated the implementation and timings inplace. For your example in #53515 (comment), I now get:

21.2 ms ± 134 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)  <- df.stack
18.8 ms ± 50.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)  <- new_stack

Edit: It also seems to be working on inputs with a MultiIndex, though this needs a lot more testing