Datasets >= 5 ruins sharding when shuffle() is called

The shuffle change in v5 appears to be intentional, but whether this particular behavior is intended is much more questionable:

@lhoestq

I can reproduce the 8 shards -> shuffle() -> 1 shard -> one effective Data worker

behavior independently.

My current read is:

IterableDataset.shuffle()

so that the shuffle buffer can be fed from multiple input shards at once;num_shards

For an 8-file Parquet stream, the most practical routes I would try are:

dataset = dataset.reshard()
dataset = dataset.shuffle(seed=42, buffer_size=...)

or, if preserving the pre-v5 behavior is more important:

dataset = dataset.shuffle(
    seed=42,
    buffer_size=...,
    max_buffer_input_shards=1,
)

There is also a useful middle ground for your specific 8-shard / 4-worker case:

dataset = dataset.shuffle(
    seed=42,
    buffer_size=...,
    max_buffer_input_shards=2,
)

In my small reproduction, that retained 4 logical shards and all 4 Data workers, while still allowing the shuffle buffer to draw from more than one input shard at a time.

I would probably try reshard()

first for Parquet if it gives you enough logical shards, because that preserves the new v5 cross-shard mixing behavior. max_buffer_input_shards=1

is the clean compatibility option if you specifically want the old semantics.

A controlled reproduction I ran looked like this:

| Datasets version / configuration | num_shards after shuffle | Data workers that actually yielded examples | |---|---|---| 4.8.5, default shuffle() | 8 | 4 | 5.0.0, default shuffle() | 1 | 1 | 5.0.1, default shuffle() | 1 | 1 | current main (5.0.2.dev0 when tested), default | 1 | 1 | 5.x, max_buffer_input_shards=1 | 8 | 4 | 5.x, max_buffer_input_shards=2 | 4 | 4 | 5.x, max_buffer_input_shards=4 | 2 | 2 | 5.x, reshard() then default shuffle | 64 → 6 | 4 |

The test dataset was eight local Parquet files, 512 rows each, with eight row groups per file. Every tested configuration still returned all 4096 unique rows with zero duplicates, so in that small case the thing that changed was the execution/sharding topology, not dataset coverage.

I would not read much into throughput numbers from such a small local-file test; active worker count and correctness are the useful observations here.

Why this happens in v5So, for your original questions:

Why the change from Datasets 4 to 5?

To improve streaming shuffle quality by filling the shuffle buffer from multiple input shards instead of effectively processing one input shard at a time. That change is intentional and documented in the 5.0 release and #8194.

Is the 8 → 1 behavior itself a bug?

The 8 -> 1

logical-shard result follows from the current implementation and is reproducible. What is much less clear is whether losing Data worker parallelism as a consequence was an intended part of that API change. Given the current worker documentation, I would consider this worth an upstream clarification/issue rather than assuming it is expected.

Do you need to change how the streaming dataset is instantiated?

Probably not fundamentally. For Parquet, I would first try:

dataset = dataset.reshard()
dataset = dataset.shuffle(...)

If you need exact pre-v5 behavior:

dataset = dataset.shuffle(
    ...,
    max_buffer_input_shards=1,
)

And for eight source shards with four workers, max_buffer_input_shards=2

is also a reasonable low-cost experiment if you want to keep some of the new cross-shard mixing without collapsing below four logical shards.