Title: `ArrowInvalid` occurs while running `Dataset.map()` function
State: closed

It seems to fail to process the final batch. This [colab](https://colab.research.google.com/drive/1_byLZRHwGP13PHMkJWo62Wp50S_Z2HMD?usp=sharing) can reproduce the error.

Code:

```python
# train_ds = Dataset(features: {
#     'title': Value(dtype='string', id=None), 
#     'score': Value(dtype='float64', id=None)
# }, num_rows: 99999)

# suggested in #665 
class PicklableTokenizer(BertJapaneseTokenizer):
    def __getstate__(self):
        state = dict(self.__dict__)
        state['do_lower_case'] = self.word_tokenizer.do_lower_case
        state['never_split'] = self.word_tokenizer.never_split
        del state['word_tokenizer']
        return state
    
    def __setstate(self):
        do_lower_case = state.pop('do_lower_case')
        never_split = state.pop('never_split')
        self.__dict__ = state
        self.word_tokenizer = MecabTokenizer(
            do_lower_case=do_lower_case, never_split=never_split
        )

t = PicklableTokenizer.from_pretrained('bert-base-japanese-whole-word-masking')

encoded = train_ds.map(
    lambda examples: {'tokens': t.encode(examples['title'], max_length=1000)}, batched=True, batch_size=1000
)
```

Error Message:

```
 99% 99/100 [00:22<00:00, 39.07ba/s]
---------------------------------------------------------------------------
ArrowInvalid                              Traceback (most recent call last)
<timed exec> in <module>

/usr/local/lib/python3.6/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)
   1242                 fn_kwargs=fn_kwargs,
   1243                 new_fingerprint=new_fingerprint,
-> 1244                 update_data=update_data,
   1245             )
   1246         else:

/usr/local/lib/python3.6/site-packages/datasets/arrow_dataset.py in wrapper(*args, **kwargs)
    151             "output_all_columns": self._output_all_columns,
    152         }
--> 153         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    154         if new_format["columns"] is not None:
    155             new_format["columns"] = list(set(new_format["columns"]) & set(out.column_names))

/usr/local/lib/python3.6/site-packages/datasets/fingerprint.py in wrapper(*args, **kwargs)
    161             # Call actual function
    162 
--> 163             out = func(self, *args, **kwargs)
    164 
    165             # Update fingerprint of in-place transforms + update in-place history of transforms

/usr/local/lib/python3.6/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, update_data)
   1496                     if update_data:
   1497                         batch = cast_to_python_objects(batch)
-> 1498                         writer.write_batch(batch)
   1499             if update_data:
   1500                 writer.finalize()  # close_stream=bool(buf_writer is None))  # We only close if we are writing in a file

/usr/local/lib/python3.6/site-packages/datasets/arrow_writer.py in write_batch(self, batch_examples, writer_batch_size)
    271             typed_sequence = TypedSequence(batch_examples[col], type=col_type, try_type=col_try_type)
    272             typed_sequence_examples[col] = typed_sequence
--> 273         pa_table = pa.Table.from_pydict(typed_sequence_examples)
    274         self.write_table(pa_table)
    275 

/usr/local/lib/python3.6/site-packages/pyarrow/table.pxi in pyarrow.lib.Table.from_pydict()

/usr/local/lib/python3.6/site-packages/pyarrow/table.pxi in pyarrow.lib.Table.from_arrays()

/usr/local/lib/python3.6/site-packages/pyarrow/table.pxi in pyarrow.lib.Table.validate()

/usr/local/lib/python3.6/site-packages/pyarrow/error.pxi in pyarrow.lib.check_status()

ArrowInvalid: Column 4 named tokens expected length 999 but got length 1000
```


Comments: