[`to_json`] add multi-proc sharding support
CONTRIBUTOR
Hi @stas00, 
I want to work on this issue and I was thinking why don't we use `imap` [in this loop](https://github.com/huggingface/datasets/blob/440b14d0dd428ae1b25881aa72ba7bbb8ad9ff84/src/datasets/io/json.py#L99)? This way, using offset (which is being used to slice the pyarrow table) we can convert  pyarrow table to `json` using multiprocessing. I've a small code snippet for some clarity:
```
result = list(
            pool.imap(self._apply_df, [(offset, batch_size) for offset in range(0, len(self.dataset), batch_size)])
        )
```
`_apply_df` is a function which will return `batch.to_pandas().to_json(path_or_buf=None, orient="records", lines=True)` which is basically json version of the batched pyarrow table. Later on we can concatenate it to form json file? 

I think the only downside here is to write file from `imap` output (output would be a list and we'll need to iterate over it and write in a file) which might add a little overhead cost. What do you think about this?
2021-07-31T20:35:13Z
https://github.com/huggingface/datasets/issues/2663#issuecomment-890401313
https://api.github.com/repos/huggingface/datasets/issues/2663
IC_kwDODunzps41EnIh
https://api.github.com/repos/huggingface/datasets/issues/comments/890401313/reactions
2021-08-01T17:26:06Z
https://api.github.com/repos/huggingface/datasets/issues/comments/890401313
https://avatars.githubusercontent.com/u/19718818?v=4
https://api.github.com/users/bhavitvyamalik/events{/privacy}
https://api.github.com/users/bhavitvyamalik/followers
https://api.github.com/users/bhavitvyamalik/following{/other_user}
https://api.github.com/users/bhavitvyamalik/gists{/gist_id}
https://github.com/bhavitvyamalik
bhavitvyamalik
MDQ6VXNlcjE5NzE4ODE4
https://api.github.com/users/bhavitvyamalik/orgs
https://api.github.com/users/bhavitvyamalik/received_events
https://api.github.com/users/bhavitvyamalik/repos
https://api.github.com/users/bhavitvyamalik/starred{/owner}{/repo}
https://api.github.com/users/bhavitvyamalik/subscriptions
User
https://api.github.com/users/bhavitvyamalik
public
CONTRIBUTOR
Followed up in https://github.com/huggingface/datasets/pull/2747
2021-08-03T21:29:21Z
https://github.com/huggingface/datasets/issues/2663#issuecomment-892178828
https://api.github.com/repos/huggingface/datasets/issues/2663
IC_kwDODunzps41LZGM
https://api.github.com/repos/huggingface/datasets/issues/comments/892178828/reactions
2021-08-03T21:29:21Z
https://api.github.com/repos/huggingface/datasets/issues/comments/892178828
https://avatars.githubusercontent.com/u/10676103?v=4
https://api.github.com/users/stas00/events{/privacy}
https://api.github.com/users/stas00/followers
https://api.github.com/users/stas00/following{/other_user}
https://api.github.com/users/stas00/gists{/gist_id}
https://github.com/stas00
stas00
MDQ6VXNlcjEwNjc2MTAz
https://api.github.com/users/stas00/orgs
https://api.github.com/users/stas00/received_events
https://api.github.com/users/stas00/repos
https://api.github.com/users/stas00/starred{/owner}{/repo}
https://api.github.com/users/stas00/subscriptions
User
https://api.github.com/users/stas00
public
As discussed on slack it appears that `to_json` is quite slow on huge datasets like OSCAR.

I implemented sharded saving, which is much much faster - but the tqdm bars all overwrite each other, so it's hard to make sense of the progress, so if possible ideally this multi-proc support could be implemented internally in `to_json` via `num_proc` argument. I guess `num_proc` will be the number of shards?

I think the user will need to use this feature wisely, since too many processes writing to say normal style HD is likely to be slower than one process.

I'm not sure whether the user should be responsible to concatenate the shards at the end  or `datasets`, either way works for my needs.

The code I was using:

```
from multiprocessing import cpu_count, Process, Queue

[...]

filtered_dataset = concat_dataset.map(filter_short_documents, batched=True, batch_size=256, num_proc=cpu_count())

DATASET_NAME = "oscar"
SHARDS = 10
def process_shard(idx):
    print(f"Sharding {idx}")
    ds_shard = filtered_dataset.shard(SHARDS, idx, contiguous=True)
    # ds_shard = ds_shard.shuffle() # remove contiguous=True above if shuffling
    print(f"Saving {DATASET_NAME}-{idx}.jsonl")
    ds_shard.to_json(f"{DATASET_NAME}-{idx}.jsonl", orient="records", lines=True, force_ascii=False)

queue = Queue()
processes = [Process(target=process_shard, args=(idx,)) for idx in range(SHARDS)]
for p in processes:
    p.start()

for p in processes:
    p.join()
```

Thank you!

@lhoestq 