`.to_json` is extremely slow after `.select`
MEMBER
Hi ! It's slower indeed because a datasets on which `select`/`shard`/`train_test_split`/`shuffle` has been called has to do additional steps to retrieve the data of the dataset table in the right order.

Indeed, if you call `dataset.select([0, 5, 10])`, the underlying table of the dataset is not altered to keep the examples at index 0, 5, and 10. Instead, an indices mapping is added on top of the table, that says that the first example is at index 0, the second at index 5 and the last one at index 10.

Therefore accessing the examples of the dataset is slower because of the additional step that uses the indices mapping.

The step that takes the most time is to query the dataset table from a list of indices here:

https://github.com/huggingface/datasets/blob/047dc756ed20fbf06e6bcaf910464aba0e20610a/src/datasets/formatting/formatting.py#L61-L63

In your case it can be made significantly faster by checking if the indices are contiguous. If they're contiguous, we could pass a python `slice` or `range` instead of a list of integers to `_query_table`. This way `_query_table` will do only one lookup to get the queried batch instead of `batch_size` lookups.

Given that calling `select` with contiguous indices is a common use case I'm in favor of implementing such an optimization :)
Let me know what you think
2021-12-14T13:53:38Z
https://github.com/huggingface/datasets/issues/3419#issuecomment-993564758
https://api.github.com/repos/huggingface/datasets/issues/3419
IC_kwDODunzps47OJhW
https://api.github.com/repos/huggingface/datasets/issues/comments/993564758/reactions
2021-12-14T13:53:38Z
https://api.github.com/repos/huggingface/datasets/issues/comments/993564758
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
CONTRIBUTOR
Hi, thanks for the response!
I still don't understand why it is so much slower than iterating and saving:
```python
from datasets import load_dataset

original = load_dataset("squad", split="train")
original.to_json("from_original.json")  # Takes 0 seconds

selected_subset1 = original.select([i for i in range(len(original))])
selected_subset1.to_json("from_select1.json")  # Takes 99 seconds

selected_subset2 = original.select([i for i in range(int(len(original) / 2))])
selected_subset2.to_json("from_select2.json")  # Takes 47 seconds

selected_subset3 = original.select([i for i in range(len(original)) if i % 2 == 0])
selected_subset3.to_json("from_select3.json")  # Takes 49 seconds

import json
import time
def fast_to_json(dataset, path):
    start = time.time()
    with open(path, mode="w") as f:
        for example in dataset:
            f.write(json.dumps(example, separators=(',', ':')) + "\n")
    end = time.time()
    print(f"Saved {len(dataset)} examples to {path} in {end - start} seconds.")

fast_to_json(original, "from_original_fast.json")
fast_to_json(selected_subset1, "from_select1_fast.json")
fast_to_json(selected_subset2, "from_select2_fast.json")
fast_to_json(selected_subset3, "from_select3_fast.json")
```
```
Saved 87599 examples to from_original_fast.json in 8 seconds.
Saved 87599 examples to from_select1_fast.json in 10 seconds.
Saved 43799 examples to from_select2_fast.json in 6 seconds.
Saved 43800 examples to from_select3_fast.json in 5 seconds.
```
2021-12-14T23:07:44Z
https://github.com/huggingface/datasets/issues/3419#issuecomment-994124249
https://api.github.com/repos/huggingface/datasets/issues/3419
IC_kwDODunzps47QSHZ
https://api.github.com/repos/huggingface/datasets/issues/comments/994124249/reactions
2021-12-14T23:07:44Z
https://api.github.com/repos/huggingface/datasets/issues/comments/994124249
https://avatars.githubusercontent.com/u/13485709?v=4
https://api.github.com/users/eladsegal/events{/privacy}
https://api.github.com/users/eladsegal/followers
https://api.github.com/users/eladsegal/following{/other_user}
https://api.github.com/users/eladsegal/gists{/gist_id}
https://github.com/eladsegal
eladsegal
MDQ6VXNlcjEzNDg1NzA5
https://api.github.com/users/eladsegal/orgs
https://api.github.com/users/eladsegal/received_events
https://api.github.com/users/eladsegal/repos
https://api.github.com/users/eladsegal/starred{/owner}{/repo}
https://api.github.com/users/eladsegal/subscriptions
User
https://api.github.com/users/eladsegal
public
MEMBER
There are slight differences between what you're doing and what `to_json` is actually doing.
In particular `to_json` currently converts batches of rows (as an arrow table) to a pandas dataframe, and then to JSON Lines. From your benchmark it looks like it's faster if we don't use pandas.

Thanks for investigating, I think we can optimize `to_json` significantly thanks to your test.
2021-12-15T18:44:49Z
https://github.com/huggingface/datasets/issues/3419#issuecomment-995070890
https://api.github.com/repos/huggingface/datasets/issues/3419
IC_kwDODunzps47T5Oq
https://api.github.com/repos/huggingface/datasets/issues/comments/995070890/reactions
2021-12-15T18:44:57Z
https://api.github.com/repos/huggingface/datasets/issues/comments/995070890
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
CONTRIBUTOR
Thanks for your observations, @eladsegal! I spent some time with this and tried different approaches. Turns out that https://github.com/huggingface/datasets/blob/bb13373637b1acc55f8a468a8927a56cf4732230/src/datasets/io/json.py#L100 is giving the problem when we use `to_json` after `select`. This is when `indices` parameter in `query_table` is not `None` (if it is `None` then `to_json` should work as expected)

In order to circumvent this problem, I found out instead of doing Arrow Table -> Pandas-> JSON we can directly go to JSON by using `to_pydict()` which is a little slower than the current approach but at least `select` works properly now. Lmk what you guys think of it @lhoestq, @eladsegal?
2021-12-17T16:03:49Z
https://github.com/huggingface/datasets/issues/3419#issuecomment-996837140
https://api.github.com/repos/huggingface/datasets/issues/3419
IC_kwDODunzps47aocU
https://api.github.com/repos/huggingface/datasets/issues/comments/996837140/reactions
2021-12-18T08:42:51Z
https://api.github.com/repos/huggingface/datasets/issues/comments/996837140
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
MEMBER
Sounds good to me ! Feel free to also share your benchmarks for reference @bhavitvyamalik 
2021-12-20T11:51:25Z
https://github.com/huggingface/datasets/issues/3419#issuecomment-997855339
https://api.github.com/repos/huggingface/datasets/issues/3419
IC_kwDODunzps47ehBr
https://api.github.com/repos/huggingface/datasets/issues/comments/997855339/reactions
2021-12-20T11:51:25Z
https://api.github.com/repos/huggingface/datasets/issues/comments/997855339
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
CONTRIBUTOR
Posting it in @eladsegal's format:

For `squad`:
Saving examples using current `to_json` in 3.63 secs
Saving examples to `from_select1_fast.json` in 5.00 secs
Saving examples to `from_select2_fast.json` in 2.45 secs
Saving examples to `from_select3_fast.json` in 2.50 secs

For `squad_v2`:
Saving examples using current `to_json` in 5.26 secs
Saving examples to `from_select1_fast.json` in 7.54 secs
Saving examples to `from_select2_fast.json` in 3.80 secs
Saving examples to `from_select3_fast.json` in 3.67 secs
2021-12-21T15:49:07Z
https://github.com/huggingface/datasets/issues/3419#issuecomment-998889996
https://api.github.com/repos/huggingface/datasets/issues/3419
IC_kwDODunzps47idoM
https://api.github.com/repos/huggingface/datasets/issues/comments/998889996/reactions
2021-12-21T15:49:07Z
https://api.github.com/repos/huggingface/datasets/issues/comments/998889996
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
## Describe the bug
Saving a dataset to JSON with `to_json` is extremely slow after using `.select` on the original dataset.

## Steps to reproduce the bug
```python
from datasets import load_dataset

original = load_dataset("squad", split="train")
original.to_json("from_original.json")  # Takes 0 seconds

selected_subset1 = original.select([i for i in range(len(original))])
selected_subset1.to_json("from_select1.json")  # Takes 212 seconds

selected_subset2 = original.select([i for i in range(int(len(original) / 2))])
selected_subset2.to_json("from_select2.json")  # Takes 90 seconds
```

## Environment info
<!-- You can run the command `datasets-cli env` and copy-and-paste its output below. -->
- `datasets` version: master (https://github.com/huggingface/datasets/commit/6090f3cfb5c819f441dd4a4bb635e037c875b044)
- Platform: Linux-4.4.0-19041-Microsoft-x86_64-with-glibc2.27
- Python version: 3.9.7
- PyArrow version: 6.0.0
