PERF: DataFrame fragmentation when calling DataFrame.assign() with large kwargs

[y-koj]

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 running the following code, pandas emits PerformanceWarning which describes DataFrame fragmentation.

import pandas as pd
import numpy as np
from timeit import timeit

N = 1_000_000
df = pd.DataFrame({"x": np.random.rand(N)})

# this part causes fragmentation
dict = {
    "x_" + str(key): np.random.rand(N) for key in range(100)
}
df = df.assign(**dict)

# defragmented DataFrame
df2 = df.copy()

print(timeit(lambda: df.sort_values(by="x_99"), number=10))  # 4.16105025 seconds
print(timeit(lambda: df2.sort_values(by="x_99"), number=10))  # 2.6979567919999994 seconds

Note

I guess this problem is caused by the for loop in DataFrame.assign() implementation:

pandas/pandas/core/frame.py

Lines 5238 to 5240 in 0691c5c

	for k, v in kwargs.items():
	data[k] = com.apply_if_callable(v, data)
	return data

Full program output

$ python pd.py
/Users/yk/tmp/pd.py:11: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
  df = df.assign(**dict)
4.16105025
2.6979567919999994

Installed Versions

INSTALLED VERSIONS

commit : 0691c5c
python : 3.9.6
python-bits : 64
OS : Darwin
OS-release : 24.0.0
Version : Darwin Kernel Version 24.0.0: Tue Sep 24 23:38:45 PDT 2024; root:xnu-11215.1.12~1/RELEASE_ARM64_T8122
machine : arm64
processor : arm
byteorder : little
LC_ALL : None
LANG : en_US.UTF-8
LOCALE : None.UTF-8

pandas : 2.2.3
numpy : 2.0.2
pytz : 2024.2
dateutil : 2.9.0.post0
pip : 24.2
Cython : None
sphinx : None
IPython : None
adbc-driver-postgresql: None
adbc-driver-sqlite : None
bs4 : None
blosc : None
bottleneck : None
dataframe-api-compat : None
fastparquet : None
fsspec : 2024.9.0
html5lib : None
hypothesis : None
gcsfs : None
jinja2 : 3.1.4
lxml.etree : None
matplotlib : 3.9.2
numba : None
numexpr : None
odfpy : None
openpyxl : None
pandas_gbq : None
psycopg2 : None
pymysql : None
pyarrow : None
pyreadstat : None
pytest : None
python-calamine : None
pyxlsb : None
s3fs : None
scipy : 1.13.1
sqlalchemy : None
tables : None
tabulate : None
xarray : None
xlrd : None
xlsxwriter : None
zstandard : None
tzdata : 2024.2
qtpy : None
pyqt5 : None

Prior Performance

No response

[samukweku]

The error message points you a bit in the right direction:

PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`

A better approach (to avoid the fragmentation) as suggested is to concatenate:

pd.concat([_dict, df],axis=1)

[y-koj]

The error message points you a bit in the right direction:

That's right.
But in my opinion, df.assign() will be more useful if it concatenates kwargs internally to avoid fragmentation. Although there is concern about overheads for small kwargs, it worth considering.

[samukweku]

Cc @mroeschke @phofl for your thoughts

[rhshadrach]

I'm positive on using concat in the general case.

Although there is concern about overheads for small kwargs, it worth considering.

For len(kwargs) == 1, we should not use concat. For len(kwargs) == 3, I'm seeing the perf of setattr vs concat quite narrow already (6%). For length 2, I'm seeing 34%, and that this appears to be independent of the number of columns in the input DataFrame (due to CoW?).

[phofl]

concat will keep the columns separate with CoW, you still have to copy. That makes it quite cheap but won't help much with fragmentation