Applying map after transform with multiprocessing will cause OOM
MEMBER
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
2025-07-09T10:33:18Z
https://github.com/huggingface/datasets/issues/7662#issuecomment-3052117613
https://api.github.com/repos/huggingface/datasets/issues/7662
IC_kwDODunzps61655t
https://api.github.com/repos/huggingface/datasets/issues/comments/3052117613/reactions
2025-07-09T10:33:18Z
https://api.github.com/repos/huggingface/datasets/issues/comments/3052117613
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
NONE
> 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)
2025-07-09T11:48:29Z
https://github.com/huggingface/datasets/issues/7662#issuecomment-3052335751
https://api.github.com/repos/huggingface/datasets/issues/7662
IC_kwDODunzps617vKH
https://api.github.com/repos/huggingface/datasets/issues/comments/3052335751/reactions
2025-07-09T11:48:29Z
https://api.github.com/repos/huggingface/datasets/issues/comments/3052335751
https://avatars.githubusercontent.com/u/26482910?v=4
https://api.github.com/users/JunjieLl/events{/privacy}
https://api.github.com/users/JunjieLl/followers
https://api.github.com/users/JunjieLl/following{/other_user}
https://api.github.com/users/JunjieLl/gists{/gist_id}
https://github.com/JunjieLl
JunjieLl
MDQ6VXNlcjI2NDgyOTEw
https://api.github.com/users/JunjieLl/orgs
https://api.github.com/users/JunjieLl/received_events
https://api.github.com/users/JunjieLl/repos
https://api.github.com/users/JunjieLl/starred{/owner}{/repo}
https://api.github.com/users/JunjieLl/subscriptions
User
https://api.github.com/users/JunjieLl
public
MEMBER
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
2025-07-09T12:00:15Z
https://github.com/huggingface/datasets/issues/7662#issuecomment-3052383419
https://api.github.com/repos/huggingface/datasets/issues/7662
IC_kwDODunzps6176y7
https://api.github.com/repos/huggingface/datasets/issues/comments/3052383419/reactions
2025-07-09T12:01:10Z
https://api.github.com/repos/huggingface/datasets/issues/comments/3052383419
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
NONE
> 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?
2025-07-09T12:26:44Z
https://github.com/huggingface/datasets/issues/7662#issuecomment-3052479181
https://api.github.com/repos/huggingface/datasets/issues/7662
IC_kwDODunzps618SLN
https://api.github.com/repos/huggingface/datasets/issues/comments/3052479181/reactions
2025-07-09T12:26:44Z
https://api.github.com/repos/huggingface/datasets/issues/comments/3052479181
https://avatars.githubusercontent.com/u/26482910?v=4
https://api.github.com/users/JunjieLl/events{/privacy}
https://api.github.com/users/JunjieLl/followers
https://api.github.com/users/JunjieLl/following{/other_user}
https://api.github.com/users/JunjieLl/gists{/gist_id}
https://github.com/JunjieLl
JunjieLl
MDQ6VXNlcjI2NDgyOTEw
https://api.github.com/users/JunjieLl/orgs
https://api.github.com/users/JunjieLl/received_events
https://api.github.com/users/JunjieLl/repos
https://api.github.com/users/JunjieLl/starred{/owner}{/repo}
https://api.github.com/users/JunjieLl/subscriptions
User
https://api.github.com/users/JunjieLl
public
NONE
> 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.


2025-07-10T06:17:39Z
https://github.com/huggingface/datasets/issues/7662#issuecomment-3055756719
https://api.github.com/repos/huggingface/datasets/issues/7662
IC_kwDODunzps62IyWv
https://api.github.com/repos/huggingface/datasets/issues/comments/3055756719/reactions
2025-07-10T06:17:39Z
https://api.github.com/repos/huggingface/datasets/issues/comments/3055756719
https://avatars.githubusercontent.com/u/26482910?v=4
https://api.github.com/users/JunjieLl/events{/privacy}
https://api.github.com/users/JunjieLl/followers
https://api.github.com/users/JunjieLl/following{/other_user}
https://api.github.com/users/JunjieLl/gists{/gist_id}
https://github.com/JunjieLl
JunjieLl
MDQ6VXNlcjI2NDgyOTEw
https://api.github.com/users/JunjieLl/orgs
https://api.github.com/users/JunjieLl/received_events
https://api.github.com/users/JunjieLl/repos
https://api.github.com/users/JunjieLl/starred{/owner}{/repo}
https://api.github.com/users/JunjieLl/subscriptions
User
https://api.github.com/users/JunjieLl
public
### 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