Title: OverflowError: There was an overflow in the <class 'pyarrow.lib.ListArray'>. Try to reduce writer_batch_size to have batches smaller than 2GB
State: open
Labels: description: Something isn't working
name: bug

## Describe the bug
Despite having batches way under 2Gb when running `datasets.map()`, after processing correctly the data of the first batch without fuss and irrespective of writer_batch_size (say 2,4,8,16,32,64 and 128 in my case), it returns the following error :

> OverflowError: There was an overflow in the <class 'pyarrow.lib.ListArray'>. Try to reduce writer_batch_size to have batches smaller than 2GB

Note that I always run `batch_size=writer_batch_size` :

## Steps to reproduce the bug
```python
datasets.map(lambda example : {"column_name" : function(arguments)}, batched=False, remove_columns = datasets.column_names, batch_size=batch_size, writer_batch_size=batch_size, disable_nullable=True, num_proc=None, desc="blablabla")
```
## Introspecting CUDA memory during bug

Placed within `function(arguments)` the following statement to introspect memory usage, merely a little over 1/4 of 2Gb
`print(torch.cuda.memory_summary(device=device, abbreviated=False))`

> |===========================================================================|
 |                  PyTorch CUDA memory summary, device ID 0                 |
 |---------------------------------------------------------------------------|
 |            CUDA OOMs: 0            |        cudaMalloc retries: 0         |
 |===========================================================================|
 |        Metric         | Cur Usage  | Peak Usage | Tot Alloc  | Tot Freed  |
 |---------------------------------------------------------------------------|
 | Allocated memory      |  541418 KB |  545725 KB |  555695 KB |   14276 KB |
 |       from large pool |  540672 KB |  544431 KB |  544431 KB |    3759 KB |
 |       from small pool |     746 KB |    1714 KB |   11264 KB |   10517 KB |
 |---------------------------------------------------------------------------|
 | Active memory         |  541418 KB |  545725 KB |  555695 KB |   14276 KB |
 |       from large pool |  540672 KB |  544431 KB |  544431 KB |    3759 KB |
 |       from small pool |     746 KB |    1714 KB |   11264 KB |   10517 KB |
 |---------------------------------------------------------------------------|
 | GPU reserved memory   |  598016 KB |  598016 KB |  598016 KB |       0 B  |
 |       from large pool |  595968 KB |  595968 KB |  595968 KB |       0 B  |
 |       from small pool |    2048 KB |    2048 KB |    2048 KB |       0 B  |
 |---------------------------------------------------------------------------|
 | Non-releasable memory |   36117 KB |   52292 KB |  274275 KB |  238158 KB |
|       from large pool |   34816 KB |   51537 KB |  261713 KB |  226897 KB |
 |       from small pool |    1301 KB |    2045 KB |   12562 KB |   11261 KB |
 |---------------------------------------------------------------------------|
 | Allocations           |     198    |     224    |     478    |     280    |
 |       from large pool |      74    |      75    |      75    |       1    |
 |       from small pool |     124    |     150    |     403    |     279    |
 |---------------------------------------------------------------------------|
 | Active allocs         |     198    |     224    |     478    |     280    |
 |       from large pool |      74    |      75    |      75    |       1    |
 |       from small pool |     124    |     150    |     403    |     279    |
 |---------------------------------------------------------------------------|
 | GPU reserved segments |      21    |      21    |      21    |       0    |
 |       from large pool |      20    |      20    |      20    |       0    |
 |       from small pool |       1    |       1    |       1    |       0    |
 |---------------------------------------------------------------------------|
 | Non-releasable allocs |      18    |      23    |     166    |     148    |
 |       from large pool |      17    |      18    |      19    |       2    |
 |       from small pool |       1    |       6    |     147    |     146    |
 |===========================================================================|

## Expected results
Efficiently process the datasets and write it down to disk.

## Actual results
--------------------------------------------------------------------------
OverflowError                             Traceback (most recent call last)
~\anaconda3\envs\xxx\lib\site-packages\datasets\arrow_dataset.py in _map_single(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, new_fingerprint, rank, offset, disable_tqdm, desc, cache_only)
   2390                             else:
-> 2391                                 writer.write(example)
   2392                 else:

~\anaconda3\envs\xxx\lib\site-packages\datasets\arrow_writer.py in write(self, example, key, writer_batch_size)
    367 
--> 368             self.write_examples_on_file()
    369 

~\anaconda3\envs\xxx\lib\site-packages\datasets\arrow_writer.py in write_examples_on_file(self)
    316                 if not isinstance(pa_array[0], pa.lib.FloatScalar):
--> 317                     raise OverflowError(
    318                         "There was an overflow in the {}. Try to reduce writer_batch_size to have batches smaller than 2GB".format(

OverflowError: There was an overflow in the <class 'pyarrow.lib.ListArray'>. Try to reduce writer_batch_size to have batches smaller than 2GB

During handling of the above exception, another exception occurred:

OverflowError                             Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_16268/2456940807.py in <module>
      3 #tracker = OfflineEmissionsTracker(country_iso_code="FRA", project_name='xxx'+time_stamp,output_dir='./codecarbon')
      4 #tracker.start()
----> 5 process_datasets(source_datasets_paths, dataset_dir, LM_tokenizer, LMhead_model, datasets_selection=['wikipedia'], from_scratch=True,
      6                  clean_sentences=False, negative_sampling=False, translate=False, tokenize=False, generate_embeddings=True, concatenate_embeddings=False,
      7                  max_sample=10000, padding='do_not_pad', truncation=True, cpu_batch_size=1000, gpu_batch_size=2, cpu_writer_batch_size=1000, gpu_writer_batch_size=2, disable_nullable=True, num_proc=None) #

~\xxx\xxx.py in process_datasets(source_datasets_paths, dataset_dir, LM_tokenizer, LMhead_model, datasets_selection, from_scratch, clean_sentences, translate, negative_sampling, tokenize, generate_embeddings, concatenate_embeddings, max_sample, padding, truncation, cpu_batch_size, gpu_batch_size, cpu_writer_batch_size, gpu_writer_batch_size, disable_nullable, num_proc)
    481             for column in tqdm(dataset.column_names, desc=f'Processing column', leave=False):
    482                 if "xxx_" in column:
--> 483                     dataset = dataset.map(lambda example :
    484                                                 {"embeddings_"+str(column).replace("translated_",""):function(input_ids=example[column],
    485                                                                                                                             token_type_ids=example[column.replace("input_ids","token_type_ids")],

~\anaconda3\envs\xxx\lib\site-packages\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, desc)
   2034 
   2035         if num_proc is None or num_proc == 1:
-> 2036             return self._map_single(
   2037                 function=function,
   2038                 with_indices=with_indices,

~\anaconda3\envs\xxx\lib\site-packages\datasets\arrow_dataset.py in wrapper(*args, **kwargs)
    501             self: "Dataset" = kwargs.pop("self")
    502         # apply actual function
--> 503         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    504         datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    505         for dataset in datasets:

~\anaconda3\envs\xxx\lib\site-packages\datasets\arrow_dataset.py in wrapper(*args, **kwargs)
    468         }
    469         # apply actual function
--> 470         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    471         datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    472         # re-apply format to the output

~\anaconda3\envs\xxx\lib\site-packages\datasets\fingerprint.py in wrapper(*args, **kwargs)
    404             # Call actual function
    405 
--> 406             out = func(self, *args, **kwargs)
    407 
    408             # Update fingerprint of in-place transforms + update in-place history of transforms

~\anaconda3\envs\xxx\lib\site-packages\datasets\arrow_dataset.py in _map_single(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, new_fingerprint, rank, offset, disable_tqdm, desc, cache_only)
   2425                 if update_data:
   2426                     if writer is not None:
-> 2427                         writer.finalize()
   2428                     if tmp_file is not None:
   2429                         tmp_file.close()

~\anaconda3\envs\xxx\lib\site-packages\datasets\arrow_writer.py in finalize(self, close_stream)
    440             # Re-intializing to empty list for next batch
    441             self.hkey_record = []
--> 442         self.write_examples_on_file()
    443         if self.pa_writer is None:
    444             if self._schema is not None:

~\anaconda3\envs\xxx\lib\site-packages\datasets\arrow_writer.py in write_examples_on_file(self)
    315                 # This check fails with FloatArrays with nans, which is not what we want, so account for that:
    316                 if not isinstance(pa_array[0], pa.lib.FloatScalar):
--> 317                     raise OverflowError(
    318                         "There was an overflow in the {}. Try to reduce writer_batch_size to have batches smaller than 2GB".format(
    319                             type(pa_array)

OverflowError: There was an overflow in the <class 'pyarrow.lib.ListArray'>. Try to reduce writer_batch_size to have batches smaller than 2GB

## Environment info
<!-- You can run the command `datasets-cli env` and copy-and-paste its output below. -->
- `datasets` version: 1.13.3
- Platform: Windows-10-10.0.19042-SP0
- Python version: 3.8.11
- PyArrow version: 3.0.0

##Next steps
Testing on Linux.
@albertvillanova 


Comments:
I am very unsure on why you tagged me here. I am not a maintainer of the Datasets library and have no idea how to help you.

fixed

Ok got it, tensor full of NaNs, cf.

~\anaconda3\envs\xxx\lib\site-packages\datasets\arrow_writer.py in write_examples_on_file(self)
315 # This check fails with FloatArrays with nans, which is not what we want, so account for that:

Actually this is is a live bug, documented yet still live so reopening