IterableDataset's state_dict shard_example_idx is always equal to the number of samples in a shard

[bruno-hays]

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

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

[Harry-Yang0518]

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

[bruno-hays]

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

[Harry-Yang0518]

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

[Harry-Yang0518]

#self-assign

[lhoestq]

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

datasets/src/datasets/arrow_dataset.py

Lines 5313 to 5317 in 94ccd1b

	ex_iterable = ArrowExamplesIterable(
	Dataset._generate_tables_from_shards,
	kwargs={"shards": shards, "batch_size": config.DEFAULT_MAX_BATCH_SIZE},
	)
	ds = IterableDataset(ex_iterable, info=DatasetInfo(features=self.features))

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

[Harry-Yang0518]

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:

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:

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:

{'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!

[lhoestq]

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

[Harry-Yang0518]

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.