`to_tf_dataset` consumes too much memory
MEMBER
Cc @amyeroberts @Rocketknight1 

Indded I think it's because it does something like this under the hood when there's no multiprocessing:

```python
tf_dataset = tf_dataset.shuffle(len(dataset))
```

PS: with multiprocessing it appears to be different:

```python
indices = np.arange(len(dataset))
if shuffle:
    np.random.shuffle(indices)
```
2023-05-15T09:57:03Z
https://github.com/huggingface/datasets/issues/5855#issuecomment-1547544863
https://api.github.com/repos/huggingface/datasets/issues/5855
IC_kwDODunzps5cPakf
https://api.github.com/repos/huggingface/datasets/issues/comments/1547544863/reactions
2023-05-15T09:57:03Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1547544863
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
MEMBER
Hi @massquantity, the dataset being shuffled there is not the full dataset. If you look at [the line above](https://github.com/huggingface/datasets/blob/main/src/datasets/utils/tf_utils.py#L182), the dataset is actually just a single indices array at that point, and that array is the only thing that gets fully loaded into memory and shuffled. We then load samples from the dataset by applying a transform function to the shuffled dataset, which fetches samples based on the indices it receives.

If your dataset is **really** gigantic, then this index tensor might be a memory issue, but since it's just an int64 tensor it will only use 1GB of memory per 125 million samples.

Still, if you're encountering memory issues, there might be another cause here - can you share some code to reproduce the error, or does it depend on some internal/proprietary dataset?
2023-05-15T12:20:20Z
https://github.com/huggingface/datasets/issues/5855#issuecomment-1547756997
https://api.github.com/repos/huggingface/datasets/issues/5855
IC_kwDODunzps5cQOXF
https://api.github.com/repos/huggingface/datasets/issues/comments/1547756997/reactions
2023-05-15T12:20:20Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1547756997
https://avatars.githubusercontent.com/u/12866554?v=4
https://api.github.com/users/Rocketknight1/events{/privacy}
https://api.github.com/users/Rocketknight1/followers
https://api.github.com/users/Rocketknight1/following{/other_user}
https://api.github.com/users/Rocketknight1/gists{/gist_id}
https://github.com/Rocketknight1
Rocketknight1
MDQ6VXNlcjEyODY2NTU0
https://api.github.com/users/Rocketknight1/orgs
https://api.github.com/users/Rocketknight1/received_events
https://api.github.com/users/Rocketknight1/repos
https://api.github.com/users/Rocketknight1/starred{/owner}{/repo}
https://api.github.com/users/Rocketknight1/subscriptions
User
https://api.github.com/users/Rocketknight1
public
NONE
Hi @Rocketknight1, you're right and I also noticed that only indices are used in shuffling. My data has shape (50000000, 10), but really the problem doesn't relate to a specific dataset. Simply running the following code costs me 10GB of memory.

```python
from datasets import Dataset

def gen():
    for i in range(50000000):
        yield {"data": i}

ds = Dataset.from_generator(gen, cache_dir="./huggingface")

tf_ds = ds.to_tf_dataset(
    batch_size=1,
    shuffle=True,
    drop_remainder=False,
    prefetch=True,
)
tf_ds = iter(tf_ds)
next(tf_ds)
# {'data': <tf.Tensor: shape=(1,), dtype=int64, numpy=array([0])>}
```

I just realized maybe it was an issue from tensorflow (I'm using tf 2.12). So I tried the following code, and it used 10GB of memory too.
```python
import numpy as np
import tensorflow as tf

data_size = 50000000
tf_dataset = tf.data.Dataset.from_tensor_slices(np.arange(data_size))
tf_dataset = iter(tf_dataset.shuffle(data_size))
next(tf_dataset)
# <tf.Tensor: shape=(), dtype=int64, numpy=24774043>
```

By the way, as @lhoestq mentioned, multiprocessing uses numpy shuffling, and it uses less than 1 GB of memory:
```python
tf_ds_mp = ds.to_tf_dataset(
    batch_size=1,
    shuffle=True,
    drop_remainder=False,
    prefetch=True,
    num_workers=2,
)
```
2023-05-15T13:06:47Z
https://github.com/huggingface/datasets/issues/5855#issuecomment-1547828178
https://api.github.com/repos/huggingface/datasets/issues/5855
IC_kwDODunzps5cQfvS
https://api.github.com/repos/huggingface/datasets/issues/comments/1547828178/reactions
2023-05-15T13:26:20Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1547828178
https://avatars.githubusercontent.com/u/28751760?v=4
https://api.github.com/users/massquantity/events{/privacy}
https://api.github.com/users/massquantity/followers
https://api.github.com/users/massquantity/following{/other_user}
https://api.github.com/users/massquantity/gists{/gist_id}
https://github.com/massquantity
massquantity
MDQ6VXNlcjI4NzUxNzYw
https://api.github.com/users/massquantity/orgs
https://api.github.com/users/massquantity/received_events
https://api.github.com/users/massquantity/repos
https://api.github.com/users/massquantity/starred{/owner}{/repo}
https://api.github.com/users/massquantity/subscriptions
User
https://api.github.com/users/massquantity
public
MEMBER
Thanks for that reproduction script - I've confirmed the same issue is occurring for me. Investigating it now!
2023-05-15T13:25:30Z
https://github.com/huggingface/datasets/issues/5855#issuecomment-1547857482
https://api.github.com/repos/huggingface/datasets/issues/5855
IC_kwDODunzps5cQm5K
https://api.github.com/repos/huggingface/datasets/issues/comments/1547857482/reactions
2023-05-15T13:25:30Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1547857482
https://avatars.githubusercontent.com/u/12866554?v=4
https://api.github.com/users/Rocketknight1/events{/privacy}
https://api.github.com/users/Rocketknight1/followers
https://api.github.com/users/Rocketknight1/following{/other_user}
https://api.github.com/users/Rocketknight1/gists{/gist_id}
https://github.com/Rocketknight1
Rocketknight1
MDQ6VXNlcjEyODY2NTU0
https://api.github.com/users/Rocketknight1/orgs
https://api.github.com/users/Rocketknight1/received_events
https://api.github.com/users/Rocketknight1/repos
https://api.github.com/users/Rocketknight1/starred{/owner}{/repo}
https://api.github.com/users/Rocketknight1/subscriptions
User
https://api.github.com/users/Rocketknight1
public
MEMBER
Update: The memory usage is occurring in creation of the index and shuffle buffer. You can reproduce it very simply with:

```python
import tensorflow as tf
indices = tf.range(50_000_000, dtype=tf.int64)
dataset = tf.data.Dataset.from_tensor_slices(indices)
dataset = dataset.shuffle(len(dataset))
print(next(iter(dataset))
```
When I wrote this code I thought `tf.data` had an optimization for shuffling an entire tensor that wouldn't create the entire shuffle buffer, but evidently it's just creating the enormous buffer in memory. I'll see if I can find a more efficient way to do this - we might end up moving everything to the `numpy` multiprocessing path to avoid it.
2023-05-15T13:45:15Z
https://github.com/huggingface/datasets/issues/5855#issuecomment-1547887714
https://api.github.com/repos/huggingface/datasets/issues/5855
IC_kwDODunzps5cQuRi
https://api.github.com/repos/huggingface/datasets/issues/comments/1547887714/reactions
2023-05-15T13:45:15Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1547887714
https://avatars.githubusercontent.com/u/12866554?v=4
https://api.github.com/users/Rocketknight1/events{/privacy}
https://api.github.com/users/Rocketknight1/followers
https://api.github.com/users/Rocketknight1/following{/other_user}
https://api.github.com/users/Rocketknight1/gists{/gist_id}
https://github.com/Rocketknight1
Rocketknight1
MDQ6VXNlcjEyODY2NTU0
https://api.github.com/users/Rocketknight1/orgs
https://api.github.com/users/Rocketknight1/received_events
https://api.github.com/users/Rocketknight1/repos
https://api.github.com/users/Rocketknight1/starred{/owner}{/repo}
https://api.github.com/users/Rocketknight1/subscriptions
User
https://api.github.com/users/Rocketknight1
public
MEMBER
I opened a PR to fix this - will continue the discussion there!
2023-05-15T15:36:23Z
https://github.com/huggingface/datasets/issues/5855#issuecomment-1548095696
https://api.github.com/repos/huggingface/datasets/issues/5855
IC_kwDODunzps5cRhDQ
https://api.github.com/repos/huggingface/datasets/issues/comments/1548095696/reactions
2023-05-15T15:36:23Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1548095696
https://avatars.githubusercontent.com/u/12866554?v=4
https://api.github.com/users/Rocketknight1/events{/privacy}
https://api.github.com/users/Rocketknight1/followers
https://api.github.com/users/Rocketknight1/following{/other_user}
https://api.github.com/users/Rocketknight1/gists{/gist_id}
https://github.com/Rocketknight1
Rocketknight1
MDQ6VXNlcjEyODY2NTU0
https://api.github.com/users/Rocketknight1/orgs
https://api.github.com/users/Rocketknight1/received_events
https://api.github.com/users/Rocketknight1/repos
https://api.github.com/users/Rocketknight1/starred{/owner}{/repo}
https://api.github.com/users/Rocketknight1/subscriptions
User
https://api.github.com/users/Rocketknight1
public
### Describe the bug

Hi, I'm using `to_tf_dataset` to convert a _large_ dataset to `tf.data.Dataset`. I observed that the data loading *before* training took a lot of time and memory, even with `batch_size=1`.

After some digging, i believe the reason lies in the shuffle behavior. The [source code](https://github.com/huggingface/datasets/blob/main/src/datasets/utils/tf_utils.py#L185) uses `len(dataset)` as the `buffer_size`, which may load all the data into the memory, and the [tf.data doc](https://www.tensorflow.org/guide/data#randomly_shuffling_input_data) also states that "While large buffer_sizes shuffle more thoroughly, they can take a lot of memory, and significant time to fill".

### Steps to reproduce the bug

```python
from datasets import Dataset

def gen():  # some large data
    for i in range(50000000):
        yield {"data": i}

ds = Dataset.from_generator(gen, cache_dir="./huggingface")

tf_ds = ds.to_tf_dataset(
    batch_size=64,
    shuffle=False,  # no shuffle
    drop_remainder=False,
    prefetch=True,
)

# fast and memory friendly 🤗
for batch in tf_ds: 
    ...

tf_ds_shuffle = ds.to_tf_dataset(
    batch_size=64,
    shuffle=True,
    drop_remainder=False,
    prefetch=True,
)

# slow and memory hungry for simple iteration 😱
for batch in tf_ds_shuffle: 
    ...

```

### Expected behavior

Shuffling should not load all the data into the memory. Would adding a `buffer_size` parameter in the `to_tf_dataset` API alleviate the problem?

### Environment info

- `datasets` version: 2.11.0
- Platform: Linux-5.17.1-051701-generic-x86_64-with-glibc2.17
- Python version: 3.8.13
- Huggingface_hub version: 0.13.4
- PyArrow version: 11.0.0
- Pandas version: 1.4.3
