`Dataset.from_pandas` allocates ~1 MB per Arrow chunk, Out of Memory errors

[nenadom]

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

• datasets 5.0.0, pandas 3.0.4, pyarrow 24.0.0, Python 3.12.13
• macOS 14.6, Apple M3 Pro (arm64), 36 GB RAM

[SuryanshSS1011]

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:

step	peak
pa.Table.from_pandas(df)	~0 MB
Dataset.from_pandas(df)	1499 MB
generate_fingerprint(ds) alone	1496 MB

So ~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:

chunks (same rows, same bytes)	generate_fingerprint peak
50	171 MB
200	670 MB
600	2071 MB

generate_fingerprint at fingerprint.py:235 hashes dataset.__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) and table.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]

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 :)