Title: Optimize contiguous shard and select
Author: lhoestq (MEMBER)
State: closed
Created: 2022-06-09T13:45:39Z

Currently `.shard()` and `.select()` always create an indices mapping. However if the requested data are contiguous, it's much more optimized to simply slice the Arrow table instead of building an indices mapping. In particular:
- the shard/select operation will be much faster
- reading speed will be much faster in the resulting dataset, since it won't have to do a lookup step in the indices mapping

Since `.shard()` is also used for `.map()` with `num_proc>1`, it will also significantly improve the reading speed of multiprocessed `.map()` operations

Here is an example of speed-up:
```python
>>> import io
>>> import numpy as np
>>> from datasets import Dataset
>>> ds = Dataset.from_dict({"a": np.random.rand(10_000_000)})
>>> shard = ds.shard(num_shards=4, index=0, contiguous=True)  # this calls `.select(range(2_500_000))`
>>> buf = io.BytesIO()
>>> %time dd.to_json(buf)
Creating json from Arrow format: 100%|██████████████████| 100/100 [00:00<00:00, 376.17ba/s]
CPU times: user 258 ms, sys: 9.06 ms, total: 267 ms
Wall time: 266 ms
```
while previously it was
```python
Creating json from Arrow format: 100%|███████████████████| 100/100 [00:03<00:00, 29.41ba/s]
CPU times: user 3.33 s, sys: 69.1 ms, total: 3.39 s
Wall time: 3.4 s
```

In this simple case the speed-up is x10, but @sayakpaul experienced a x100 speed-up on its data when exporting to JSON.

## Implementation details

I mostly improved `.select()`: it now checks if the input corresponds to a contiguous chunk of data and then it slices the main Arrow table (or the indices mapping table if it exists). To check if the input indices are contiguous it checks two possibilities:
- if the indices is of type `range`, it checks that start >= 0 and step = 1
- otherwise in the general case, it iterates over the indices. If all the indices are contiguous then we're good, otherwise we have to build an indices mapping.

Having to iterate over the indices doesn't cause performance issues IMO because:
- either they are contiguous and in this case the cost of iterating over the indices is much less than the cost of creating an indices mapping
- or they are not contiguous, and then iterating generally stops quickly when it first encounters the first indice that is not contiguous.

Comments:
_The documentation is not available anymore as the PR was closed or merged._

I thought of just mentioning the benefits I got. Here's the code that @lhoestq provided:

```py
import os
from datasets import load_dataset
from tqdm.auto import tqdm

ds = load_dataset("squad", split="train")
os.makedirs("tmp")

num_shards = 5
for index in tqdm(range(num_shards)):
    size = len(ds) // num_shards
    shard = Dataset(ds.data.slice(size * index, size), fingerprint=f"{ds._fingerprint}_{index}")
    shard.to_json(f"tmp/data_{index}.jsonl")
```

It is 1.64s. Previously the code was:

```py
num_shards = 5
for index in tqdm(range(num_shards)):
    shard = ds.shard(num_shards=num_shards, index=index, contiguous=True)
    shard.to_json(f"tmp/data_{index}.jsonl")
    # upload_to_gcs(f"tmp/data_{index}.jsonl")
```

It was 2min31s. 

I ran it on my humble MacBook Pro:

<img width="574" alt="image" src="https://user-images.githubusercontent.com/22957388/172864881-f1db489a-2305-47f2-a07f-7d3df610b1b8.png">


I addressed your comments @albertvillanova , let me know what you think :)