[Refactor] Use in-memory/memory-mapped/concatenation tables in Dataset
Opened by lhoestq (closed)

## Intro

Currently there is one assumption that we need to change: a dataset is either fully in memory (dataset._data_files is empty), or the dataset can be reloaded from disk with memory mapping (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

## Issues

Because of this assumption, we can't easily implement methods like `Dataset.add_item` to append more rows to a dataset, or `dataset.add_column` to add a column, since we can't mix data from memory and data from the disk.
Moreover, `concatenate_datasets` doesn't work if the datasets to concatenate are not all from memory, or all form the disk.

## Solution provided in this PR

I changed this by allowing several types of Table to be used in the Dataset object.
More specifically I added three pyarrow Table wrappers: InMemoryTable, MemoryMappedTable and ConcatenationTable.
The in-memory and memory-mapped tables implement the pickling behavior described above.
The ConcatenationTable can be made from several tables (either in-memory or memory mapped) called "blocks". Pickling a ConcatenationTable simply pickles the underlying blocks.

## Implementation details

The three tables classes mentioned above all inherit from a `Table` class defined in `table.py`, which is a wrapper of a pyarrow table. The `Table` wrapper implements all the attributes and methods of the underlying pyarrow table.

Regarding the MemoryMappedTable:
Reloading a pyarrow table from the disk makes you lose all the changes you may have applied (slice, rename_columns, drop, cast etc.). Therefore the MemoryMappedTable implements a "replay" mechanism to re-apply the changes when reloading the pyarrow table from the disk.

## Checklist

- [x] add InMemoryTable
- [x] add MemoryMappedTable
- [x] add ConcatenationTable
- [x] Update the ArrowReader to use these new tables depending on the `in_memory` parameter
- [x] Update Dataset.from_xxx methods
- [x] Update load_from_disk and save_to_disk
- [x] Backward compatibility of load_from_disk
- [x] Add tests for the new tables
- [x] Update current tests
- [ ] Documentation

----------

I would be happy to discuss the design of this PR :)

Close #1877 

Comments: There is one more thing I would love to see. Let's say we iteratively keep updating a data source that loaded from **load_dataset** or **load_from_disk**. Now we need to save it to the same location by overriding the previous file inorder to save the disk space. At the moment **save_to_disk** can not assign a name. So I do not see an easy way to override the previous files.  @lhoestq  is this possible?



p.s one last thing?

Is there a way to flush out any connection to a data source loaded from **load_from_disk** or **load_dataset** methods?  At the moment I suspect when we use any of those functions, it will always keep a pointer although we override it again with a new version of the dataset source. This is really useful in an iterative process. 



> There is one more thing I would love to see. Let's say we iteratively keep updating a data source that loaded from **load_dataset** or **load_from_disk**. Now we need to save it to the same location by overriding the previous file inorder to save the disk space. At the moment **save_to_disk** can not assign a name. So I do not see an easy way to override the previous files. @lhoestq is this possible?

In the new save_to_disk, the filename of the arrow file is fixed: `dataset.arrow`.
This way is will be overwritten if you save your dataset again

> Is there a way to flush out any connection to a data source loaded from **load_from_disk** or **load_dataset** methods? At the moment I suspect when we use any of those functions, it will always keep a pointer although we override it again with a new version of the dataset source. This is really useful in an iterative process.

If you update an arrow file, then you must reload it with `load_from_disk` for example in order to have the updated data.
Does that answer the question ? How does this "pointer" behavior manifest exactly on your side ?

Apparently the usage of the compute layer of pyarrow requires pyarrow>=1.0.0 (otherwise there are some issues on windows with file permissions when doing dataset concatenation).

I'll bump the pyarrow requirement from, 0.17.1 to 1.0.0


> If you update an arrow file, then you must reload it with `load_from_disk` for example in order to have the updated data.
> Does that answer the question? How does this "pointer" behavior manifest exactly on your side?

Yes,  I checked this behavior.. if we update the .arrow file it kind of flushes out the previous one. So your solution is perfect <3. 

Sorry for spamming, there's a a bug that only happens on the CI so I have to re-run it several times

Alright I finally added all the tests I wanted !
I also fixed all the bugs and now all the tests are passing :)

Let me know if you have comments.

I also noticed that two methods in pyarrow seem to bring some data in memory even for a memory mapped table: filter and cast:
- for filter I took a look at the C++ code on the arrow's side and found [this part](https://github.com/apache/arrow/blob/55c8d74d5556b25238fb2028e9fb97290ea24684/cpp/src/arrow/compute/kernels/vector_selection.cc#L93-L160) that "builds" the array during filter. It seems to indicate that it allocates new memory for the filtered array but not 100% sure.
- regarding cast I noticed that it happens when changing the precision of an array of integers. Not sure if there are other cases.


Maybe we'll need to investigate this a bit for your PR on improving `filter` @theo-m , since we don't want to fill the users memory.

> Maybe we'll need to investigate this a bit for your PR on improving `filter` @theo-m , since we don't want to fill the users memory.

I'm a bit unclear on this, I thought the point of the refactor was to use `Table.filter` to speed up our own `.filter` and stop using `.map` that offloaded too much stuff on disk. 
At some point I recall we decided to use `keep_in_memory=True`  as the expectations were that it would be hard to fill the memory?

> I'm a bit unclear on this, I thought the point of the refactor was to use Table.filter to speed up our own .filter and stop using .map that offloaded too much stuff on disk.
> At some point I recall we decided to use keep_in_memory=True as the expectations were that it would be hard to fill the memory?

Yes it's ok to have the mask in memory, but not the full table. I was not aware that the table returned by filter could actually be in memory (it's not part of the pyarrow documentation afaik).
To be more specific I noticed that every time you call `filter`, the pyarrow total allocated memory increases.
I haven't checked on a big dataset though, but it would be nice to see how much memory it uses with respect to the size of the dataset.

I have addressed your comments @theo-m @albertvillanova ! Thanks for the suggestions

I totally agree with you. I would have loved to use inheritance instead.
However because `pa.Table` is a cython class without proper initialization methods (you can't call `__init__` for example): you can't instantiate a subclass of `pa.Table` in python.
To be more specific, you actually can try to instantiate a subclass of `pa.Table` with no data BUT this is not a valid table so you get an error.
And since `pa.Table` objects are immutable you can't even set the data in `__new__` or `__init__`.

EDIT: one could make a new cython class that inherits from `pa.Table` with proper initialization methods, so that we can inherit from this class instead in python. We can do that in the future if we plan to use cython in `datasets`.
(see: https://arrow.apache.org/docs/python/extending.html)

@lhoestq, but in which cases you would like to instantiate directly either `InMemoryTable` or `MemoryMappedTable`? You normally use one of their `from_xxx` class methods...

Yes I was thinking of these cases. The issue is that they return `pa.Table` objects even from a subclass of `pa.Table`

That is indeed a weird behavior...

I guess that in this case, the best approach is as you did, using composition over inheritance...

https://github.com/apache/arrow/pull/5322

@lhoestq I think you forgot to add the new classes to the docs?

Yes you're right, let me add them