GitHub issue #5541: Flattening indices in selected datasets is extremely inefficient
State: closed | Created: 2023-02-17T01:52:24Z

### Describe the bug

If we perform a `select`  (or `shuffle`, `train_test_split`, etc.) operation on a dataset , we end up with a dataset with an `indices_table`. Currently, flattening such dataset consumes a lot of memory and the resulting flat dataset contains ChunkedArrays with as many chunks as there are rows. This is extremely inefficient and slows down the operations on the flat dataset, e.g., saving/loading the dataset to disk becomes really slow. 

Perhaps more importantly, loading the dataset back from disk basically loads the whole table into RAM, as it cannot take advantage of memory mapping.

### Steps to reproduce the bug

The following script reproduces the issue:

```python
import gc
import os
import psutil
import tempfile
import time

from datasets import Dataset

DATASET_SIZE = 5000000


def profile(func):
    def wrapper(*args, **kwargs):
        mem_before = psutil.Process(os.getpid()).memory_info().rss / (1024 * 1024)
        start = time.time()

        # Run function here
        out = func(*args, **kwargs)

        end = time.time()
        mem_after = psutil.Process(os.getpid()).memory_info().rss / (1024 * 1024)
        print(f"{func.__name__} -- RAM memory used: {mem_after - mem_before} MB -- Total time: {end - start:.6f} s")

        return out

    return wrapper


def main():
    ds = Dataset.from_list([{'col': i} for i in range(DATASET_SIZE)])
    print(f"Num chunks for original ds: {ds.data['col'].num_chunks}")

    with tempfile.TemporaryDirectory() as tmpdir:
        path1 = os.path.join(tmpdir, 'ds1')
        print("Original ds save/load")
        profile(ds.save_to_disk)(path1)

        ds_loaded = profile(Dataset.load_from_disk)(path1)
        print(f"Num chunks for original ds after reloading: {ds_loaded.data['col'].num_chunks}")
        print("")

        ds_select = ds.select(reversed(range(len(ds))))
        print(f"Num chunks for selected ds: {ds_select.data['col'].num_chunks}")
        del ds
        del ds_loaded
        gc.collect()

        # This would happen anyway when we call save_to_disk
        ds_select = profile(ds_select.flatten_indices)()
        print(f"Num chunks for selected ds after flattening: {ds_select.data['col'].num_chunks}")
        print("")

        path2 = os.path.join(tmpdir, 'ds2')
        print("Selected ds save/load")
        profile(ds_select.save_to_disk)(path2)
        del ds_select
        gc.collect()

        ds_select_loaded = profile(Dataset.load_from_disk)(path2)
        print(f"Num chunks for selected ds after reloading: {ds_select_loaded.data['col'].num_chunks}")


if __name__ == '__main__':
    main()
```

Sample result:
```
Num chunks for original ds: 1
Original ds save/load
save_to_disk -- RAM memory used: 0.515625 MB -- Total time: 0.253888 s
load_from_disk -- RAM memory used: 42.765625 MB -- Total time: 0.015176 s
Num chunks for original ds after reloading: 5000

Num chunks for selected ds: 1
flatten_indices -- RAM memory used: 4852.609375 MB -- Total time: 46.116774 s
Num chunks for selected ds after flattening: 5000000

Selected ds save/load
save_to_disk -- RAM memory used: 1326.65625 MB -- Total time: 42.309825 s
load_from_disk -- RAM memory used: 2085.953125 MB -- Total time: 11.659137 s
Num chunks for selected ds after reloading: 5000000
```

### Expected behavior

Saving/loading the dataset should be much faster and consume almost no extra memory thanks to pyarrow memory mapping.

### Environment info

- `datasets` version: 2.9.1.dev0
- Platform: macOS-13.1-arm64-arm-64bit
- Python version: 3.10.8
- PyArrow version: 11.0.0
- Pandas version: 1.5.3