datasets massively slows data reads, even in memory
Opened by lrast (open)

### Describe the bug

Loading image data in a huggingface dataset results in very slow read speeds, approximately 1000 times longer than reading the same data from a pytorch dataset. This applies even when the dataset is loaded into RAM using a `keep_in_memory=True` flag.

The following script reproduces the result with random data, but it applies equally to datasets that are loaded from the hub.

### Steps to reproduce the bug

The following script should reproduce the behavior
```
import torch
import time
from datasets import Dataset

images = torch.randint(0, 255, (1000, 3, 224, 224), dtype=torch.uint8)
labels = torch.randint(0, 200, (1000,), dtype=torch.uint8)

pt_dataset = torch.utils.data.TensorDataset(images, labels)

hf_dataset = Dataset.from_dict({'image': images, 'label':labels})
hf_dataset.set_format('torch', dtype=torch.uint8)
hf_in_memory = hf_dataset.map(lambda x: x, keep_in_memory=True)

# measure access speeds
def time_access(dataset, img_col):
    start_time = time.time()
    for i in range(1000):
        _ = dataset[i][img_col].shape
    end_time = time.time()
    return end_time - start_time


print(f"In-memory Tensor access: {time_access(pt_dataset, 0):.4f} seconds")
print(f"HF Dataset access: {time_access(hf_dataset, 'image'):.4f} seconds")
print(f"In-memory HF Dataset access: {time_access(hf_in_memory, 'image'):.4f} seconds")
```



### Expected behavior

For me, the above script produces
```
In-memory Tensor access: 0.0025 seconds
HF Dataset access: 2.9317 seconds
In-memory HF Dataset access: 2.8082 seconds
```

I think that this difference is larger than expected.


### Environment info

- `datasets` version: 4.0.0
- Platform: macOS-14.7.7-arm64-arm-64bit
- Python version: 3.12.11
- `huggingface_hub` version: 0.34.3
- PyArrow version: 18.0.0
- Pandas version: 2.2.3
- `fsspec` version: 2024.9.0

Comments:
Hi ! you should try

```python
from datasets import Array3D, Dataset, Features, Value

features = Features({"image": Array3D(shape=(3, 224, 224), dtype="uint8"), "label": Value("uint8")})
hf_dataset = Dataset.from_dict({'image': images, 'label':labels}, features=features)
```

otherwise the type of the "image" column is List(List(List(Value("uint8")))) and is less efficient.

Thanks! This leads to a 10x speedup:
```python
import torch
import time
from datasets import Array3D, Dataset, Features, Value

images = torch.randint(0, 255, (1000, 3, 224, 224), dtype=torch.uint8)
labels = torch.randint(0, 200, (1000,), dtype=torch.uint8)

pt_dataset = torch.utils.data.TensorDataset(images, labels)

features = Features({"image": Array3D(shape=(3, 224, 224), dtype="uint8"), "label": Value("uint8")})
hf_dataset = Dataset.from_dict({'image': images, 'label':labels}, features=features)
hf_in_memory = hf_dataset.map(lambda x: x, keep_in_memory=True)

hf_dataset.set_format('torch', dtype=torch.uint8)
hf_in_memory.set_format('torch', dtype=torch.uint8)

# measure access speeds
def time_access(dataset, img_col):
    start_time = time.time()
    for i in range(1000):
        _ = dataset[i][img_col].shape
    end_time = time.time()
    return end_time - start_time


print(f"In-memory Tensor access: {time_access(pt_dataset, 0):.4f} seconds")
print(f"HF Dataset access: {time_access(hf_dataset, 'image'):.4f} seconds")
print(f"In-memory HF Dataset access: {time_access(hf_in_memory, 'image'):.4f} seconds")
```
Produces
```
In-memory Tensor access: 0.0026 seconds
HF Dataset access: 0.2070 seconds
In-memory HF Dataset access: 0.2112 seconds
```

Curious if there is a reason why this is not the default behavior for huggingface image processors?
```python
from transformers import ViTImageProcessor
from transformers import AutoImageProcessor

from datasets import load_dataset
# Load the dataset
ds = load_dataset('ylecun/mnist', split='train[0:100]')

# Instantiate the processor, explicitly requesting NumPy arrays
processor1 = ViTImageProcessor.from_pretrained('facebook/vit-mae-base',  do_convert_rgb=True)
processor2 = AutoImageProcessor.from_pretrained("facebook/detr-resnet-50", use_fast=True)

processed1 = ds.map(lambda row: processor1(row['image']))
processed2 = ds.map(lambda row: processor2(row['image']))

print( type(processed1['pixel_values'][0]), type(processed1['pixel_values'][0]))
```
produces
```
<class 'list'> <class 'list'>
```

I can, of course, manually manipulate the dataset to the use the correct format, but this is fairly standard for images, and the performance implications seem large.