huggingface / huggingface/datasets

Support concatenating multiple streaming datasets while preserving the sum of shards

Open
#8,293 1 comment 0 reactions 0 assignees View on GitHub
enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.