Issue: offset overflow when multiprocessing batched map on large datasets.
State: closed
Labels: name: bug

It only happened when "multiprocessing" + "batched" + "large dataset" at the same time.

```
def bprocess(examples):
  examples['len'] = []
  for text in examples['text']:
    examples['len'].append(len(text))
  return examples
wiki.map(brpocess, batched=True, num_proc=8)
```
```
---------------------------------------------------------------------------
RemoteTraceback                           Traceback (most recent call last)
RemoteTraceback: 
"""
Traceback (most recent call last):
  File "/home/yisiang/miniconda3/envs/ml/lib/python3.7/multiprocessing/pool.py", line 121, in worker
    result = (True, func(*args, **kwds))
  File "/home/yisiang/datasets/src/datasets/arrow_dataset.py", line 153, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
  File "/home/yisiang/datasets/src/datasets/fingerprint.py", line 163, in wrapper
    out = func(self, *args, **kwargs)
  File "/home/yisiang/datasets/src/datasets/arrow_dataset.py", line 1486, in _map_single
    batch = self[i : i + batch_size]
  File "/home/yisiang/datasets/src/datasets/arrow_dataset.py", line 1071, in __getitem__
    format_kwargs=self._format_kwargs,
  File "/home/yisiang/datasets/src/datasets/arrow_dataset.py", line 972, in _getitem
    data_subset = self._data.take(indices_array)
  File "pyarrow/table.pxi", line 1145, in pyarrow.lib.Table.take
  File "/home/yisiang/miniconda3/envs/ml/lib/python3.7/site-packages/pyarrow/compute.py", line 268, in take
    return call_function('take', [data, indices], options)
  File "pyarrow/_compute.pyx", line 298, in pyarrow._compute.call_function
  File "pyarrow/_compute.pyx", line 192, in pyarrow._compute.Function.call
  File "pyarrow/error.pxi", line 122, in pyarrow.lib.pyarrow_internal_check_status
  File "pyarrow/error.pxi", line 84, in pyarrow.lib.check_status
pyarrow.lib.ArrowInvalid: offset overflow while concatenating arrays
"""

The above exception was the direct cause of the following exception:

ArrowInvalid                              Traceback (most recent call last)
 in 
     30   owt = datasets.load_dataset('/home/yisiang/datasets/datasets/openwebtext/openwebtext.py', cache_dir='./datasets')['train']
     31   print('load/create data from OpenWebText Corpus for ELECTRA')
---> 32   e_owt = ELECTRAProcessor(owt, apply_cleaning=False).map(cache_file_name=f"electra_owt_{c.max_length}.arrow")
     33   dsets.append(e_owt)
     34 

~/Reexamine_Attention/electra_pytorch/_utils/utils.py in map(self, **kwargs)
    126       writer_batch_size=10**4,
    127       num_proc=num_proc,
--> 128       **kwargs
    129     )
    130 

~/hugdatafast/hugdatafast/transform.py in my_map(self, *args, **kwargs)
     21     if not cache_file_name.endswith('.arrow'): cache_file_name += '.arrow'
     22     if '/' not in cache_file_name: cache_file_name = os.path.join(self.cache_directory(), cache_file_name)
---> 23   return self.map(*args, cache_file_name=cache_file_name, **kwargs)
     24 
     25 @patch

~/datasets/src/datasets/arrow_dataset.py in map(self, function, with_indices, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, suffix_template, new_fingerprint)
   1285                 logger.info("Spawning {} processes".format(num_proc))
   1286                 results = [pool.apply_async(self.__class__._map_single, kwds=kwds) for kwds in kwds_per_shard]
-> 1287                 transformed_shards = [r.get() for r in results]
   1288                 logger.info("Concatenating {} shards from multiprocessing".format(num_proc))
   1289                 result = concatenate_datasets(transformed_shards)

~/datasets/src/datasets/arrow_dataset.py in (.0)
   1285                 logger.info("Spawning {} processes".format(num_proc))
   1286                 results = [pool.apply_async(self.__class__._map_single, kwds=kwds) for kwds in kwds_per_shard]
-> 1287                 transformed_shards = [r.get() for r in results]
   1288                 logger.info("Concatenating {} shards from multiprocessing".format(num_proc))
   1289                 result = concatenate_datasets(transformed_shards)

~/miniconda3/envs/ml/lib/python3.7/multiprocessing/pool.py in get(self, timeout)
    655             return self._value
    656         else:
--> 657             raise self._value
    658 
    659     def _set(self, i, obj):

ArrowInvalid: offset overflow while concatenating arrays
```

This should be fixed with #645 

Feel free to re-open if it still occurs

This has just happened to me while working with a large (65GB) Parquet dataset.
```
[rank0]: Traceback (most recent call last):
[rank0]:   File "/app/LLaMA-Factory/src/llamafactory/launcher.py", line 23, in <module>
[rank0]:     launch()
[rank0]:   File "/app/LLaMA-Factory/src/llamafactory/launcher.py", line 19, in launch
[rank0]:     run_exp()
[rank0]:   File "/app/LLaMA-Factory/src/llamafactory/train/tuner.py", line 110, in run_exp
[rank0]:     _training_function(config={"args": args, "callbacks": callbacks})
[rank0]:   File "/app/LLaMA-Factory/src/llamafactory/train/tuner.py", line 72, in _training_function
[rank0]:     run_sft(model_args, data_args, training_args, finetuning_args, generating_args, callbacks)
[rank0]:   File "/app/LLaMA-Factory/src/llamafactory/train/sft/workflow.py", line 51, in run_sft
[rank0]:     dataset_module = get_dataset(template, model_args, data_args, training_args, stage="sft", **tokenizer_module)
[rank0]:                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/app/LLaMA-Factory/src/llamafactory/data/loader.py", line 306, in get_dataset
[rank0]:     dataset = _get_merged_dataset(data_args.dataset, model_args, data_args, training_args, stage)
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/app/LLaMA-Factory/src/llamafactory/data/loader.py", line 184, in _get_merged_dataset
[rank0]:     datasets[dataset_name] = _load_single_dataset(dataset_attr, model_args, data_args, training_args)
[rank0]:                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/app/LLaMA-Factory/src/llamafactory/data/loader.py", line 164, in _load_single_dataset
[rank0]:     return align_dataset(dataset, dataset_attr, data_args, training_args)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/app/LLaMA-Factory/src/llamafactory/data/converter.py", line 279, in align_dataset
[rank0]:     return dataset.map(
[rank0]:            ^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_dataset.py", line 557, in wrapper
[rank0]:     out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
[rank0]:                                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/datasets/arrow_dataset.py", line 3171, in map
[rank0]:     for rank, done, content in iflatmap_unordered(
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/datasets/utils/py_utils.py", line 728, in iflatmap_unordered
[rank0]:     [async_result.get(timeout=0.05) for async_result in async_results]
[rank0]:      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/usr/local/lib/python3.12/dist-packages/multiprocess/pool.py", line 774, in get
[rank0]:     raise self._value
[rank0]: pyarrow.lib.ArrowInvalid: offset overflow while concatenating arrays
```

Probably also worth mentioning that my dataset is multimodal; there is one text column and another with PIL image lists (singletons.)

Hi ! Arrow has a limitation in the max size of binary data (e.g. images) per batch, you can try to lower the Arrow writer batch size e.g. `datasets.config.DEFAULT_MAX_BATCH_SIZE = 100` (default is 1000)

I've already tried to set that via the keyword argument `writer_batch_size=4` and the same issue appears; I'm currently testing preprocessing with no multiprocessing.