Allow concatenation of both in-memory and on-disk datasets
Author: lhoestq
State: closed

This is a prerequisite for the addition of the `add_item` feature (see #1870).
Currently there is one assumption that we would need to change: a dataset is either fully in memory (dataset._data_files is empty), or the dataset can be reloaded from disk (using the dataset._data_files).
This assumption is used for pickling for example:
- in-memory dataset can just be pickled/unpickled in-memory
- on-disk dataset can be unloaded to only keep the filepaths when pickling, and then reloaded from the disk when unpickling

Maybe let's have a design that allows a Dataset to have a Table that can be rebuilt from heterogenous sources like in-memory tables or on-disk tables ? This could also be further extended in the future

One idea would be to define a list of sources and each source implements a way to reload its corresponding pyarrow Table.
Then the dataset would be the concatenation of all these tables.

Depending on the source type, the serialization using pickle would be different. In-memory data would be copied while on-disk data would simply be replaced by the path to these data.

If you have some ideas you would like to share about the design/API feel free to do so :)

cc @albertvillanova 

Comments: I started working on this. My idea is to first add the pyarrow Table wrappers InMemoryTable and MemoryMappedTable that both implement what's necessary regarding copy/pickle. Then have another wrapper that takes the concatenation of InMemoryTable/MemoryMappedTable objects.

What's important here is that concatenating two tables into one doesn't double the memory used (`total_allocated_bytes()` stays the same).

Hi @lhoestq @albertvillanova,

I checked the linked issues and PR, this seems like a great idea. Would you mind elaborating on the in-memory and memory-mapped datasets? 
Based on my understanding, it is something like this, please correct me if I am wrong:
1. For in-memory datasets, we don't have any dataset files so the entire dataset is pickled to the cache during loading, and then whenever required it is unpickled .
2. For on-disk/memory-mapped datasets, we have the data files provided, so they can be re-loaded from the paths, and only the file-paths are stored while pickling.

If this is correct, will the feature also handle pickling/unpickling of a concatenated dataset? Will this be cached?

This also leads me to ask whether datasets are chunked during pickling? 

Thanks,
Gunjan

Hi ! Yes you're totally right about your two points :)

And in the case of a concatenated dataset, then we should reload each sub-table depending on whether it's in-memory or memory mapped. That means the dataset will be made of several blocks in order to keep track of what's from memory and what's memory mapped. This allows to pickle/unpickle concatenated datasets

Hi @lhoestq

Thanks, that sounds nice. Can you explain where the issue of the double memory may arise? Also, why is the existing `concatenate_datasets` not sufficient for this purpose?

Hi @lhoestq,

Will the `add_item` feature also help with lazy writing (or no caching) during `map`/`filter`?

> Can you explain where the issue of the double memory may arise?

We have to keep each block (in-memory vs memory mapped) separated in order to be able to reload them with pickle.
On the other hand we also need to have the full table from mixed in-memory and memory mapped data in order to iterate or extract data conveniently. That means that each block is accessible twice: once in the full table, and once in the separated blocks. But since pyarrow tables concatenation doesn't double the memory, then building the full table doesn't cost memory which is what we want :)

> Also, why is the existing concatenate_datasets not sufficient for this purpose?

The existing `concatenate_datasets` doesn't support having both in-memory and memory mapped data together (there's no fancy block separation logic). It works for datasets fully in-memory or fully memory mapped but not a mix of the two.

> Will the add_item feature also help with lazy writing (or no caching) during map/filter?

It will enable the implementation of the fast, masked filter from this discussion: https://github.com/huggingface/datasets/issues/1949
However I don't think this will affect map.