Title: ShuffledDataSourcesArrowExamplesIterable cannot properly resume from checkpoint
State: open

### Describe the bug

ShuffledDataSourcesArrowExamplesIterable cannot properly resume from checkpoint

### Steps to reproduce the bug

1. The reproducible code is as follows：

```
from datasets import Dataset, concatenate_datasets, interleave_datasets
ds = Dataset.from_dict({"a": range(12)}).to_iterable_dataset(num_shards=1)
ds = ds.shuffle(seed=42)

for idx, example in enumerate(ds):
    print(example)
    if idx == 2: #The checkpoint can be loaded correctly only when idx <= 1.
        state_dict = ds.state_dict()
        print("checkpoint")
        break

print("state_dict: ",state_dict)

ds.load_state_dict(state_dict)
print(f"restart from checkpoint")
for example in ds:
    print(example)
```
2. The error message is as follows:
```
{'a': 0}
{'a': 7}
{'a': 6}
checkpoint
state_dict:  {'examples_iterable': {'examples_iterable': {'examples_iterable': {'shard_idx': 1, 'shard_example_idx': 0, 'type': 'ShuffledDataSourcesArrowExamplesIterable'}, 'previous_state': {'shard_idx': 0, 'shard_example_idx': 0, 'type': 'ShuffledDataSourcesArrowExamplesIterable'}, 'batch_idx': 12, 'num_chunks_since_previous_state': 12, 'cropped_chunk_length': 0, 'type': 'RebatchedArrowExamplesIterable'}, 'previous_state': {'examples_iterable': {'shard_idx': 1, 'shard_example_idx': 0, 'type': 'ShuffledDataSourcesArrowExamplesIterable'}, 'previous_state': {'shard_idx': 0, 'shard_example_idx': 0, 'type': 'ShuffledDataSourcesArrowExamplesIterable'}, 'batch_idx': 12, 'num_chunks_since_previous_state': 12, 'cropped_chunk_length': 0, 'type': 'RebatchedArrowExamplesIterable'}, 'batch_idx': 3, 'num_chunks_since_previous_state': 2, 'cropped_chunk_length': 0, 'type': 'RebatchedArrowExamplesIterable'}, 'epoch': 0}
restart from checkpoint
Loading a state dict of a shuffle buffer of a dataset without the buffer content.The shuffle buffer will be refilled before starting to yield new examples.
```
### Expected behavior

I want a correct resume from any checkpoint, but currently the checkpoint can only be loaded correctly when idx <= 1.

### Environment info

datasets Version: 4.4.1

@lhoestq 

Comments:
Hi ! As you can read in the logs, the shuffle buffer content is lost when resuming a shuffled dataset. The default size is 1000 examples, but you can tweak it

e.g. if you run your code with this

```diff
- ds = Dataset.from_dict({"a": range(12)}).to_iterable_dataset(num_shards=1)
- ds = ds.shuffle(seed=42)
+ ds = Dataset.from_dict({"a": range(100)}).to_iterable_dataset(num_shards=1)
+ ds = ds.shuffle(seed=42, buffer_size=10)
```

then you get

```
{'a': 0}
{'a': 7}
{'a': 6}
checkpoint
restart from checkpoint
Loading a state dict of a shuffle buffer of a dataset without the buffer content.The shuffle buffer will be refilled before starting to yield new examples.
{'a': 17}
{'a': 15}
{'a': 24}
{'a': 19}
{'a': 21}
{'a': 23}
...
```

where you only lose 10 rows ([1, 2, 3, 4, 5, 8, 9, 10, 11, 12])

> Hi ! As you can read in the logs, the shuffle buffer content is lost when resuming a shuffled dataset. The default size is 1000 examples, but you can tweak it
> 
> e.g. if you run your code with this
> 
> - ds = Dataset.from_dict({"a": range(12)}).to_iterable_dataset(num_shards=1)
> - ds = ds.shuffle(seed=42)
> + ds = Dataset.from_dict({"a": range(100)}).to_iterable_dataset(num_shards=1)
> + ds = ds.shuffle(seed=42, buffer_size=10)
> then you get
> 
> ```
> {'a': 0}
> {'a': 7}
> {'a': 6}
> checkpoint
> restart from checkpoint
> Loading a state dict of a shuffle buffer of a dataset without the buffer content.The shuffle buffer will be refilled before starting to yield new examples.
> {'a': 17}
> {'a': 15}
> {'a': 24}
> {'a': 19}
> {'a': 21}
> {'a': 23}
> ...
> ```
> 
> where you only lose 10 rows ([1, 2, 3, 4, 5, 8, 9, 10, 11, 12])

Thank you for your answer. So, when ShuffledDataSourcesArrowExamplesIterable resumes training, it will definitely discard unused data in buffer_size?

Yes correct. This is because the state_dict doesn't save the content of the buffer, so when resuming the buffer starts empty and the examples that were in the buffer are lost.