GitHub issue #4349: Dataset.map()'s fails at any value of parameter writer_batch_size 
State: closed | Created: 2022-05-13T16:55:12Z

## Describe the bug
If the the value of `writer_batch_size` is less than the total number of instances in the dataset it will fail at that same number of instances. If it is greater than the total number of instances, it fails on the last instance.

Context:
I am attempting to fine-tune a pre-trained HuggingFace transformers model called LayoutLMv2. This model takes three inputs: document images, words and word bounding boxes. [The Processor for this model has two options](https://huggingface.co/docs/transformers/model_doc/layoutlmv2#usage-layoutlmv2processor), the default is passing a document to the Processor and allowing it to create images of the document and use PyTesseract to perform OCR and generate words/bounding boxes. The other option is to provide `revision="no_ocr"` to the pre-trained model which allows you to use your own OCR results (in my case, Amazon Textract) so you have to provide the image, words and bounding boxes yourself. I am using this second option which might be good context for the bug.

I am using the Dataset.map() paradigm to create these three inputs, encode them and save the dataset. Note that my documents (data instances) on average are fairly large and can range from 1 page up to 300 pages.
Code I am using is provided below

## Steps to reproduce the bug
I do not have explicit sample code, but I will paste the code I'm using in case reading it helps. When `.map()` is called, the dataset has 2933 rows, many of which represent large pdf documents.
```python
def get_encoded_data(data):
    dataset = Dataset.from_pandas(data)
    unique_labels = data['label'].unique()
    features = Features({
        'image': Array3D(dtype="int64", shape=(3, 224, 224)),
        'input_ids': Sequence(feature=Value(dtype='int64')),
        'attention_mask': Sequence(Value(dtype='int64')),
        'token_type_ids': Sequence(Value(dtype='int64')),
        'bbox': Array2D(dtype="int64", shape=(512, 4)),
        'label': ClassLabel(num_classes=len(unique_labels), names=unique_labels),
    })

    encoded_dataset = dataset.map(preprocess_data, features=features, remove_columns=dataset.column_names, writer_batch_size=dataset.num_rows+1)
    encoded_dataset.save_to_disk(TRAINING_DATA_PATH + ENCODED_DATASET_NAME)
    encoded_dataset.set_format(type="torch")
    return encoded_dataset
```
```python
PROCESSOR = LayoutLMv2Processor.from_pretrained(MODEL_PATH, revision="no_ocr", use_fast=False)

def preprocess_data(examples):
    directory = os.path.join(FILES_PATH, examples['file_location'])
    images_dir = os.path.join(directory, PDF_IMAGE_DIR)
    textract_response_path = os.path.join(directory, 'textract.json')
    doc_meta_path = os.path.join(directory, 'doc_meta.json')
    textract_document = get_textract_document(textract_response_path, doc_meta_path)
    images, words, bboxes = get_doc_training_data(images_dir, textract_document)
    encoded_inputs = PROCESSOR(images, words, boxes=bboxes, padding="max_length", truncation=True)
    # https://github.com/NielsRogge/Transformers-Tutorials/issues/36
    encoded_inputs["image"] = np.array(encoded_inputs["image"])
    encoded_inputs["label"] = examples['label_id']
    return encoded_inputs
```

## Expected results
My expectation is that `writer_batch_size` allows one to simply trade off performance and memory requirements, not that it must be a specific number for `.map()` to function correctly.

## Actual results
If writer_batch_size is set to a value less than the number of rows, I get either:

```
OverflowError: There was an overflow with type <class 'list'>. Try to reduce writer_batch_size to have batches smaller than 2GB.
(offset overflow while concatenating arrays)
```
or simply

```
zsh: killed     python doc_classification.py

UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown
```

If it is greater than the number of rows, i get the `zsh: killed` error above

## Environment info
<!-- You can run the command `datasets-cli env` and copy-and-paste its output below. -->
- `datasets` version: 2.1.0
- Platform: macOS-12.2.1-arm64-arm-64bit
- Python version: 3.9.12
- PyArrow version: 6.0.1
- Pandas version: 1.4.2