Weird memory usage of shuffling in `tf.data.Dataset` 
Hi massquantity,

After reviewing your issue and the code you provided, I believe I have found a solution that can help reduce the memory usage during data shuffling in TensorFlow.

The key to reducing memory usage is to explicitly set the buffer size to a reasonable value. To resolve this issue, you can explicitly set the buffer_size to a smaller value that is appropriate for your use case. As a rule of thumb, the buffer size should be set to a value that is small enough to fit into memory but large enough to provide good shuffling performance. A buffer size of 1024, as you mentioned, may be too small and may result in poor shuffling performance. A buffer size of 10000 is generally a good default value.

you can use the prefetch() method in conjunction with shuffle(). This method allows you to overlap the shuffling process with other parts of your data pipeline, which can help reduce memory usage. 

Here's an updated version of your code that incorporates these changes:

```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 = tf_dataset.shuffle(buffer_size=10000).prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
tf_dataset = iter(tf_dataset)

print(next(tf_dataset))
```
**Memory Usage**
![Figure_3](https://github.com/tensorflow/tensorflow/assets/83278020/1bfa294d-0791-4f4d-9763-061f454e1d7a)

I tested this code on my machine and found that it reduces memory usage to less than 500MiB.
Please let me know if this solution works for you. 

Best regards,
Nitya Thanks to Nitya, for your kind and detailed explanation! I'm also aware that the buffer_size should be set as large as possible for shuffling performance. `buffer_size=1024` is just an example, and what I'm seeking in this issue is some clarification on the memory usage of buffer_size. 

When a user set `buffer_size=1024`, she expects the underlying code will only use the memory of 1024 element rather than the memory of equivalent to 10240 elements, which is a "surprise" we generally want to avoid when using an open source library like TensorFlow. That's why I'm surprised when seeing the full shuffling of 50,000,000 `int` data cost 10 GB of memory.

If the code uses more memory than the specified buffer_size, that's totally OK to me, provided that the doc is clear about this. But the [shuffle doc](https://www.tensorflow.org/api_docs/python/tf/data/Dataset#shuffle) says "This dataset fills a buffer with buffer_size elements, then randomly samples elements from this buffer, replacing the selected elements with new elements."  From my understanding, it only uses `buffer_size` elements of memory. @nitya-khuntia 

If the `.shuffle(10_000)` operation increases memory usage by 500 megabyte then each number is using up around 50 kilobyte each. They should not. If we assume they are 64 bit numbers then they should be using 8 byte plus overhead. If the overhead is 50'000 byte minus 8 byte then that seems like a lot. It's about 6000 x the actual data. I did a small benchmark to find the actual value of the shuffle memory overhead in a more realistic scenario: took 1000 samples of 1,69 MB each and tested different shuffle buffer sizes to see the delta memory/buffer_size ratio which turned out to be pretty constant (the first data point is at 10):

![image](https://github.com/tensorflow/tensorflow/assets/11353865/78953b37-82d2-4e96-acff-265efad81a14)

So, as a general rule of thumb, if you do a shuffle, plan to have 8.4 (!) times your buffer_size memory available. 😉 


<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

No

### Source

binary

### Tensorflow Version

tf 2.12

### Custom Code

No

### OS Platform and Distribution

Ubuntu 22.04

### Mobile device

_No response_

### Python version

3.8

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

Hi, I've read the [tf.data doc](https://www.tensorflow.org/guide/data#randomly_shuffling_input_data), which says using a large `buffer size` in data shuffling is not recommended, but still the shuffling behavior costs way more memory than I would expect. For example, I expect the full shuffling of 50,000,000 `int` data may only use 1 GB of memory, but the following code after `print` essentially uses 10 GB.

This leads to some practical concerns. If I set `buffer_size=1024` in data shuffling, would the *actual* memory usage of the buffer size be 10 times that of 1024 elements?

### Standalone code to reproduce the issue

```shell
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))
print(next(tf_dataset))
```


### Relevant log output

```shell
<tf.Tensor: shape=(), dtype=int64, numpy=24774043>
```
</details>