Add batching to `IterableDataset`
MEMBER
Cool ! Thanks for diving into it :)

Your implementation is great and indeed supports shuffling and batching, you just need to additionally account for state_dict (for dataset [checkpointing+resuming](https://huggingface.co/docs/datasets/main/en/use_with_pytorch#checkpoint-and-resume))

That being said, I believe the implementation can be made simpler by relying on `IterableDataset.map()` which already implements all this. Maybe something like

```python

def batch(self, batch_size: int, drop_last_batch: bool = False) -> "IterableDataset":
    def batch(unbatched: dict[str, list]) -> dict[str, list]:
        return {k: [v] for k, v in unbatched}

    return self.map(batch, batched=True, batch_size=batch_size, drop_last_batch=drop_last_batch)
```

And this way no need to reimplement everything !

(my only small concern is that it's not an Arrow-optimized function so it requires the examples to be manipulated as python objects even if the original data is in Arrow format (e.g. when streaming Parquet files) but it's not a big deal and we can see later if we need to optimize this)
2024-07-19T14:34:02Z
https://github.com/huggingface/datasets/pull/7054#issuecomment-2239323638
https://api.github.com/repos/huggingface/datasets/issues/7054
IC_kwDODunzps6FeV32
https://api.github.com/repos/huggingface/datasets/issues/comments/2239323638/reactions
2024-07-19T15:10:51Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2239323638
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
CONTRIBUTOR
Thanks a lot for the feedback @lhoestq! I definitely could have saved some time looking into it properly first. 😅 

Implemented the `.batch()` method, added a proper docsrtring for documentation, and added tests.

Let me know what you think and if this needs some update.
2024-07-22T11:33:55Z
https://github.com/huggingface/datasets/pull/7054#issuecomment-2242738804
https://api.github.com/repos/huggingface/datasets/issues/7054
IC_kwDODunzps6FrXp0
https://api.github.com/repos/huggingface/datasets/issues/comments/2242738804/reactions
2024-07-22T11:33:55Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2242738804
https://avatars.githubusercontent.com/u/61876623?v=4
https://api.github.com/users/lappemic/events{/privacy}
https://api.github.com/users/lappemic/followers
https://api.github.com/users/lappemic/following{/other_user}
https://api.github.com/users/lappemic/gists{/gist_id}
https://github.com/lappemic
lappemic
MDQ6VXNlcjYxODc2NjIz
https://api.github.com/users/lappemic/orgs
https://api.github.com/users/lappemic/received_events
https://api.github.com/users/lappemic/repos
https://api.github.com/users/lappemic/starred{/owner}{/repo}
https://api.github.com/users/lappemic/subscriptions
User
https://api.github.com/users/lappemic
public
NONE
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/datasets/pr_7054). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.
2024-07-22T11:39:47Z
https://github.com/huggingface/datasets/pull/7054#issuecomment-2242748668
https://api.github.com/repos/huggingface/datasets/issues/7054
IC_kwDODunzps6FraD8
https://api.github.com/repos/huggingface/datasets/issues/comments/2242748668/reactions
2024-07-22T11:39:47Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2242748668
https://avatars.githubusercontent.com/u/99929124?v=4
https://api.github.com/users/HuggingFaceDocBuilderDev/events{/privacy}
https://api.github.com/users/HuggingFaceDocBuilderDev/followers
https://api.github.com/users/HuggingFaceDocBuilderDev/following{/other_user}
https://api.github.com/users/HuggingFaceDocBuilderDev/gists{/gist_id}
https://github.com/HuggingFaceDocBuilderDev
HuggingFaceDocBuilderDev
U_kgDOBfTMJA
https://api.github.com/users/HuggingFaceDocBuilderDev/orgs
https://api.github.com/users/HuggingFaceDocBuilderDev/received_events
https://api.github.com/users/HuggingFaceDocBuilderDev/repos
https://api.github.com/users/HuggingFaceDocBuilderDev/starred{/owner}{/repo}
https://api.github.com/users/HuggingFaceDocBuilderDev/subscriptions
User
https://api.github.com/users/HuggingFaceDocBuilderDev
public
CONTRIBUTOR
Thanks for the feedbak @lhoestq!

Applied it and referenced the `batched=True` option in the `map` function and highlighted the difference. Hope i got this right.
2024-07-23T07:21:07Z
https://github.com/huggingface/datasets/pull/7054#issuecomment-2244447649
https://api.github.com/repos/huggingface/datasets/issues/7054
IC_kwDODunzps6Fx42h
https://api.github.com/repos/huggingface/datasets/issues/comments/2244447649/reactions
2024-07-23T07:21:07Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2244447649
https://avatars.githubusercontent.com/u/61876623?v=4
https://api.github.com/users/lappemic/events{/privacy}
https://api.github.com/users/lappemic/followers
https://api.github.com/users/lappemic/following{/other_user}
https://api.github.com/users/lappemic/gists{/gist_id}
https://github.com/lappemic
lappemic
MDQ6VXNlcjYxODc2NjIz
https://api.github.com/users/lappemic/orgs
https://api.github.com/users/lappemic/received_events
https://api.github.com/users/lappemic/repos
https://api.github.com/users/lappemic/starred{/owner}{/repo}
https://api.github.com/users/lappemic/subscriptions
User
https://api.github.com/users/lappemic
public
NONE
<details>
<summary>Show benchmarks</summary>

PyArrow==8.0.0

<details>
<summary>Show updated benchmarks!</summary>

### Benchmark: benchmark_array_xd.json

| metric | read_batch_formatted_as_numpy after write_array2d | read_batch_formatted_as_numpy after write_flattened_sequence | read_batch_formatted_as_numpy after write_nested_sequence | read_batch_unformated after write_array2d | read_batch_unformated after write_flattened_sequence | read_batch_unformated after write_nested_sequence | read_col_formatted_as_numpy after write_array2d | read_col_formatted_as_numpy after write_flattened_sequence | read_col_formatted_as_numpy after write_nested_sequence | read_col_unformated after write_array2d | read_col_unformated after write_flattened_sequence | read_col_unformated after write_nested_sequence | read_formatted_as_numpy after write_array2d | read_formatted_as_numpy after write_flattened_sequence | read_formatted_as_numpy after write_nested_sequence | read_unformated after write_array2d | read_unformated after write_flattened_sequence | read_unformated after write_nested_sequence | write_array2d | write_flattened_sequence | write_nested_sequence |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.005181 / 0.011353 (-0.006172) | 0.003714 / 0.011008 (-0.007294) | 0.063060 / 0.038508 (0.024552) | 0.030885 / 0.023109 (0.007776) | 0.239060 / 0.275898 (-0.036838) | 0.262480 / 0.323480 (-0.061000) | 0.004103 / 0.007986 (-0.003883) | 0.002696 / 0.004328 (-0.001632) | 0.048706 / 0.004250 (0.044456) | 0.042577 / 0.037052 (0.005525) | 0.249928 / 0.258489 (-0.008561) | 0.283252 / 0.293841 (-0.010589) | 0.029304 / 0.128546 (-0.099242) | 0.012001 / 0.075646 (-0.063646) | 0.204467 / 0.419271 (-0.214804) | 0.035639 / 0.043533 (-0.007894) | 0.243850 / 0.255139 (-0.011289) | 0.261609 / 0.283200 (-0.021590) | 0.018302 / 0.141683 (-0.123381) | 1.096040 / 1.452155 (-0.356115) | 1.135917 / 1.492716 (-0.356800) |

### Benchmark: benchmark_getitem\_100B.json

| metric | get_batch_of\_1024\_random_rows | get_batch_of\_1024\_rows | get_first_row | get_last_row |
|--------|---|---|---|---|
| new / old (diff) | 0.091976 / 0.018006 (0.073970) | 0.296396 / 0.000490 (0.295906) | 0.000203 / 0.000200 (0.000003) | 0.000043 / 0.000054 (-0.000011) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.018405 / 0.037411 (-0.019007) | 0.062470 / 0.014526 (0.047944) | 0.073340 / 0.176557 (-0.103216) | 0.119474 / 0.737135 (-0.617661) | 0.075750 / 0.296338 (-0.220588) |

### Benchmark: benchmark_iterating.json

| metric | read 5000 | read 50000 | read_batch 50000 10 | read_batch 50000 100 | read_batch 50000 1000 | read_formatted numpy 5000 | read_formatted pandas 5000 | read_formatted tensorflow 5000 | read_formatted torch 5000 | read_formatted_batch numpy 5000 10 | read_formatted_batch numpy 5000 1000 | shuffled read 5000 | shuffled read 50000 | shuffled read_batch 50000 10 | shuffled read_batch 50000 100 | shuffled read_batch 50000 1000 | shuffled read_formatted numpy 5000 | shuffled read_formatted_batch numpy 5000 10 | shuffled read_formatted_batch numpy 5000 1000 |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.279586 / 0.215209 (0.064377) | 2.768542 / 2.077655 (0.690887) | 1.449158 / 1.504120 (-0.054962) | 1.328760 / 1.541195 (-0.212435) | 1.336338 / 1.468490 (-0.132152) | 0.732582 / 4.584777 (-3.852195) | 2.325558 / 3.745712 (-1.420154) | 2.898077 / 5.269862 (-2.371784) | 1.893107 / 4.565676 (-2.672569) | 0.078788 / 0.424275 (-0.345487) | 0.005273 / 0.007607 (-0.002335) | 0.334887 / 0.226044 (0.108842) | 3.304173 / 2.268929 (1.035244) | 1.834743 / 55.444624 (-53.609882) | 1.527463 / 6.876477 (-5.349014) | 1.538824 / 2.142072 (-0.603249) | 0.785646 / 4.805227 (-4.019581) | 0.134876 / 6.500664 (-6.365788) | 0.042894 / 0.075469 (-0.032575) |

### Benchmark: benchmark_map_filter.json

| metric | filter | map fast-tokenizer batched | map identity | map identity batched | map no-op batched | map no-op batched numpy | map no-op batched pandas | map no-op batched pytorch | map no-op batched tensorflow |
|--------|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.976635 / 1.841788 (-0.865152) | 11.217156 / 8.074308 (3.142848) | 9.616971 / 10.191392 (-0.574421) | 0.127276 / 0.680424 (-0.553148) | 0.014344 / 0.534201 (-0.519857) | 0.301896 / 0.579283 (-0.277387) | 0.259615 / 0.434364 (-0.174749) | 0.340693 / 0.540337 (-0.199645) | 0.429145 / 1.386936 (-0.957791) |

</details>
PyArrow==latest

<details>
<summary>Show updated benchmarks!</summary>

### Benchmark: benchmark_array_xd.json

| metric | read_batch_formatted_as_numpy after write_array2d | read_batch_formatted_as_numpy after write_flattened_sequence | read_batch_formatted_as_numpy after write_nested_sequence | read_batch_unformated after write_array2d | read_batch_unformated after write_flattened_sequence | read_batch_unformated after write_nested_sequence | read_col_formatted_as_numpy after write_array2d | read_col_formatted_as_numpy after write_flattened_sequence | read_col_formatted_as_numpy after write_nested_sequence | read_col_unformated after write_array2d | read_col_unformated after write_flattened_sequence | read_col_unformated after write_nested_sequence | read_formatted_as_numpy after write_array2d | read_formatted_as_numpy after write_flattened_sequence | read_formatted_as_numpy after write_nested_sequence | read_unformated after write_array2d | read_unformated after write_flattened_sequence | read_unformated after write_nested_sequence | write_array2d | write_flattened_sequence | write_nested_sequence |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.005534 / 0.011353 (-0.005819) | 0.003795 / 0.011008 (-0.007213) | 0.049761 / 0.038508 (0.011253) | 0.031311 / 0.023109 (0.008202) | 0.276032 / 0.275898 (0.000134) | 0.297316 / 0.323480 (-0.026164) | 0.004396 / 0.007986 (-0.003590) | 0.002693 / 0.004328 (-0.001635) | 0.049025 / 0.004250 (0.044775) | 0.039707 / 0.037052 (0.002654) | 0.284264 / 0.258489 (0.025775) | 0.319962 / 0.293841 (0.026121) | 0.031842 / 0.128546 (-0.096705) | 0.012192 / 0.075646 (-0.063454) | 0.059895 / 0.419271 (-0.359376) | 0.033676 / 0.043533 (-0.009856) | 0.275917 / 0.255139 (0.020778) | 0.292637 / 0.283200 (0.009437) | 0.017992 / 0.141683 (-0.123691) | 1.199329 / 1.452155 (-0.252826) | 1.259083 / 1.492716 (-0.233633) |

### Benchmark: benchmark_getitem\_100B.json

| metric | get_batch_of\_1024\_random_rows | get_batch_of\_1024\_rows | get_first_row | get_last_row |
|--------|---|---|---|---|
| new / old (diff) | 0.092770 / 0.018006 (0.074764) | 0.313363 / 0.000490 (0.312873) | 0.000212 / 0.000200 (0.000013) | 0.000052 / 0.000054 (-0.000003) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.022977 / 0.037411 (-0.014434) | 0.076839 / 0.014526 (0.062314) | 0.088289 / 0.176557 (-0.088267) | 0.128625 / 0.737135 (-0.608510) | 0.089348 / 0.296338 (-0.206990) |

### Benchmark: benchmark_iterating.json

| metric | read 5000 | read 50000 | read_batch 50000 10 | read_batch 50000 100 | read_batch 50000 1000 | read_formatted numpy 5000 | read_formatted pandas 5000 | read_formatted tensorflow 5000 | read_formatted torch 5000 | read_formatted_batch numpy 5000 10 | read_formatted_batch numpy 5000 1000 | shuffled read 5000 | shuffled read 50000 | shuffled read_batch 50000 10 | shuffled read_batch 50000 100 | shuffled read_batch 50000 1000 | shuffled read_formatted numpy 5000 | shuffled read_formatted_batch numpy 5000 10 | shuffled read_formatted_batch numpy 5000 1000 |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.300881 / 0.215209 (0.085672) | 2.946499 / 2.077655 (0.868845) | 1.599686 / 1.504120 (0.095566) | 1.479332 / 1.541195 (-0.061862) | 1.476910 / 1.468490 (0.008420) | 0.720536 / 4.584777 (-3.864241) | 0.944822 / 3.745712 (-2.800890) | 2.771864 / 5.269862 (-2.497998) | 1.886573 / 4.565676 (-2.679103) | 0.078462 / 0.424275 (-0.345813) | 0.005392 / 0.007607 (-0.002215) | 0.354984 / 0.226044 (0.128939) | 3.516449 / 2.268929 (1.247520) | 1.977033 / 55.444624 (-53.467592) | 1.671922 / 6.876477 (-5.204555) | 1.785755 / 2.142072 (-0.356318) | 0.795330 / 4.805227 (-4.009897) | 0.132895 / 6.500664 (-6.367769) | 0.041178 / 0.075469 (-0.034291) |

### Benchmark: benchmark_map_filter.json

| metric | filter | map fast-tokenizer batched | map identity | map identity batched | map no-op batched | map no-op batched numpy | map no-op batched pandas | map no-op batched pytorch | map no-op batched tensorflow |
|--------|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 1.031780 / 1.841788 (-0.810008) | 11.855600 / 8.074308 (3.781292) | 10.245599 / 10.191392 (0.054207) | 0.140649 / 0.680424 (-0.539775) | 0.015332 / 0.534201 (-0.518869) | 0.299402 / 0.579283 (-0.279881) | 0.120007 / 0.434364 (-0.314357) | 0.337770 / 0.540337 (-0.202568) | 0.433679 / 1.386936 (-0.953257) |

</details>
</details>

![](https://cml.dev/watermark.png#e83d6fa574710fcb44e341087239d2687183f62b "CML watermark")

2024-07-23T10:41:23Z
https://github.com/huggingface/datasets/pull/7054#issuecomment-2244870573
https://api.github.com/repos/huggingface/datasets/issues/7054
IC_kwDODunzps6FzgGt
Iv1.05c79e9ad1f6bdfa
2018-07-30T09:30:17Z
Automate your workflow from idea to production
branch_protection_rule check_run check_suite create delete deployment deployment_status discussion discussion_comment fork gollum issues issue_comment label merge_group milestone page_build project project_card project_column public pull_request pull_request_review pull_request_review_comment push registry_package release repository repository_dispatch status watch workflow_dispatch workflow_run
https://help.github.com/en/actions
https://github.com/apps/github-actions
GitHub Actions
MDM6QXBwMTUzNjg=
https://avatars.githubusercontent.com/u/9919?v=4
https://api.github.com/users/github/events{/privacy}
https://api.github.com/users/github/followers
https://api.github.com/users/github/following{/other_user}
https://api.github.com/users/github/gists{/gist_id}
https://github.com/github
github
MDEyOk9yZ2FuaXphdGlvbjk5MTk=
https://api.github.com/users/github/orgs
https://api.github.com/users/github/received_events
https://api.github.com/users/github/repos
https://api.github.com/users/github/starred{/owner}{/repo}
https://api.github.com/users/github/subscriptions
Organization
https://api.github.com/users/github
public
write
read
write
write
write
write
write
write
write
read
read
write
write
write
write
write
write
write
read
github-actions
2025-03-07T16:35:00Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2244870573/reactions
2024-07-23T10:41:23Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2244870573
https://avatars.githubusercontent.com/in/15368?v=4
https://api.github.com/users/github-actions%5Bbot%5D/events{/privacy}
https://api.github.com/users/github-actions%5Bbot%5D/followers
https://api.github.com/users/github-actions%5Bbot%5D/following{/other_user}
https://api.github.com/users/github-actions%5Bbot%5D/gists{/gist_id}
https://github.com/apps/github-actions
github-actions[bot]
MDM6Qm90NDE4OTgyODI=
https://api.github.com/users/github-actions%5Bbot%5D/orgs
https://api.github.com/users/github-actions%5Bbot%5D/received_events
https://api.github.com/users/github-actions%5Bbot%5D/repos
https://api.github.com/users/github-actions%5Bbot%5D/starred{/owner}{/repo}
https://api.github.com/users/github-actions%5Bbot%5D/subscriptions
Bot
https://api.github.com/users/github-actions%5Bbot%5D
public
I've taken a try at implementing a batched `IterableDataset` as requested in issue #6279. This PR adds a new `BatchedExamplesIterable` class and a `.batch()` method to the `IterableDataset` class.

The main changes are:
1. A new `BatchedExamplesIterable` that groups examples into batches.
2. A `.batch()` method for `IterableDataset` to easily create batched versions.
3. Support for shuffling and sharding to work with PyTorch DataLoader and multiple workers.

I'm not sure if this is exactly what you had in mind and also have not fully tested it atm, so I'd really appreciate your feedback. Does this seem like it's heading in the right direction? I'm happy to make any changes or explore different approaches if needed.

Pinging @lhoestq 