Mapping gets stuck at 99%
### 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.rand(30000, 3),
                "normals": torch.rand(30000, 3),
                "features": torch.rand(30000, 6),
                "scalars": torch.rand(20),
            }
```

```python
# main.py

import datasets
import torch

def compute_mean_std(
    dataset: datasets.Dataset,
) -> dict[str, torch.Tensor]:
    """Compute the mean and standard deviation of each feature of the dataset.

    Args:
        dataset (`Dataset`): A huggingface dataset.

    Returns:
        dict: A dictionary containing the mean and standard deviation of each feature.
    """
    result = {}

    for key in dataset:
        # extract data from dataset
        data: torch.Tensor = dataset[key]  # type: ignore

        # reshape data, from (a, ..., b, c) -> (*, c)
        data = data.reshape(-1, data.shape[-1])

        # compute mean and std
        mean = data.mean(dim=0)  # (c)
        std = data.std(dim=0)  # (c)

        # store in result
        result[key] = torch.stack((mean, std))

    return result


def apply_mean_std(
    dataset: datasets.Dataset,
    mean_std: datasets.Dataset,
) -> dict[str, torch.Tensor]:
    """Normalize the dataset using the mean and standard deviation of each feature.

    Args:
        dataset (`Dataset`): A huggingface dataset.
        mean_std (`Dataset`): A huggingface dataset containing the mean and standard deviation of each feature.

    Returns:
        dict: A dictionary containing the normalized dataset.
    """
    result = {}

    for key in mean_std.column_names:
        # extract data from dataset
        data: torch.Tensor = dataset[key]  # type: ignore

        # extract mean and std from dict
        mean = mean_std[key][0]  # type: ignore
        std = mean_std[key][1]  # type: ignore

        # normalize data
        normalized_data = (data - mean) / std

        result[key] = normalized_data

    return result

# hack to force the map function to use the entire dataset
MAX_MAP_BATCH_SIZE = 1_000_000_000

# get dataset
ds = datasets.load_dataset(
    path="random_data.py",
    split="train",
).with_format("torch")

# compute mean/std of each feature
mean_std = ds.map(
    desc="Computing mean/std",  # type: ignore
    remove_columns=ds.column_names,  # type: ignore
    function=compute_mean_std,
    batch_size=MAX_MAP_BATCH_SIZE,
    batched=True,
)

# normalize each feature of the dataset
ds_normalized = ds.map(
    desc="Applying mean/std",  # type: ignore
    function=apply_mean_std,
    batched=False,
    fn_kwargs={
        "mean_std": mean_std,
    },
)
```

### 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