Applying map after transform with multiprocessing will cause OOM
Opened by JunjieLl (open)

### Describe the bug

I have a 30TB dataset. When I perform add_column and cast_column operations on it and then execute a multiprocessing map, it results in an OOM (Out of Memory) error. However, if I skip the add_column and cast_column steps and directly run the map, there is no OOM. After debugging step by step, I found that the OOM is caused at this point, and I suspect it’s because the add_column and cast_column operations are not cached, which causes the entire dataset to be loaded in each subprocess, leading to the OOM. The critical line of code is: https://github.com/huggingface/datasets/blob/e71b0b19d79c7531f9b9bea7c09916b5f6157f42/src/datasets/utils/py_utils.py#L607


Note num_process=1 would not cause OOM. I'm confused.

### Steps to reproduce the bug

For reproduce, you can load  dataset and set cache_dir (for caching): amphion/Emilia-Dataset which is a veru large datasets that RAM can not fits.
And apply the map with multiprocessing after a transform operation  (e.g. add_column, cast_column).
As long as num_process>1, it must cause OOM.

### Expected behavior

It should not cause OOM.

### Environment info

- `datasets` version: 3.6.0
- Platform: Linux-5.10.134-16.101.al8.x86_64-x86_64-with-glibc2.35
- Python version: 3.10.12
- `huggingface_hub` version: 0.33.1
- PyArrow version: 20.0.0
- Pandas version: 2.3.0
- `fsspec` version: 2024.6.1

Comments:
Hi ! `add_column` loads the full column data in memory:

https://github.com/huggingface/datasets/blob/bfa497b1666f4c58bd231c440d8b92f9859f3a58/src/datasets/arrow_dataset.py#L6021-L6021

a workaround to add the new column is to include the new data in the map() function instead, which only loads one batch at a time

> Hi ! `add_column` loads the full column data in memory:
> 
> [datasets/src/datasets/arrow_dataset.py](https://github.com/huggingface/datasets/blob/bfa497b1666f4c58bd231c440d8b92f9859f3a58/src/datasets/arrow_dataset.py#L6021-L6021)
> 
> Line 6021 in [bfa497b](/huggingface/datasets/commit/bfa497b1666f4c58bd231c440d8b92f9859f3a58)
> 
>  column_table = InMemoryTable.from_pydict({name: column}, schema=pyarrow_schema) 
> a workaround to add the new column is to include the new data in the map() function instead, which only loads one batch at a time


How about cast_column，since map cannot apply type transformation, e.g. Audio(16000) to Audio(24000)

cast_column calls `pyarrow.Table.cast` on the full dataset which I believe the memory usage depends on the source and target types but should be low in general

casting from Audio(16000) to Audio(24000) is cheap since the source and target arrow types are the same

> cast_column calls `pyarrow.Table.cast` on the full dataset which I believe the memory usage depends on the source and target types but should be low in general
> 
> casting from Audio(16000) to Audio(24000) is cheap since the source and target arrow types are the same

Thanks for replying. So the OOM is caused by add_column operation. When I skip the operation, low memory will be achieved. Right?

> Hi ! `add_column` loads the full column data in memory:
> 
> [datasets/src/datasets/arrow_dataset.py](https://github.com/huggingface/datasets/blob/bfa497b1666f4c58bd231c440d8b92f9859f3a58/src/datasets/arrow_dataset.py#L6021-L6021)
> 
> Line 6021 in [bfa497b](/huggingface/datasets/commit/bfa497b1666f4c58bd231c440d8b92f9859f3a58)
> 
>  column_table = InMemoryTable.from_pydict({name: column}, schema=pyarrow_schema) 
> a workaround to add the new column is to include the new data in the map() function instead, which only loads one batch at a time


Note num_process=1 would not cause OOM. I'm confused.