Optimize contiguous shard and select
NONE
_The documentation is not available anymore as the PR was closed or merged._
2022-06-09T13:53:50Z
https://github.com/huggingface/datasets/pull/4466#issuecomment-1151151184
https://api.github.com/repos/huggingface/datasets/issues/4466
IC_kwDODunzps5EnSxQ
https://api.github.com/repos/huggingface/datasets/issues/comments/1151151184/reactions
2022-06-14T16:04:30Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1151151184
https://avatars.githubusercontent.com/u/99929124?v=4
https://api.github.com/users/HuggingFaceDocBuilderDev/events{/privacy}
https://api.github.com/users/HuggingFaceDocBuilderDev/followers
https://api.github.com/users/HuggingFaceDocBuilderDev/following{/other_user}
https://api.github.com/users/HuggingFaceDocBuilderDev/gists{/gist_id}
https://github.com/HuggingFaceDocBuilderDev
HuggingFaceDocBuilderDev
U_kgDOBfTMJA
https://api.github.com/users/HuggingFaceDocBuilderDev/orgs
https://api.github.com/users/HuggingFaceDocBuilderDev/received_events
https://api.github.com/users/HuggingFaceDocBuilderDev/repos
https://api.github.com/users/HuggingFaceDocBuilderDev/starred{/owner}{/repo}
https://api.github.com/users/HuggingFaceDocBuilderDev/subscriptions
User
https://api.github.com/users/HuggingFaceDocBuilderDev
public
MEMBER
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">

2022-06-09T13:57:34Z
https://github.com/huggingface/datasets/pull/4466#issuecomment-1151155480
https://api.github.com/repos/huggingface/datasets/issues/4466
IC_kwDODunzps5EnT0Y
https://api.github.com/repos/huggingface/datasets/issues/comments/1151155480/reactions
2022-06-09T13:57:34Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1151155480
https://avatars.githubusercontent.com/u/22957388?v=4
https://api.github.com/users/sayakpaul/events{/privacy}
https://api.github.com/users/sayakpaul/followers
https://api.github.com/users/sayakpaul/following{/other_user}
https://api.github.com/users/sayakpaul/gists{/gist_id}
https://github.com/sayakpaul
sayakpaul
MDQ6VXNlcjIyOTU3Mzg4
https://api.github.com/users/sayakpaul/orgs
https://api.github.com/users/sayakpaul/received_events
https://api.github.com/users/sayakpaul/repos
https://api.github.com/users/sayakpaul/starred{/owner}{/repo}
https://api.github.com/users/sayakpaul/subscriptions
User
https://api.github.com/users/sayakpaul
public
MEMBER
I addressed your comments @albertvillanova , let me know what you think :)
2022-06-14T14:50:49Z
https://github.com/huggingface/datasets/pull/4466#issuecomment-1155300439
https://api.github.com/repos/huggingface/datasets/issues/4466
IC_kwDODunzps5E3HxX
https://api.github.com/repos/huggingface/datasets/issues/comments/1155300439/reactions
2022-06-14T14:50:55Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1155300439
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
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.