GitHub issue: Slow iteration speeds when using IterableDataset.shuffle with load_dataset(data_files=..., streaming=True)
Author: lajd (NONE)
State: open | Created: 2024-08-14 21:44:44

### Describe the bug

When I load a dataset from a number of arrow files, as in:

```
random_dataset = load_dataset(
    "arrow",
    data_files={split: shard_filepaths},
    streaming=True,
    split=split,
)
```

I'm able to get fast iteration speeds when iterating over the dataset without shuffling. 

When I shuffle the dataset, the iteration speed is reduced by ~1000x.

It's very possible the way I'm loading dataset shards is not appropriate; if so please advise!

Thanks for the help

### Steps to reproduce the bug

Here's full code to reproduce the issue:

- Generate a random dataset
- Create shards of data independently using Dataset.save_to_disk()
- The below will generate 16 shards (arrow files), of 512 examples each


```
import time
from pathlib import Path
from multiprocessing import Pool, cpu_count

import torch
from datasets import Dataset, load_dataset

split = "train"
split_save_dir = "/tmp/random_split"

def generate_random_example():
    return {
        'inputs': torch.randn(128).tolist(),
        'indices': torch.randint(0, 10000, (2, 20000)).tolist(),
        'values': torch.randn(20000).tolist(), 
    }

def generate_shard_dataset(examples_per_shard: int = 512):
    dataset_dict = {
        'inputs': [],
        'indices': [],
        'values': []
    }

    for _ in range(examples_per_shard):
        example = generate_random_example()
        dataset_dict['inputs'].append(example['inputs'])
        dataset_dict['indices'].append(example['indices'])
        dataset_dict['values'].append(example['values'])

    return Dataset.from_dict(dataset_dict)

def save_shard(shard_idx, save_dir, examples_per_shard):
    shard_dataset = generate_shard_dataset(examples_per_shard)
    shard_write_path = Path(save_dir) / f"shard_{shard_idx}"
    shard_dataset.save_to_disk(shard_write_path)
    return str(Path(shard_write_path) / "data-00000-of-00001.arrow")

def generate_split_shards(save_dir, num_shards: int = 16, examples_per_shard: int = 512):
    with Pool(cpu_count()) as pool:
        args = [(m, save_dir, examples_per_shard) for m in range(num_shards)]
        shard_filepaths = pool.starmap(save_shard, args)
    
    return shard_filepaths

shard_filepaths = generate_split_shards(split_save_dir)
```

Load the dataset as IterableDataset:


```
random_dataset = load_dataset(
    "arrow",
    data_files={split: shard_filepaths},
    streaming=True,
    split=split,
)
random_dataset = random_dataset.with_format("numpy")
```

Observe the iterations/second when iterating over the dataset directly, and applying shuffling before iterating:

Without shuffling, this gives ~1500 iterations/second

```
start_time = time.time()
for count, item in enumerate(random_dataset):
    if count > 0 and count % 100 == 0:
        elapsed_time = time.time() - start_time
        iterations_per_second = count / elapsed_time
        print(f"Processed {count} items at an average of {iterations_per_second:.2f} iterations/second")
```

```
Processed 100 items at an average of 705.74 iterations/second
Processed 200 items at an average of 1169.68 iterations/second
Processed 300 items at an average of 1497.97 iterations/second
Processed 400 items at an average of 1739.62 iterations/second
Processed 500 items at an average of 1931.11 iterations/second`
```

When shuffling, this gives ~3 iterations/second:


```

random_dataset = random_dataset.shuffle(buffer_size=100,seed=42)

start_time = time.time()
for count, item in enumerate(random_dataset):
    if count > 0 and count % 100 == 0:
        elapsed_time = time.time() - start_time
        iterations_per_second = count / elapsed_time
        print(f"Processed {count} items at an average of {iterations_per_second:.2f} iterations/second")
```


```
Processed 100 items at an average of 3.75 iterations/second
Processed 200 items at an average of 3.93 iterations/second
```

### Expected behavior

Iterations per second should be barely affected by shuffling, especially with a small buffer size

### Environment info

Datasets version: 2.21.0
Python 3.10
Ubuntu 22.04