Repository navigation
Dataset.from_pandas allocates ~1 MB per Arrow chunk, Out of Memory errors #8327
Description
Describe the bug
Dataset.from_pandas memory usage scales with the chunk count of the backing Arrow array (around 1 MB per chunk).
Dataset.shuffle() followed by to_pandas() produces string columns with one chunk per row, so a 25k-row DataFrame holding ~32 MB of text allocates ~30 GB on the way back through from_pandas and the process gets killed silently. Encountered in Jupyter.
pyarrow.Table.from_pandas on the identical DataFrame completes in ~140 MB, so the amplification is likely in Dataset.
Steps to reproduce the bug
The following allocates ~30GB.
Open Activity Monitor (or equivalent), observe memory pressure when running the following code:
from datasets import Dataset
texts = ['lorem ipsum dolor sit amet consectetur ' * 34 for _ in range(25_000)] # ~1.3 KB each
ds = Dataset.from_dict({'text': texts, 'label': [0] * 25_000})
df = ds.shuffle(seed=42).to_pandas()
df['text'].array._pa_array.num_chunks # 25k, chunk per row
Dataset.from_pandas(df)Measurements:
| scenario, 25k rows × ~1.3 KB | chunks | result |
|---|---|---|
shuffle() → to_pandas() → from_pandas() |
25,000 | OOM |
| same, 6,000 rows | 6,000 | OOM (>6.5 GB) |
no shuffle → from_pandas() |
25 | OK, 139 MB |
shuffle() → flatten_indices() → from_pandas() |
25 | OK, 184 MB |
shuffled df → pa.Table.from_pandas(df) |
25,000 | OK, 141 MB |
Dataset.from_dict({'text': df['text'].tolist(), ...}) |
— | OK, 468 MB |
Expected behavior
Memory proportional to data size, independent of chunk layout (ex. pyarrow.Table.from_pandas).
Possible workaround is to flatten_indices() before to_pandas(), or from_dict with .tolist() columns.
Environment info
datasets5.0.0,pandas3.0.4,pyarrow24.0.0, Python 3.12.13- macOS 14.6, Apple M3 Pro (arm64), 36 GB RAM
Activity
SuryanshSS1011 commentedon Jul 16, 2026
The measurements in the issue point at
from_pandas, but I think the allocation is actually in fingerprinting, and not in the Arrow conversion at all.On datasets 5.0.0 I get:
pa.Table.from_pandas(df)Dataset.from_pandas(df)generate_fingerprint(ds)aloneSo ~99.8% of it is
generate_fingerprint.To confirm it tracks chunk count and not data size, I held the data constant at 600 rows and only varied how many chunks it was split into:
generate_fingerprintpeakgenerate_fingerprintatfingerprint.py:235hashesdataset.__dict__, which includes_data, and the Hasher's dill serialization of a ChunkedArray scales with the number of chunks.shuffle().to_pandas()gives one chunk per row, so the chunk count explodes and the hash goes with it.I also checked
IndexedTableMixin.__init__(0.4 MB) andtable.cast(0.7 MB), so neither of those is involved.@lhoestq as I work on a PR containing the fix for this, I would love to know if there is a direction you prefer, and what you think about this in general? The two directions I see are making the table hash chunk-count-independent, or calling
combine_chunks()before hashing. Since fingerprints key the cache, I didn't want to pick that myself.lhoestq commentedon Jul 16, 2026
Hi ! making the table hash chunk-count-independent makes sense and would avoid the extra copy from
combine_chunks(). Though at the same time I understand it's maybe not as easy to do, let me know how this goes if you plan to look into it :)SuryanshSS1011 commentedon Jul 22, 2026
@lhoestq I looked into it and implemented the discussed approach in #8339. It takes the repro from a ~1.5 GB allocation down to a few MB. One thing worth flagging is that it changes fingerprint values once, so it invalidates existing caches. More details are in the PR body. Would love your read on this when you have a moment.