Title: [Regression] IterableDataset is broken on 2.20.0
Author: AjayP13 (NONE)
State: closed
Created: 2024-07-31T13:01:59Z

### Describe the bug

In the latest version of datasets there is a major regression, after creating an `IterableDataset` from a generator and applying a few operations (`map`, `select`), you can no longer iterate through the dataset multiple times.

The issue seems to stem from the recent addition of "resumable IterableDatasets" (#6658) (@lhoestq). It seems like it's keeping state when it shouldn't.



### Steps to reproduce the bug

Minimal Reproducible Example (comparing `datasets==2.17.0` and `datasets==2.20.0`)

```
#!/bin/bash

# List of dataset versions to test
versions=("2.17.0" "2.20.0")

# Loop through each version
for version in "${versions[@]}"; do
    # Install the specific version of the datasets library
    pip3 install -q datasets=="$version" 2>/dev/null

    # Run the Python script
    python3 - <<EOF
from datasets import IterableDataset
from datasets.features.features import Features, Value

def test_gen():
    yield from [{"foo": i} for i in range(10)]

features = Features([("foo", Value("int64"))])
d = IterableDataset.from_generator(test_gen, features=features)
mapped = d.map(lambda row: {"foo": row["foo"] * 2})
column = mapped.select_columns(["foo"])
print("Version $version - Iterate Once:", list(column))
print("Version $version - Iterate Twice:", list(column))
EOF
done
```

The output looks like this:
```
Version 2.17.0 - Iterate Once: [{'foo': 0}, {'foo': 2}, {'foo': 4}, {'foo': 6}, {'foo': 8}, {'foo': 10}, {'foo': 12}, {'foo': 14}, {'foo': 16}, {'foo': 18}]
Version 2.17.0 - Iterate Twice: [{'foo': 0}, {'foo': 2}, {'foo': 4}, {'foo': 6}, {'foo': 8}, {'foo': 10}, {'foo': 12}, {'foo': 14}, {'foo': 16}, {'foo': 18}]
Version 2.20.0 - Iterate Once: [{'foo': 0}, {'foo': 2}, {'foo': 4}, {'foo': 6}, {'foo': 8}, {'foo': 10}, {'foo': 12}, {'foo': 14}, {'foo': 16}, {'foo': 18}]
Version 2.20.0 - Iterate Twice: []
```

### Expected behavior

The expected behavior is it version 2.20.0 should behave the same as 2.17.0.

### Environment info

`datasets==2.20.0` on any platform.

Comments:
@lhoestq I detected this regression over on [DataDreamer](https://github.com/datadreamer-dev/DataDreamer)'s test suite. I put in these [monkey patches](https://github.com/datadreamer-dev/DataDreamer/blob/4cbaf9f39cf7bedde72bbaa68346e169788fbecb/src/_patches/datasets_reset_state_hack.py) in case that fixed it our tests failing in case it helps you figure out where this is coming from. I found it hard to reason through the resumable IterableDataset code though, so hopefully you have more intuition to implement a proper fix.

I believe these lines in `TypedExamplesIterable` are responsible for stopping the re-iteration of `IterableDataset`:

https://github.com/huggingface/datasets/blob/ebec2691fb1e40145429f63375cef3f46d3011ab/src/datasets/iterable_dataset.py#L1616-L1619

In contrast to other `Iterable`s, there is no check on whether `self._state_dict` is None or not. This particular case stands out and seems less straightforward to comprehend why.  @lhoestq could you please assist us with this? Your help is much appreciated.

Thanks for reporting for investigating - your assumption was correct @VeryLazyBoy !