PERF: wide_to_long taking a long time, but I found why !

[Simon-Free]

Pandas version checks

• I have checked that this issue has not already been reported.

• I have confirmed this issue exists on the latest version of pandas.

• I have confirmed this issue exists on the main branch of pandas.

Reproducible Example

When wide_to_long has a large number of columns passed as i (my dataframe has 400000 lines, 50 cols of index of multiple types), the function can take a lot of time. I found that doing something like:

1. Separate i_dataframe, and stubname_dataframe. Synchronize index between the two.
2. Reset index on stubname_dataframe
3. Perform wide_to_long on stubname_dataframe, with i=index, to obtain long_dataframe
4. Merge i_dataframe on long_dataframe using index.

Installed Versions

Replace this line with the output of pd.show_versions() --> getting an AssertionError, but using pandas 1.3.5. Though, I checked the changelogs, no performances changes were made on pd.wide_to_long between 1.3.5 and 1.5.0

Prior Performance

No response

[mroeschke]

Thanks for the issue. Could you provide:

1. A before/after timing of the change?
2. The specific lines of codes that would need changing?
3. Also update the issue title what specifically is slow with wide_to_long?

[Simon-Free]

I will try to make a nice clean example this weekend :)

[phofl]

Closing for now, please ping to reopen, when you can provide something reproducible

[samukweku]

pinging @phofl @mroeschke to reopen this:

np.random.seed(123)

df = pd.DataFrame({"A1970" : {0 : "a", 1 : "b", 2 : "c"},

                   "A1980" : {0 : "d", 1 : "e", 2 : "f"},

                   "B1970" : {0 : 2.5, 1 : 1.2, 2 : .7},

                   "B1980" : {0 : 3.2, 1 : 1.3, 2 : .1},

                   "X"     : dict(zip(range(3), np.random.randn(3)))

                  })

df = pd.concat([df]*10_000, ignore_index=True)
df["id"] = df.index
df
      A1970 A1980  B1970  B1980         X     id
0         a     d    2.5    3.2 -1.085631      0
1         b     e    1.2    1.3  0.997345      1
2         c     f    0.7    0.1  0.282978      2
3         a     d    2.5    3.2 -1.085631      3
4         b     e    1.2    1.3  0.997345      4
...     ...   ...    ...    ...       ...    ...
29995     b     e    1.2    1.3  0.997345  29995
29996     c     f    0.7    0.1  0.282978  29996
29997     a     d    2.5    3.2 -1.085631  29997
29998     b     e    1.2    1.3  0.997345  29998
29999     c     f    0.7    0.1  0.282978  29999

[30000 rows x 6 columns]

pd.wide_to_long(df, ["A", "B"], i="id", j="year")
Out[74]:
                   X  A    B
id    year
0     1970 -1.085631  a  2.5
1     1970  0.997345  b  1.2
2     1970  0.282978  c  0.7
3     1970 -1.085631  a  2.5
4     1970  0.997345  b  1.2
...              ... ..  ...
29995 1980  0.997345  e  1.3
29996 1980  0.282978  f  0.1
29997 1980 -1.085631  d  3.2
29998 1980  0.997345  e  1.3
29999 1980  0.282978  f  0.1

[60000 rows x 3 columns]

 %timeit pd.wide_to_long(df, ["A", "B"], i="id", j="year")
56.6 ms ± 240 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%%timeit
ff = df.set_index(['id','X'])
ff.columns = pd.MultiIndex.from_frame(ff.columns.str.extract(r"(.)(.+)"), names=[None, 'year'])
ff = ff.stack(level='year', future_stack=True)
4.07 ms ± 62.4 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

I believe there is room for optimisation within wide_to_long as well as extending its flexibility.

[samukweku]

Pinging @rhshadrach @MarcoGorelli @lukemanley

[MarcoGorelli]

Hi - there's a lot of open issues, but if the suggested fix passes all tests and improves performance, you may want to directly submit a PR