Custom formatting for lazy map + arrow data extraction refactor
State: closed | Author: lhoestq (MEMBER)

Hi !

This PR refactors the way data are extracted from pyarrow tables to extend it to the use of custom formatting functions.

While the internal storage of the dataset is always the Apache Arrow format, by setting a specific format on a dataset, you can cast the output of `datasets.Dataset.__getitem__` in NumPy/pandas/PyTorch/TensorFlow, on-the-fly.

A specific format can be activated with `datasets.Dataset.set_format`. For example: `dataset.set_format(type='torch', columns=['label'])`.

### What's new:

You can now also define your own formatting function that is applied on-the-fly. To do so you can pass your formatting function in the `transform` parameter of `datasets.Dataset.set_format`, and keep `type` to `None`.
A formatting function is a callable that takes a batch (as a dict, formatted as python) as input and returns a batch.

Here is an example to tokenize and pad tokens on-the-fly when accessing the samples:

```python
from datasets import load_dataset
from transformers import BertTokenizer

tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
def encode(batch):
    return tokenizer(batch["sentence1"], padding="longest", truncation=True, max_length=512, return_tensors="pt")

dataset = load_dataset("glue", "mrpc", split="train")
dataset.set_format(transform=encode)
dataset.format
# {'type': 'custom', 'format_kwargs': {'transform': <function __main__.encode(batch)>}, 'columns': ['idx', 'label', 'sentence1', 'sentence2'], 'output_all_columns': False}
dataset[:2]
# {'input_ids': tensor([[  101,  2572,  3217, ... 102]]), 'token_type_ids': tensor([[0, 0, 0, ... 0]]), 'attention_mask': tensor([[1, 1, 1, ... 1]])}
```

Let me know what you think of this API !
We can still change it if we want to.

Especially @sgugger since this may be useful when using `datasets` to train models.

EDIT: this was changed to `dataset.set_transform(encode)`

-------------------

Note:

I had to refactor the way data are extracted and formatted from pyarrow tables and I made it more robust and flexible. In particular I modularized it to be able to unit-test it properly. This was very helpful since I detected some bugs in the previous implementation and was able to fix them.

Some bugs I found and fixed:
- certain slices/ranges were not supported because negative ids were passed to pyarrow
- formatting as numpy/torch/tensorflow a column would make it lose its precision information (for example a column as `Value("float32")`) would be returned as a tensor of float64 (default behavior for numpy)
- on windows  integers formatted as numpy/torch/tensorflow were not always int64 tensors by default but were int32 

The unit tests for those are now really extensive :)

Comments: This PR is amazing!!!

I only looked at `arrow_dataset.py` and `formatting/formatting.py` but those look good to me.

My only (tiny) concern is the name of the function: I don't think it's self-evident that `set_format` applies a generic transformation, and some people might not look too far into the doc.

Maybe we could have an `apply_transform` or `process_columns` method which is called by `set_format` (to keep backward compatibility)?

What about something like `.set_format` and `.set_transform` ?
- set_format would be the same as right now, i.e. defined by a format type.
- set_transform would define the transformation that is applied on output batches on-the-fly.

I was also thinking about `._with_format` and `.with_transform`. It could be their equivalent but would create a **new** dataset with the corresponding format or transform ? I know @sgugger was interested in something like that.

Yup, I think that would make all of these options very clear!

I like all those options as well (as long as the `_` in `_with_format` is a typo ;-) )

Yes it's a typo indeed ;)

Alright I'll do the changes !

I took all your suggestions into account, thanks :)
Let me know if you have more comments

Hi @lhoestq , thanks for offering the set_transform() function. It is very handy to process large datasets on the fly. But I ran into a problem when using it (error message shown below). Since we are working with a large collection, there's no way to filter all invalid data points beforehand. Those invalid data points will be problematic with the set_transform and I don't find a good work-around to ignore them. I wonder if you can offer some advice on dealing with invalid data points in this case. Thank you!

```
Traceback (most recent call last):
  File "/export/share/ruimeng/env/anaconda/envs/ir/lib/python3.8/site-packages/torch/utils/data/_utils/worker.py", line 198, in _worker_loop
    data = fetcher.fetch(index)
  File "/export/share/ruimeng/env/anaconda/envs/ir/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 44, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/export/share/ruimeng/env/anaconda/envs/ir/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 44, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/export/share/ruimeng/env/anaconda/envs/ir/lib/python3.8/site-packages/datasets/arrow_dataset.py", line 1763, in __getitem__
    return self._getitem(
  File "/export/share/ruimeng/env/anaconda/envs/ir/lib/python3.8/site-packages/datasets/arrow_dataset.py", line 1748, in _getitem
    formatted_output = format_table(
  File "/export/share/ruimeng/env/anaconda/envs/ir/lib/python3.8/site-packages/datasets/formatting/formatting.py", line 532, in format_table
    return formatter(pa_table, query_type=query_type)
  File "/export/share/ruimeng/env/anaconda/envs/ir/lib/python3.8/site-packages/datasets/formatting/formatting.py", line 281, in __call__
    return self.format_row(pa_table)
  File "/export/share/ruimeng/env/anaconda/envs/ir/lib/python3.8/site-packages/datasets/formatting/formatting.py", line 391, in format_row
    raise TypeError(
TypeError: Custom formatting function must return a dict to be able to pick a row, but got None

```


> Hi @lhoestq , thanks for offering the set_transform() function. It is very handy to process large datasets on the fly. But I ran into a problem when using it (error message shown below). Since we are working with a large collection, there's no way to filter all invalid data points beforehand. Those invalid data points will be problematic with the set_transform and I don't find a good work-around to ignore them. I wonder if you can offer some advice on dealing with invalid data points in this case. Thank you!
> 
> ```
> Traceback (most recent call last):
>   File "/export/share/ruimeng/env/anaconda/envs/ir/lib/python3.8/site-packages/torch/utils/data/_utils/worker.py", line 198, in _worker_loop
>     data = fetcher.fetch(index)
>   File "/export/share/ruimeng/env/anaconda/envs/ir/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 44, in fetch
>     data = [self.dataset[idx] for idx in possibly_batched_index]
>   File "/export/share/ruimeng/env/anaconda/envs/ir/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 44, in <listcomp>
>     data = [self.dataset[idx] for idx in possibly_batched_index]
>   File "/export/share/ruimeng/env/anaconda/envs/ir/lib/python3.8/site-packages/datasets/arrow_dataset.py", line 1763, in __getitem__
>     return self._getitem(
>   File "/export/share/ruimeng/env/anaconda/envs/ir/lib/python3.8/site-packages/datasets/arrow_dataset.py", line 1748, in _getitem
>     formatted_output = format_table(
>   File "/export/share/ruimeng/env/anaconda/envs/ir/lib/python3.8/site-packages/datasets/formatting/formatting.py", line 532, in format_table
>     return formatter(pa_table, query_type=query_type)
>   File "/export/share/ruimeng/env/anaconda/envs/ir/lib/python3.8/site-packages/datasets/formatting/formatting.py", line 281, in __call__
>     return self.format_row(pa_table)
>   File "/export/share/ruimeng/env/anaconda/envs/ir/lib/python3.8/site-packages/datasets/formatting/formatting.py", line 391, in format_row
>     raise TypeError(
> TypeError: Custom formatting function must return a dict to be able to pick a row, but got None
> ```

I found this trick can be helpful: return an empty dict in exception:
```
def transform_fn(example):
    try:
        process_your_data(example)
    except Exception as e:
        print(e)
        return {'input_ids': [[]], 'token_type_ids': [[]], 'attention_mask': [[]]}
train_dataset = datasets.load_dataset(...)
train_dataset = train_dataset.with_transform(parse_fn)
```