Datasets: offset overflow when multiprocessing batched map on large datasets.

[richarddwang]

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

[lhoestq]

This should be fixed with #645

[lhoestq]

Feel free to re-open if it still occurs