PERF: `.iloc[:, :len(columns)+1]` in case of multiple numeric blocks works much slower since pandas 2.1.0

[dchigarev]

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

Here in the reproducer, I'm splitting a dataframe into several even row chunks.

Case 1: I only pass a row indexer into .iloc (works fast).
Case 2: I pass a row indexer and slice(None) as a column indexer to .iloc. This overall shouldn't and doesn't change the result (also works fast).
Case 3: I pass a row indexer and a full-axis slice as a column indexer to .iloc. In my understanding, this is not a big change compared with 'case 2' and 'case 1', since I still expect to get the same result, however, in this case, performance drops dramatically starting from pandas 2.1.0

The problem is only reproducible with multiple number of numeric blocks, so I do columns concatenation beforehand in my reproducer.

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

NROWS = 30_000_000
NCOLS = 4
# concating multiple single-column dataframes to create several blocks
df = pd.concat([pd.DataFrame({f"col{i}": np.arange(NROWS)}) for i in range(NCOLS)], axis=1)

print("==== blocks info ====")
print(f"num blocks={len(df._mgr.blocks)}; {df._mgr.blocks=}")

print("\n==== df info ====")
print(df.shape) # (30_000_000, 4)
print(df.dtypes) # all np.int64

NCHUNKS = 16
row_chunksize = len(df) // NCHUNKS
print(f"{row_chunksize=}") # 1_875_000

print("\ncase 1: only access rows: .iloc[start:stop] (fast)")
t1 = timer()
parts = [
    df.iloc[i : i + row_chunksize]
    for i in range(0, len(df), row_chunksize)
]
print(f"splitting time: {timer() - t1}s.") # pandas 2.0.3 ~0.0004s. | pandas 2.1.0 ~0.0004s.
print(f"num parts: {len(parts)}") # 16

print("\ncase 2: access rows and do a none-slice on columns: .iloc[start:stop, :] (fast)")
t1 = timer()
parts = [
    df.iloc[i : i + row_chunksize, :]
    for i in range(0, len(df), row_chunksize)
]
print(f"splitting time: {timer() - t1}s.") # pandas 2.0.3 ~0.0004s. | pandas 2.1.0 ~0.0004s.
print(f"num parts: {len(parts)}") # 16

print("\ncase 3: access rows and do a full-axis slice on columns: .iloc[start:stop, :32] (super slow)")
t1 = timer()
parts = [
    df.iloc[i : i + row_chunksize, :32]
    for i in range(0, len(df), row_chunksize)
]
print(f"splitting time: {timer() - t1}s.") # pandas 2.0.3 ~0.49s. | pandas 2.1.0 ~7.15s.
print(f"num parts: {len(parts)}") # 16

Installed Versions

INSTALLED VERSIONS

commit : ba1cccd
python : 3.11.5.final.0
python-bits : 64
OS : Linux
OS-release : 5.15.0-76-generic
Version : #83-Ubuntu SMP Thu Jun 15 19:16:32 UTC 2023
machine : x86_64
processor : x86_64
byteorder : little
LC_ALL : None
LANG : en_US.UTF-8
LOCALE : en_US.UTF-8

pandas : 2.1.0
numpy : 1.26.0
pytz : 2023.3.post1
dateutil : 2.8.2
setuptools : 68.2.2
pip : 23.2.1
Cython : None
pytest : None
hypothesis : None
sphinx : None
blosc : None
feather : None
xlsxwriter : None
lxml.etree : None
html5lib : None
pymysql : None
psycopg2 : None
jinja2 : None
IPython : None
pandas_datareader : None
bs4 : None
bottleneck : None
dataframe-api-compat: None
fastparquet : None
fsspec : None
gcsfs : None
matplotlib : None
numba : None
numexpr : None
odfpy : None
openpyxl : None
pandas_gbq : None
pyarrow : None
pyreadstat : None
pyxlsb : None
s3fs : None
scipy : None
sqlalchemy : None
tables : None
tabulate : None
xarray : None
xlrd : None
zstandard : None
tzdata : 2023.3
qtpy : None
pyqt5 : None

(I verified that the numpy version was the same in both runs when I compared pandas 2.0.3 and 2.1.0)

Prior Performance

No response

[jbrockmendel]

Trying to boil down the reproducer, @dchigarev pls confirm that this is a reasonable reproducer of the original issue:

import pandas as pd
import numpy as np

NROWS = 30_000_000
NCOLS = 4
NCHUNKS = 16
row_chunksize = NROWS // NCHUNKS

df = pd.concat([pd.DataFrame({f"col{i}": np.arange(NROWS)}) for i in range(NCOLS)], axis=1)
assert df._mgr.nblocks == 4

%timeit df.iloc[ : row_chunksize, :32]
27 μs ± 80.6 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)  # <- main, numpy 2.3.4
27.2 μs ± 46.9 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)  # <- main, numpy 1.26.4
6.97 ms ± 39.7 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)  # <- 2.0.3, numpy 1.26.4
142 ms ± 3.05 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)  # <- 2.3.3, numpy 1.26.4

%timeit [df.iloc[i : i+row_chunksize, :32] for i in range(0, len(df), row_chunksize)]
439 μs ± 1.58 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)  # <- main, numpy 2.3.4
444 μs ± 4.91 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each) # <- main, numpy 1.26.4
142 ms ± 1.61 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  # <- 2.0.3, numpy 1.26.4
2.82 s ± 88.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)  # <- 2.3.3, numpy 1.26.4

@dchigarev can you confirm the issue is resolved on main?

[dchigarev]

Yep, the problem seems to be resolved on main. Closing the issue