Fix IterableDataset state_dict shard_example_idx counting

[Harry-Yang0518]

Fix IterableDataset's state_dict shard_example_idx reporting

Description

This PR fixes issue #7475 where the shard_example_idx value in IterableDataset's state_dict() always equals the number of samples in a shard, even if only a few examples have been consumed.

The issue is in the _iter_arrow method of the ArrowExamplesIterable class where it updates the shard_example_idx state by the full length of the batch (len(pa_table)) even when we're only partway through processing the examples.

Changes

Modified the _iter_arrow method of ArrowExamplesIterable to:

1. Track the actual number of examples processed
2. Only increment the shard_example_idx by the number of examples actually yielded
3. Handle partial batches correctly

How to Test

I've included a simple test case that demonstrates the fix:

from datasets import Dataset

# Create a test dataset
ds = Dataset.from_dict({"a": range(6)}).to_iterable_dataset(num_shards=1)

# Iterate through part of the dataset
for idx, example in enumerate(ds):
    print(example)
    if idx == 2:  # Stop after 3 examples (0, 1, 2)
        state_dict = ds.state_dict()
        print("Checkpoint state_dict:", state_dict)
        break

# Before the fix, the output would show shard_example_idx: 6
# After the fix, it shows shard_example_idx: 3, correctly reflecting the 3 processed examples

Implementation Details

1. Added logic to track the number of examples actually seen in the current shard
2. Modified the state update to only count examples actually yielded
3. Improved handling of partial batches and skipped examples

This fix ensures that checkpointing and resuming works correctly with exactly the expected number of examples, rather than skipping ahead to the end of the batch.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[lhoestq]

Cool ! I left some comments :)

Feel free to also update your branch to include CI fixes from main

[lhoestq]

is this expected ? I think ShuffledDataSourcesArrowExamplesIterable has its own _iter_arrow() implementation

[lhoestq]

how is it different from shard_example_idx ?

[lhoestq]

is this needed ? we always yield full tables, so it's unlikely we end up with a shard_example_idx that doesn't land exactly on a table boundary (except if the dataset state is manually crafted maybe)