Offset overflow when executing Dataset.map
Opened by jan-pair (open)

### Describe the bug

Hi, I'm trying to use `.map` method to cache multiple random crops from the image to speed up data processing during training, as the image size is too big. 

The map function executes all iterations, and then returns the following error:

```bash
Traceback (most recent call last):                                                                                                                                                                                                                                                                                          
  File "/home/ubuntu/miniconda3/envs/enhancement/lib/python3.8/site-packages/datasets/arrow_dataset.py", line 3353, in _map_single
    writer.finalize()  # close_stream=bool(buf_writer is None))  # We only close if we are writing in a file
  File "/home/ubuntu/miniconda3/envs/enhancement/lib/python3.8/site-packages/datasets/arrow_writer.py", line 582, in finalize
    self.write_examples_on_file()
  File "/home/ubuntu/miniconda3/envs/enhancement/lib/python3.8/site-packages/datasets/arrow_writer.py", line 446, in write_examples_on_file
    self.write_batch(batch_examples=batch_examples)
  File "/home/ubuntu/miniconda3/envs/enhancement/lib/python3.8/site-packages/datasets/arrow_writer.py", line 555, in write_batch
    self.write_table(pa_table, writer_batch_size)
  File "/home/ubuntu/miniconda3/envs/enhancement/lib/python3.8/site-packages/datasets/arrow_writer.py", line 567, in write_table
    pa_table = pa_table.combine_chunks()
  File "pyarrow/table.pxi", line 3315, in pyarrow.lib.Table.combine_chunks
  File "pyarrow/error.pxi", line 144, in pyarrow.lib.pyarrow_internal_check_status
  File "pyarrow/error.pxi", line 100, in pyarrow.lib.check_status
pyarrow.lib.ArrowInvalid: offset overflow while concatenating arrays
```

Here is the minimal code (`/home/datasets/DIV2K_train_HR` is just a folder of images that can be replaced by any appropriate):

### Steps to reproduce the bug

```python
from glob import glob

import torch
from datasets import Dataset, Image
from torchvision.transforms import PILToTensor, RandomCrop

file_paths = glob("/home/datasets/DIV2K_train_HR/*")

to_tensor = PILToTensor()
crop_transf = RandomCrop(size=256)

def prepare_data(example):
    tensor = to_tensor(example["image"].convert("RGB"))
    return {"hr": torch.stack([crop_transf(tensor) for _ in range(25)])}

train_data = Dataset.from_dict({"image": file_paths}).cast_column("image", Image())
train_data = train_data.map(
    prepare_data,
    cache_file_name="/home/datasets/DIV2K_train_HR_crops.tmp",
    desc="Caching multiple random crops of image",
    remove_columns="image",
)
print(train_data[0].keys(), train_data[0]["hr"].shape)
```

### Expected behavior

Cached file is stored at `"/home/datasets/DIV2K_train_HR_crops.tmp"`, output is `dict_keys(['hr']) torch.Size([25, 3, 256, 256])`

### Environment info

- `datasets` version: 2.10.1
- Platform: Linux-5.15.0-67-generic-x86_64-with-glibc2.10
- Python version: 3.8.16
- PyArrow version: 11.0.0
- Pandas version: 1.5.3
- Pytorch version: 2.0.0+cu117
- torchvision version: 0.15.1+cu117

Comments:
Upd. the above code works if we replace `25` with `1`, but the result value at key "hr" is not a tensor but a list of lists of lists of uint8.

Adding `train_data.set_format("torch")` after map fixes this, but the original issue remains



As a workaround, one can replace
`return {"hr": torch.stack([crop_transf(tensor) for _ in range(25)])}`
with
`return {f"hr_crop_{i}": crop_transf(tensor) for i in range(25)}`
and then choose appropriate crop randomly in further processing, but I still don't understand why the original approach doesn't work(