Mapping gets stuck at 99%
Author: Laurent2916 | Pull request: no | State: open

### Describe the bug

Hi !

I'm currently working with a large (~150GB) unnormalized dataset at work.
The dataset is available on a read-only filesystem internally, and I use a [loading script](https://huggingface.co/docs/datasets/dataset_script) to retreive it.

I want to normalize the features of the dataset, meaning I need to compute the mean and standard deviation metric for each feature of the entire dataset. I cannot load the entire dataset to RAM as it is too big, so following [this discussion on the huggingface discourse](https://discuss.huggingface.co/t/copy-columns-in-a-dataset-and-compute-statistics-for-a-column/22157) I am using a [map operation](https://huggingface.co/docs/datasets/v2.14.0/en/package_reference/main_classes#datasets.Dataset.map) to first compute the metrics and a second map operation to apply them on the dataset.

The problem lies in the second mapping, as it gets stuck at ~99%. By checking what the process does (using `htop` and `strace`) it seems to be doing a lot of I/O operations, and I'm not sure why.

Obviously, I could always normalize the dataset externally and then load it using a loading script. However, since the internal dataset is updated fairly frequently, using the library to perform normalization automatically would make it much easier for me.


### Steps to reproduce the bug

I'm able to reproduce the problem using the following scripts:

```python
# random_data.py

import datasets
import torch

_VERSION = "1.0.0"


class RandomDataset(datasets.GeneratorBasedBuilder):
    def _info(self):
        return datasets.DatasetInfo(
            version=_VERSION,
            supervised_keys=None,
            features=datasets.Features(
                {
                    "positions": datasets.Array2D(
                        shape=(30000, 3),
                        dtype="float32",
                    ),
                    "normals": datasets.Array2D(
                        shape=(30000, 3),
                        dtype="float32",
                    ),
                    "features": datasets.Array2D(
                        shape=(30000, 6),
                        dtype="float32",
                    ),
                    "scalars": datasets.Sequence(
                        feature=datasets.Value("float32"),
                        length=20,
                    ),
                },
            ),
        )

    def _split_generators(self, dl_manager):
        return [
            datasets.SplitGenerator(
                name=datasets.Split.TRAIN,  # type: ignore
                gen_kwargs={"nb_samples": 1000},
            ),
            datasets.SplitGenerator(
                name=datasets.Split.TEST,  # type: ignore
                gen_kwargs={"nb_samples": 100},
            ),
        ]

    def _generate_examples(self, nb_samples: int):
        for idx in range(nb_samples):
            yield idx, {
                "positions": torch.randn(30000, 3),
                "normals": torch.randn(30000, 3),
                "features": torch.randn(30000, 6),
                "scalars": torch.randn(20),
            }
```

```python
# main.py

import datasets
import torch


def apply_mean_std(
    dataset: datasets.Dataset,
    means: dict[str, torch.Tensor],
    stds: dict[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
    """Normalize the dataset using the mean and standard deviation of each feature.

    Args:
        dataset (`Dataset`): A huggingface dataset.
        mean (`dict[str, Tensor]`): A dictionary containing the mean of each feature.
        std (`dict[str, Tensor]`): A dictionary containing the standard deviation of each feature.

    Returns:
        dict: A dictionary containing the normalized dataset.
    """
    result = {}

    for key in means.keys():
        # extract data from dataset
        data: torch.Tensor = dataset[key]  # type: ignore

        # extract mean and std from dict
        mean = means[key]  # type: ignore
        std = stds[key]  # type: ignore

        # normalize data
        normalized_data = (data - mean) / std

        result[key] = normalized_data

    return result


# get dataset
ds = datasets.load_dataset(
    path="random_data.py",
    split="train",
).with_format("torch")

# compute mean (along last axis)
means = {key: torch.zeros(ds[key][0].shape[-1]) for key in ds.column_names}
means_sq = {key: torch.zeros(ds[key][0].shape[-1]) for key in ds.column_names}

for batch in ds.iter(batch_size=8):
    for key in ds.column_names:
        data = batch[key]
        batch_size = data.shape[0]
        data = data.reshape(-1, data.shape[-1])
        means[key] += data.mean(dim=0) / len(ds) * batch_size
        means_sq[key] += (data**2).mean(dim=0) / len(ds) * batch_size

# compute std (along last axis)
stds = {key: torch.sqrt(means_sq[key] - means[key] ** 2) for key in ds.column_names}

# normalize each feature of the dataset
ds_normalized = ds.map(
    desc="Applying mean/std",  # type: ignore
    function=apply_mean_std,
    batched=False,
    fn_kwargs={
        "means": means,
        "stds": stds,
    },
)
```

### Expected behavior

Using the previous scripts, the `ds_normalized` mapping completes in ~5 minutes, but any subsequent use of `ds_normalized` is really really slow, for example reapplying `apply_mean_std` to `ds_normalized` takes forever. This is very strange, I'm sure I must be missing something, but I would still expect this to be faster.

### Environment info

- `datasets` version: 2.13.1
- Platform: Linux-3.10.0-1160.66.1.el7.x86_64-x86_64-with-glibc2.17
- Python version: 3.10.12
- Huggingface_hub version: 0.15.1
- PyArrow version: 12.0.0
- Pandas version: 2.0.2