Adding support for  generic multi dimensional tensors and auxillary image data for multimodal datasets

nlp/features.py:

The main factory class is MultiArray, every single time this class is called, a corresponding pyarrow extension array and type class is generated (and added to the list of globals for future use) for a given root data type and set of dimensions/shape. I provide examples on working with this in datasets/lxmert_pretraining_beta/test_multi_array.py

src/nlp/arrow_writer.py

I had to add a method for writing batches that include extension array types because despite having a unique class for each multidimensional array shape, pyarrow is unable to write any other "array-like" data class to a batch object unless it is of the type pyarrow.ExtensionType. The problem in this is that when writing multiple batches, the order of the schema and data to be written get mixed up (where the pyarrow datatype in the schema only refers to as ExtensionAray, but each ExtensionArray subclass has a different shape) ...  possibly I am missing something here and would be grateful if anyone else could take a look!

datasets/lxmert_pretraining_beta/lxmert_pretraining_beta.py & datasets/lxmert_pretraining_beta/to_arrow_data.py:

I have begun adding the data from the original LXMERT paper (https://arxiv.org/abs/1908.07490)  hosted here: (https://github.com/airsplay/lxmert). The reason I am not pulling from the source of truth for each individual dataset is because it seems that there will also need to be functionality to aggregate multimodal datasets to create a pre-training corpus (:sleepy: ). 
For now, this is just being used to test and run edge-cases for the MultiArray feature, so ive labeled it as "beta_pretraining"!

(still working on the pretraining, just wanted to push out the new functionality sooner than later)


Thank you! I just marked this as a draft PR. It probably would be better to create specific Array2D and Array3D classes as needed instead of a generic MultiArray for now, it should simplify the code a lot too so, I'll update it as such. Also i was meaning to reply earlier, but I wanted to thank you for the testing script you sent me earlier since it ended up being tremendously helpful. 

Okay, I just converted the MultiArray class to Array2D, and got rid of all those "globals()"! 

The main issues I had were that when including a "pa.ExtensionType" as a column, the ordinary methods to batch the data would not work and it would throw me some mysterious error, so I first cleaned up my code to order the row to match the schema (because when including extension types the row is disordered ) and then made each row a pa.Table and then concatenated all the tables. Also each n-dimensional vector class we implement will be size invariant which is some good news. 

Okay awesome! I just added your suggestions and changed up my recursive functions. 

Here is the traceback for the when I use the original code in the write_on_file method:

```
Traceback (most recent call last):
  File "<stdin>", line 33, in <module>
  File "/home/eltoto/nlp/src/nlp/arrow_writer.py", line 214, in finalize
    self.write_on_file()
  File "/home/eltoto/nlp/src/nlp/arrow_writer.py", line 134, in write_on_file
    pa_array = pa.array(self.current_rows, type=self._type)
  File "pyarrow/array.pxi", line 269, in pyarrow.lib.array
  File "pyarrow/array.pxi", line 38, in pyarrow.lib._sequence_to_array
  File "pyarrow/error.pxi", line 106, in pyarrow.lib.check_status
pyarrow.lib.ArrowNotImplementedError: MakeBuilder: cannot construct builder for type extension<arrow.py_extension_type>

shell returned 1
```

I think when trying to cast an extension array within a list of dictionaries, some method gets called that bugs out Arrow and somehow doesn't get called when adding a single row to a a table and then appending multiple tables together. I tinkered with this for a while but could not find any workaround. 

In the case that this new method causes bad compression/worse performance, we can explicitly set the batch size in the pa.Table.to_batches(***batch_size***) method, which will return a list of batches. Perhaps, we can check that the batch size is not too large converting the table to batches after X many rows are appended to it by following the batch_size check below.

> I think when trying to cast an extension array within a list of dictionaries, some method gets called that bugs out Arrow and somehow doesn't get called when adding a single row to a a table and then appending multiple tables together. I tinkered with this for a while but could not find any workaround.

Indeed that's weird.

> In the case that this new method causes bad compression/worse performance, we can explicitly set the batch size in the pa.Table.to_batches(batch_size) method, which will return a list of batches. Perhaps, we can check that the batch size is not too large converting the table to batches after X many rows are appended to it by following the batch_size check below.

The argument of `pa.Table.to_batches` is not `batch_size` but `max_chunksize`, which means that right now it would have no effects (each chunk is of length 1).

We can fix that just by doing `entries.combine_chunks().to_batches(batch_size)`. In that case it would write by chunk of 1000 which is what we want. I don't think it will slow down the writing by much, but we may have to do a benchmark just to make sure. If speed is ok we could even replace the original code to always write chunks this way.

Do you still have errors that need to be fixed ?

@lhoestq Nope all should be good! 

Would you like me to add the entries.combine_chunks().to_batch_size() code + benchmark?

> @lhoestq Nope all should be good!

Awesome :)

I think it would be good to start to add some tests then.
You already have `test_multi_array.py` which is a good start, maybe you can place it in /tests and make it a `unittest.TestCase` ?

> Would you like me to add the entries.combine_chunks().to_batch_size() code + benchmark?

That would be interesting. We don't want reading/writing to be the bottleneck of dataset processing for example in terms of speed. Maybe we could test the write + read speed of different datasets:
- write speed + read speed a dataset with `nlp.Array2D` features
- write speed + read speed a dataset with `nlp.Sequence(nlp.Sequence(nlp.Value("float32")))` features
- write speed + read speed a dataset with `nlp.Sequence(nlp.Value("float32"))` features (same data but flatten)
It will be interesting to see the influence of `.combine_chunks()` on the `Array2D` test too.

What do you think ?

Well actually it looks like we're still having the `print(dataset[0])` error no ?

I just tested your code to try to understand better.


- First thing you must know is that we've switched from `dataset._data.to_pandas` to `dataset._data.to_pydict` by default when we call `dataset[0]` in #423 . Right now it raises an error but it can be fixed by adding this method to `ExtensionArray2D`:

```python
    def to_pylist(self):
        return self.to_numpy().tolist()
```

- Second, I noticed that `ExtensionArray2D.to_numpy()` always return a (5, 5) shape in your example. I thought `ExtensionArray` was for possibly multiple examples and so I was expecting a shape like (1, 5, 5) for example. Did I miss something ?
Therefore when I apply the fix I mentioned (adding to_pylist), it returns one example per row in each image (in your example of 2 images of shape 5x5, I get `len(dataset._data.to_pydict()["image"])  == 10  # True`)

[EDIT] I changed the reshape step in `ExtensionArray2D.to_numpy()` by
```python
numpy_arr = numpy_arr.reshape(len(self), *ExtensionArray2D._construct_shape(self.storage))
```
and it did the job: `len(dataset._data.to_pydict()["image"])  == 2  # True`

- Finally, I was able to make `to_pandas` work though, by implementing custom array dtype in pandas with arrow conversion (I got inspiration from [here](https://gist.github.com/Eastsun/a59fb0438f65e8643cd61d8c98ec4c08) and [here](https://pandas.pydata.org/pandas-docs/version/1.0.0/development/extending.html#compatibility-with-apache-arrow))

Maybe you could add me in your repo so I can open a PR to add these changes to your branch ?

`combine_chunks` doesn't seem to work btw:
`ArrowNotImplementedError: concatenation of extension<arrow.py_extension_type>`

> > @lhoestq Nope all should be good!
> 
> Awesome :)
> 
> I think it would be good to start to add some tests then.
> You already have `test_multi_array.py` which is a good start, maybe you can place it in /tests and make it a `unittest.TestCase` ?
> 
> > Would you like me to add the entries.combine_chunks().to_batch_size() code + benchmark?
> 
> That would be interesting. We don't want reading/writing to be the bottleneck of dataset processing for example in terms of speed. Maybe we could test the write + read speed of different datasets:
> 
> * write speed + read speed a dataset with `nlp.Array2D` features
> * write speed + read speed a dataset with `nlp.Sequence(nlp.Sequence(nlp.Value("float32")))` features
> * write speed + read speed a dataset with `nlp.Sequence(nlp.Value("float32"))` features (same data but flatten)
>   It will be interesting to see the influence of `.combine_chunks()` on the `Array2D` test too.
> 
> What do you think ?

Ya! that should be no problem at all, Ill use the timeit module and get back to you with the results sometime over the weekend.

Thank you for all your help getting the pandas and row indexing for the dataset to work! For `print(dataset[0])`, I considered the workaround of doing  `print(dataset["col_name"][0])` a temporary solution, but ya, I was never able to figure out how to previously get it to work. I'll add you to my repo right now, let me know if you do not see the invite. Also lastly, it is strange how the to_batches method is not working, so I can check that out while I add some speed tests + add the multi dim test under the unit tests this weekend. 

I created the PR :)
I also tested `to_batches` and it works on my side

Sorry for the bit of delay! I just added the tests,  the PR into my fork, and some speed tests. It should be fairly easy to add more tests if we need. Do you think there is anything else to checkout?

Cool thanks for adding the tests :) 

Next step is merge master into this branch.
Not sure I understand what you did in your last commit, but it looks like you discarded all the changes from master ^^'

We've done some changes in the features logic on master, so let me know if you need help merging it.

As soon as we've merged from master, we'll have to make sure that we have extensive tests and we'll be good to do !
About the lxmert dataset, we can probably keep it for another PR as soon as we have working 2d features. What do you think ?

We might want to merge this after tomorrow's release though to avoid potential side effects @lhoestq 

Yep I'm sure we can have it not for tomorrow's release but for the next one ;)

haha, when I tried to rebase, I ran into some conflicts. In that last commit, I restored the features.py from the previous commit on the branch in my fork because upon updating to master, the pandasdtypemanger and pandas extension types disappeared.  If you actually could help me with merging in what is needed, that would actually help a lot. 

Other than that, ya let me go ahead and move the dataloader code out of this PR. Perhaps we could discuss in the slack channelk soon about what to do with that because we can either just support the pretraining corpus for lxmert or try to implement the full COCO and visual genome datasets (+VQA +GQA) which im sure people would be pretty happy about. 

Also we can talk more tests soon too when you are free. 

Goodluck on the release tomorrow guys!

Not sure why github thinks there are conflicts here, as I just rebased from the current master branch.
Merging into master locally works on my side without conflicts
```
git checkout master
git reset --hard origin/master
git merge --no-ff eltoto1219/support_multi_dim_tensors_for_images
Merge made by the 'recursive' strategy.
 datasets/lxmert_pretraining_beta/lxmert_pretraining_beta.py |  89 +++++++++++++++++++++++++++++++++++++
 datasets/lxmert_pretraining_beta/test_multi_array.py        |  45 +++++++++++++++++++
 datasets/lxmert_pretraining_beta/to_arrow_data.py           | 371 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 src/nlp/arrow_dataset.py                                    |  24 +++++-----
 src/nlp/arrow_writer.py                                     |  22 ++++++++--
 src/nlp/features.py                                         | 229 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++---
 tests/test_array_2d.py                                      | 210 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 7 files changed, 969 insertions(+), 21 deletions(-)
 create mode 100644 datasets/lxmert_pretraining_beta/lxmert_pretraining_beta.py
 create mode 100644 datasets/lxmert_pretraining_beta/test_multi_array.py
 create mode 100644 datasets/lxmert_pretraining_beta/to_arrow_data.py
 create mode 100644 tests/test_array_2d.py
```

I put everything inside one commit from the master branch but the merge conflicts on github'side were still there for some reason.
Closing and re-opening the PR fixed the conflict check on github's side.

Almost done ! It still needs a pass on the docs/comments and maybe a few more tests.

I had to do several changes for type inference in the ArrowWriter to make it support custom types.

Ok this is now ready for review ! Thanks for your awesome work in this @eltoto1219 

Summary of the changes:
- added new feature type `Array2D`, that can be instantiated like `Array2D("float32")` for example
- added pyarrow extension type `Array2DExtensionType` and array `Array2DExtensionArray` that take care of converting from and to arrow. `Array2DExtensionType`'s storage is a list of list of any pyarrow array.
- added pandas extension type `PandasArrayExtensionType` and array `PandasArrayExtensionArray` for conversion from and to arrow/python objects
- refactor of the `ArrowWriter` write and write_batch functions to support extension types while preserving the type inference behavior.
- added a utility object `TypedSequence` that is helpful to combine extension arrays and type inference inside the writer's methods.
- added speed test for sequences writing (printed as warnings in pytest)
- breaking: set disable_nullable to False by default as pyarrow's type inference returns nullable fields

And there are plenty of new tests, mainly in `test_array2d.py` and `test_arrow_writer.py`.

Note that there are some collisions in `arrow_dataset.py` with #513 so let's be careful when we'll merge this one.

I know this is a big PR so feel free to ask questions

I'll add Array3D, 4D.. tomorrow but it should take only a few lines. The rest won't change

I took your comments into account and I added Array[3-5]D.
I changed the storage type to fixed lengths lists. I had to update the `to_numpy` function because of that. Indeed slicing a FixedLengthListArray returns a view a of the original array, while in the previous case slicing a ListArray copies the storage.