Dataset.map gets stuck on _cast_to_python_objects
Opened by srobertjames (closed)

## Describe the bug
`Dataset.map`, when fed a Huggingface Tokenizer as its map func, can sometimes spend huge amounts of time doing casts.  A minimal example follows.

Not all usages suffer from this.  For example, I profiled the preprocessor at https://github.com/huggingface/notebooks/blob/main/examples/question_answering.ipynb , and it did _not_ have this problem.  However, I'm at a loss to figure out how it avoids it, as the example below is simple and minimal and still has this problem.

This casting, where it occurs, causes the `Dataset.map` to run approximately 7x slower than it runs for code which does not cause this casting.

This may be related to https://github.com/huggingface/datasets/issues/1046 .  However, the tokenizer is _not_ set to return Tensors.

## Steps to reproduce the bug
A minimal, self-contained example to reproduce is below:
```python
import transformers
from transformers import AutoTokenizer
from datasets import load_dataset
import torch
import cProfile

pretrained = 'distilbert-base-uncased'
tokenizer = AutoTokenizer.from_pretrained(pretrained)

squad = load_dataset('squad')
squad_train = squad['train']
squad_tiny = squad_train.select(range(5000))

assert isinstance(tokenizer, transformers.PreTrainedTokenizerFast)

def tokenize(ds):
        tokens = tokenizer(text=ds['question'],
                                text_pair=ds['context'],
                                add_special_tokens=True,
                                padding='max_length',
                                truncation='only_second',
                                max_length=160,
                                stride=32,
                                return_overflowing_tokens=True,
                                return_offsets_mapping=True,
                                )
        return tokens

cmd = 'squad_tiny.map(tokenize, batched=True, remove_columns=squad_tiny.column_names)'
cProfile.run(cmd, sort='tottime')
```

## Actual results
The code works, but takes 10-25 sec per batch (about 7x slower than non-casting code), with the following profile.  Note that `_cast_to_python_objects` is the culprit.
```
    63524075 function calls (58206482 primitive calls) in 121.836 seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
5274034/40   68.751    0.000  111.060    2.776 features.py:262(_cast_to_python_objects)
 42223832   24.077    0.000   33.310    0.000 {built-in method builtins.isinstance}
 16338/20    5.121    0.000  111.053    5.553 features.py:361(<listcomp>)
  5274135    4.747    0.000    4.749    0.000 {built-in method _abc._abc_instancecheck}
    80/40    4.731    0.059  116.292    2.907 {pyarrow.lib.array}
  5274135    4.485    0.000    9.234    0.000 abc.py:96(__instancecheck__)
2661564/2645196    2.959    0.000    4.298    0.000 features.py:1081(_check_non_null_non_empty_recursive)
        5    2.786    0.557    2.786    0.557 {method 'encode_batch' of 'tokenizers.Tokenizer' objects}
  2668052    0.930    0.000    0.930    0.000 {built-in method builtins.len}
     5000    0.930    0.000    0.938    0.000 tokenization_utils_fast.py:187(_convert_encoding)
        5    0.750    0.150    0.808    0.162 {method 'to_pydict' of 'pyarrow.lib.Table' objects}
        1    0.444    0.444  121.749  121.749 arrow_dataset.py:2501(_map_single)
       40    0.375    0.009  116.291    2.907 arrow_writer.py:151(__arrow_array__)
       10    0.066    0.007    0.066    0.007 {method 'write_batch' of 'pyarrow.lib._CRecordBatchWriter' objects}
        1    0.060    0.060  121.835  121.835 fingerprint.py:409(wrapper)
11387/5715    0.049    0.000    0.175    0.000 {built-in method builtins.getattr}
       36    0.049    0.001    0.049    0.001 {pyarrow._compute.call_function}
    15000    0.040    0.000    0.040    0.000 _collections_abc.py:719(__iter__)
        3    0.023    0.008    0.023    0.008 {built-in method _imp.create_dynamic}
       77    0.020    0.000    0.020    0.000 {built-in method builtins.dir}
       37    0.019    0.001    0.019    0.001 socket.py:543(send)
       15    0.017    0.001    0.017    0.001 tokenization_utils_fast.py:460(<listcomp>)
  432/421    0.015    0.000    0.024    0.000 traitlets.py:1388(_notify_observers)
     5000    0.015    0.000    0.018    0.000 _collections_abc.py:672(keys)
       51    0.014    0.000    0.042    0.001 traitlets.py:276(getmembers)
        5    0.014    0.003    3.775    0.755 tokenization_utils_fast.py:392(_batch_encode_plus)
      3/1    0.014    0.005    0.035    0.035 {built-in method _imp.exec_dynamic}
        5    0.012    0.002    0.950    0.190 tokenization_utils_fast.py:438(<listcomp>)
    31626    0.012    0.000    0.012    0.000 {method 'append' of 'list' objects}
1532/1001    0.011    0.000    0.189    0.000 traitlets.py:643(get)
        5    0.009    0.002    3.796    0.759 arrow_dataset.py:2631(apply_function_on_filtered_inputs)
       51    0.009    0.000    0.062    0.001 traitlets.py:1766(traits)
        5    0.008    0.002    3.784    0.757 tokenization_utils_base.py:2632(batch_encode_plus)
      368    0.007    0.000    0.044    0.000 traitlets.py:1715(_get_trait_default_generator)
       26    0.007    0.000    0.022    0.001 traitlets.py:1186(setup_instance)
       51    0.006    0.000    0.010    0.000 traitlets.py:1781(<listcomp>)
    80/32    0.006    0.000    0.052    0.002 table.py:1758(cast_array_to_feature)
      684    0.006    0.000    0.007    0.000 {method 'items' of 'dict' objects}
4344/1794    0.006    0.000    0.192    0.000 traitlets.py:675(__get__)
...
```
## Environment info
I observed this on both Google colab and my local workstation:

### Google colab
- `datasets` version: 2.3.2
- Platform: Linux-5.4.188+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.13
- PyArrow version: 6.0.1
- Pandas version: 1.3.5

### Local
- `datasets` version: 2.3.2
- Platform: Windows-7-6.1.7601-SP1
- Python version: 3.8.10
- PyArrow version: 8.0.0
- Pandas version: 1.4.3


Comments: Are you able to reproduce this? My example is small enough that it should be easy to try.

Hi! Thanks for reporting and providing a reproducible example. Indeed, by default, `datasets` performs an expensive cast on the values returned by `map` to convert them to one of the types supported by PyArrow (the underlying storage format used by `datasets`). This cast is not needed on NumPy arrays as PyArrow supports them natively, so one way to make this transform faster is to add `return_tensors="np"` to the tokenizer call. 

I think we should mention this in the docs (cc @stevhliu)

I tested this tokenize function and indeed noticed a casting. However it seems to only concerns the `offset_mapping` field, which contains a list of tuples, that is converted to a list of lists. Since `pyarrow` also supports tuples, we actually don't need to convert the tuples to lists. 

I think this can be changed here: 

https://github.com/huggingface/datasets/blob/ede72d3f9796339701ec59899c7c31d2427046fb/src/datasets/features/features.py#L382-L383

```diff
- if isinstance(obj, list): 
+ if isinstance(obj, (list, tuple)): 
```

and here: 

https://github.com/huggingface/datasets/blob/ede72d3f9796339701ec59899c7c31d2427046fb/src/datasets/features/features.py#L386-L387

```diff
- return obj if isinstance(obj, list) else [], isinstance(obj, tuple)
+ return obj, False
```

@srobertjames can you try applying these changes and let us know if it helps ? If so, feel free to open a Pull Request to contribute this improvement if you want :)

Wow, adding `return_tensors="np"` sped up my example by a **factor 17x** of and completely eliminated the casting! I'd recommend not only to document it, but to make that the default.

The code at https://github.com/huggingface/notebooks/blob/main/examples/question_answering.ipynb does not specify `return_tensors="np"` but yet avoids the casting penalty.  How does it do that?  (The ntbk seems to do `return_overflowing_tokens=True, return_offsets_mapping=True,`).

Also, surprisingly enough, using `return_tensors="pt"` (which is my eventual application) yields this error:
```
TypeError: Provided `function` which is applied to all elements of table returns a `dict` of types 
[<class 'torch.Tensor'>, <class 'torch.Tensor'>, <class 'torch.Tensor'>, <class 'torch.Tensor'>]. 
When using `batched=True`, make sure provided `function` returns a `dict` of types like 
`(<class 'list'>, <class 'numpy.ndarray'>)`.
```

Setting the output to `"np"` makes the whole pipeline fast because it moves the data buffers from rust to python to arrow using zero-copy, and also because it does eliminate the casting completely ;)

Have you had a chance to try eliminating the tuple casting using the trick above ?

@lhoestq I just benchmarked the two edits to `features.py` above, and they appear to solve the problem, bringing my original example to within 20% the speed of the output `"np"` example.  Nice!

For a pull request, do you suggest simply following https://github.com/huggingface/datasets/blob/main/CONTRIBUTING.md ?

Cool ! Sure feel free to follow these instructions to open a PR :) thanks !

#take

Resolved via https://github.com/huggingface/datasets/pull/4993.