huggingface / huggingface/datasets
Support concatenating multiple streaming datasets while preserving the sum of shards
- Dominant language
- Python
- Stars
- 22k
- Forks
- 3.4k
- Avg merge
- 5d 7h
- Merged PRs (30d)
- 17
Description
### Feature request
It would be useful if concatenating multiple `IterableDataset`s could preserve the shard information from each input dataset, so that:
- the total number of shards equals the sum of the input shards;
- `split_dataset_by_node()` and worker sharding continue to work efficiently;
- downstream distributed dataloaders can fully exploit the available parallelism.
### Motivation
Currently, it's straightforward to concatenate multiple streaming datasets using `interleave_datasets` or other approaches. However, after concatenation, the resulting dataset does not preserve the original sharding information from each input dataset.
For distributed training, especially with large-scale datasets, this leads to a reduced number of effective shards compared to the sum of the input datasets' shards. Consequently, data loading parallelism may be limited, and the available worker concurrency cannot fully utilize the original dataset partitioning.
For example:
- IterableDataset A: 1 shard
- IterableDataset B: 1 shard
After concatenation, we would expect the resulting streaming dataset to expose **2 shards**, allowing distributed workers to shard the combined dataset naturally.
### Your contribution
I'd be happy to work on this feature and submit a PR if the maintainers think this would be a useful addition.
Any feedback on the proposed design would be greatly appreciated.
Contributor guide
Research direction
Start by tracing how interleave_datasets combines IterableDatasets and how split_dataset_by_node and worker sharding consume shard information. Compare the current concatenated dataset behavior with two one-shard inputs, then define and validate a design where the combined dataset exposes their summed shards while preserving distributed loading behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- data, distributed-systems
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100