IterableDataset's state_dict shard_example_idx is always equal to the number of samples in a shard
Opened by bruno-hays (closed)

### Describe the bug

I've noticed a strange behaviour with Iterable state_dict: the value of shard_example_idx is always equal to the amount of samples in a shard.


### Steps to reproduce the bug

I am reusing the example from the doc
```python
from datasets import Dataset

ds = Dataset.from_dict({"a": range(6)}).to_iterable_dataset(num_shards=1)
state_dict = None
# Iterate through the dataset and print examples
for idx, example in enumerate(ds):
    print(example)
    if idx == 2:
        state_dict = ds.state_dict()
        print("checkpoint")
        break
print(state_dict)
```
Returns:
```
{'a': 0}
{'a': 1}
checkpoint
{'examples_iterable': {'shard_idx': 0, 'shard_example_idx': 6, 'type': 'ArrowExamplesIterable'}, 'epoch': 0}
```


### Expected behavior

shard_example_idx should be 2 instead of 6
If we run with num_shards=2, then shard_example_idx is 3 instead of 2 and so on.

### Environment info

- `datasets` version: 3.4.1
- Platform: macOS-14.6.1-arm64-arm-64bit
- Python version: 3.12.9
- `huggingface_hub` version: 0.29.3
- PyArrow version: 19.0.1
- Pandas version: 2.2.3
- `fsspec` version: 2024.12.0

Comments:
Hey, I’d love to work on this issue but I am a beginner, can I work it with you?

Hello. I'm sorry but I don't have much time to get in the details for now.
Have you managed to reproduce the issue with the code provided ?
If you want to work on it, you can self-assign and ask @lhoestq for directions

Hi Bruno, I am trying to reproduce it this later in this week and let you know what I found.

#self-assign

Good catch, I tried and if the dataset is bigger (e.g. `range(9999)`) it returns `"shard_example_idx": 1000` with is the `config.DEFAULT_MAX_BATCH_SIZE`

https://github.com/huggingface/datasets/blob/94ccd1b4fada8a92cea96dc8df4e915041d695b6/src/datasets/arrow_dataset.py#L5313-L5317

It looks like the state_dict is incorrect in that case, it should account for this and use the `RebatchedArrowExamplesIterable` which buffers the batch of 1000 rows and counts the iteration within the batch in the state_dict


Hello @lhoestq,

I’ve been debugging the `IterableDataset.state_dict()` behavior and applied a patch to `ArrowExamplesIterable._iter_arrow()` in an attempt to fix the issue described in #7475—specifically, that `shard_example_idx` always equals the number of samples in the shard, even if only a few examples have been consumed.

### What I Tried

I updated `_iter_arrow` to slice off already-consumed rows and increment the state only by the number of actual examples yielded, like this:

```python
class ArrowExamplesIterable(_BaseExamplesIterable):
    # ... __init__ and _init_state_dict as before ...

    def _iter_arrow(self):
        shard_idx_start = self._state_dict["shard_idx"] if self._state_dict else 0

        for gen_kwargs in islice(
            _split_gen_kwargs(self.kwargs, max_num_jobs=self.num_shards),
            shard_idx_start, None
        ):
            shard_example_idx_start = self._state_dict["shard_example_idx"] if self._state_dict else 0
            shard_example_idx = 0

            for key, pa_table in self.generate_tables_fn(**gen_kwargs):
                num_rows = len(pa_table)
                next_idx = shard_example_idx + num_rows

                if next_idx <= shard_example_idx_start:
                    shard_example_idx = next_idx
                    continue

                offset = max(0, shard_example_idx_start - shard_example_idx)
                sliced_table = pa_table.slice(offset)

                if self._state_dict:
                    self._state_dict["shard_example_idx"] += len(sliced_table)

                yield key, sliced_table
                shard_example_idx = next_idx

            if self._state_dict:
                self._state_dict["shard_idx"] += 1
                self._state_dict["shard_example_idx"] = 0
```

I verified that the updated code was being used, and I added debug prints to confirm the table slicing and counter updates.

###  The Issue Still Exists

Despite the changes, the behavior remains the same. Running this minimal repro:

```python
ds = Dataset.from_dict({"a": range(6)}).to_iterable_dataset(num_shards=1)
for idx, example in enumerate(ds):
    print(example)
    if idx == 2:
        print("checkpoint")
        print(ds.state_dict())
        break
```

Still outputs:

```bash
{'a': 0}
{'a': 1}
{'a': 2}
checkpoint
{'examples_iterable': {'shard_idx': 0, 'shard_example_idx': 6, 'type': 'ArrowExamplesIterable'}, 'epoch': 0}
```

Even though only 3 examples were consumed, `shard_example_idx` jumps to 6.

###  Questions

- Could there be another place (e.g., in `__iter__`, `RebatchedArrowExamplesIterable`, or the `IterableDataset` wrapper) that's still using the old logic and overriding the state?
- Is there a better location to intercept and count yielded examples?
- Would you recommend tracking a new `true_example_idx` to avoid modifying existing behavior?

Let me know your thoughts—happy to iterate further and submit a PR once we align on the right approach. Thanks again for your help and feedback!

I found a fix using RebatchedArrowExamplesIterable, let me know if it's all good for you now

Hi @lhoestq, thanks for the quick fix and for referencing RebatchedArrowExamplesIterable! 🙌

I just tested your patch locally and can confirm that shard_example_idx is now tracking correctly when only a subset of examples is consumed. This resolves the issue I was seeing in #7475.

Really appreciate the guidance earlier on where to look—it was a great learning opportunity. If there are other parts of the IterableDataset internals that could use cleanup or testing, I’d be happy to help.