[bug] resume from ckpt skips samples if .map is applied
CONTRIBUTOR
Thanks for reporting this — it looks like a separate but related bug to #7538, which involved sample loss when resuming an `IterableDataset` wrapped in `FormattedExamplesIterable`. That was resolved in #7553 by re-batching the iterable to track offset correctly.

In this case, the issue seems to arise specifically from applying `.map()` before sharding and checkpointing. That wraps the iterable in `MappedExamplesIterable`, which may not preserve or propagate `shard_example_idx` correctly across `.state_dict()` and `.load_state_dict()` calls.

You can see that without `.map()`, resume works fine — but with `.map()`, it jumps from sample 9 to 50, skipping the rest of the shard.

I'll dig deeper into how `MappedExamplesIterable` manages offsets and whether it supports proper checkpoint resumption. If not, we might need a fix similar to the one in #7553, or a wrapper to preserve resume metadata.

Happy to help fix it!

2025-06-29T07:40:05Z
https://github.com/huggingface/datasets/issues/7630#issuecomment-3016382129
https://api.github.com/repos/huggingface/datasets/issues/7630
IC_kwDODunzps6zylax
https://api.github.com/repos/huggingface/datasets/issues/comments/3016382129/reactions
2025-06-29T07:51:26Z
https://api.github.com/repos/huggingface/datasets/issues/comments/3016382129
https://avatars.githubusercontent.com/u/142811259?v=4
https://api.github.com/users/ArjunJagdale/events{/privacy}
https://api.github.com/users/ArjunJagdale/followers
https://api.github.com/users/ArjunJagdale/following{/other_user}
https://api.github.com/users/ArjunJagdale/gists{/gist_id}
https://github.com/ArjunJagdale
ArjunJagdale
U_kgDOCIMgew
https://api.github.com/users/ArjunJagdale/orgs
https://api.github.com/users/ArjunJagdale/received_events
https://api.github.com/users/ArjunJagdale/repos
https://api.github.com/users/ArjunJagdale/starred{/owner}{/repo}
https://api.github.com/users/ArjunJagdale/subscriptions
User
https://api.github.com/users/ArjunJagdale
public
CONTRIBUTOR
Let me know if a dedicated test case is required — happy to add one!
2025-06-29T07:51:32Z
https://github.com/huggingface/datasets/issues/7630#issuecomment-3016387412
https://api.github.com/repos/huggingface/datasets/issues/7630
IC_kwDODunzps6zymtU
https://api.github.com/repos/huggingface/datasets/issues/comments/3016387412/reactions
2025-06-29T07:51:32Z
https://api.github.com/repos/huggingface/datasets/issues/comments/3016387412
https://avatars.githubusercontent.com/u/142811259?v=4
https://api.github.com/users/ArjunJagdale/events{/privacy}
https://api.github.com/users/ArjunJagdale/followers
https://api.github.com/users/ArjunJagdale/following{/other_user}
https://api.github.com/users/ArjunJagdale/gists{/gist_id}
https://github.com/ArjunJagdale
ArjunJagdale
U_kgDOCIMgew
https://api.github.com/users/ArjunJagdale/orgs
https://api.github.com/users/ArjunJagdale/received_events
https://api.github.com/users/ArjunJagdale/repos
https://api.github.com/users/ArjunJagdale/starred{/owner}{/repo}
https://api.github.com/users/ArjunJagdale/subscriptions
User
https://api.github.com/users/ArjunJagdale
public
### Describe the bug

resume from ckpt skips samples if .map is applied

Maybe related: https://github.com/huggingface/datasets/issues/7538

### Steps to reproduce the bug

```python
from datasets import Dataset
from datasets.distributed import split_dataset_by_node


# Create dataset with map transformation
def create_dataset():
    ds = Dataset.from_dict({"id": list(range(100))})
    ds = ds.to_iterable_dataset(num_shards=4)
    ds = ds.map(lambda x: x) #comment it out to get desired behavior
    ds = split_dataset_by_node(ds, rank=0, world_size=2)
    return ds


ds = create_dataset()

# Iterate and save checkpoint after 10 samples
it = iter(ds)
for idx, sample in enumerate(it):
    if idx == 9:  # Checkpoint after 10 samples
        checkpoint = ds.state_dict()
        print(f"Checkpoint saved at sample: {sample['id']}")
        break

# Continue with original iterator
original_next_samples = []
for idx, sample in enumerate(it):
    original_next_samples.append(sample["id"])
    if idx >= 4:
        break

# Resume from checkpoint
ds_new = create_dataset()
ds_new.load_state_dict(checkpoint)

# Get samples from resumed iterator
it_new = iter(ds_new)
resumed_next_samples = []
for idx, sample in enumerate(it_new):
    resumed_next_samples.append(sample["id"])
    if idx >= 4:
        break

print(f"\nExpected next samples: {original_next_samples}")
print(f"Actual next samples:   {resumed_next_samples}")
print(
    f"\n❌ BUG: {resumed_next_samples[0] - original_next_samples[0]} samples were skipped!"
)

```

With map
```
Checkpoint saved at sample: 9

Expected next samples: [10, 11, 12, 13, 14]
Actual next samples:   [50, 51, 52, 53, 54]

❌ BUG: 40 samples were skipped!
```

### Expected behavior

without map
```
Expected next samples: [10, 11, 12, 13, 14]
Actual next samples:   [10, 11, 12, 13, 14]

❌ BUG: 0 samples were skipped!
```

### Environment info

datasets == 3.6.0