Add Arrow type casting to struct for Image and Audio + Support nested casting
Opened by lhoestq (closed)

## Intro

1. Currently, it's not possible to have nested features containing Audio or Image. 

2. Moreover one can keep an Arrow array as a StringArray to store paths to images, but such arrays can't be directly concatenated to another image array if it's stored an another Arrow type (typically, a StructType).

3. Allowing several Arrow types for a single HF feature type also leads to bugs like this one #3497 

4. Issues like #3247 are quite frequent and happen when Arrow fails to reorder StructArrays. 

5. Casting Audio feature type is blocking preparation for the ASR task template: https://github.com/huggingface/datasets/pull/3364

All those issues are linked together by the fact that:
- we are limited by the Arrow type casting which is lacking features for nested types.
- and especially for Audio and Image: they are not robust enough for concatenation and feature inference.

## Proposed solution

To fix 1 and 4 I implemented nested array type casting (which is missing in PyArrow).

To fix 2, 3 and 5 while having a simple implementation for nested array type casting, I changed the storage type of Audio and Image to always be a StructType. Also casting from StringType is directly implemented via a new function `cast_storage` that is defined individually for Audio and Image. I also added nested decoding.

## Implementation details

### I. Better Arrow data type casting for nested data structures

I implemented new functions `array_cast` and `table_cast` that do the exact same as `pyarrow.Array.cast` or `pyarrow.Table.cast` but support nested struct casting and array re-ordering.

These functions can be used on PyArrow objects, and are already integrated in our own `datasets.table.Table.cast` functions. So one can do `my_dataset.data.cast(pyarrow_schema_with_custom_hf_types)` directly.

### II. New image and audio extension types with custom casting

I used PyArrow extension types to be able to define what casting is allowed or not. For example both StringType->ImageExtensionType and StructType->ImageExtensionType are allowed, via the `cast_storage` method.

I factorized all the PyArrow + Pandas extension stuff in the `base_extension.py` file. This aims at separating the front-facing API code of `datasets` from the Arrow back-end which requires advanced knowledge.

### III. Nested feature decoding

I added a new function `decode_nested_example` to decode image and audio data in nested data structures. For optimization's sake, this function is only called if a column has at least one feature that requires decoding.

## Alternative considered

The casting to struct type could have been done directly with python objects using some Audio and Image methods, but bringing arrow data to python objects is expensive. The Audio and Image types could also have been able to convert the arrow data directly, but this is not convenient to use when casting a full Arrow Table with nested fields. Therefore I decided to keep the Arrow data casting logic in Arrow extension types.

## Future work

This work can be used to allow the ArrayND feature types to be nested too (see issue #887)

## TODO

- [x] fix current tests
- [x] add new tests
- [x] docstrings/comments

Comments:
Regarding the tests I'm just missing the FixedSizeListType type casting for ListArray objects, will to it tomorrow as well as adding new tests + docstrings

and also adding soundfile in the CI

While writing some tests I noticed that the ExtensionArray can't be directly concatenated - maybe we can get rid of the extension types/arrays and only keep their storages in native arrow types.

In this case the `cast_storage` functions should be the responsibility of the Image and Audio classes directly. And therefore we would need to never cast to a pyarrow type again but to a HF feature - since they'd end up being the one able to tell what's castable or not. This is fine in my opinion but let me know what you think. I can take care of this on monday I think

Alright I got rid of all the extension type stuff, I'm writing the new tests now :)

Tests are done, I'll finish the comments and docstrings tomorrow and set the PR on ready for review once it's done !

> While writing some tests I noticed that the ExtensionArray can't be directly concatenated - maybe we can get rid of the extension types/arrays and only keep their storages in native arrow types.
>
>In this case the cast_storage functions should be the responsibility of the Image and Audio classes directly. And therefore we would need two never cast to a pyarrow type again but to a HF feature - since they'd end up being the one able to tell what's castable or not. This is fine in my opinion but let me know what you think. I can take care of this on monday I think

Does this change affect performance?

> Does this change affect performance?

In general it shouldn't have a significant impact on performance since the structure of the features is rarely complex (in general we have <20 features and <4 levels of nesting)

Regarding Audio and Image specifically, casting from a StringArray is a little bit more costly since it creates the "bytes" BinaryArray with `None` values with the same length as the "path" array. From the tests I did locally this is very fast though and shouldn't affect the user experience at the current scale of the audio/image datasets we have. It also requires a little bit of RAM though


Alright this is ready for review now ! Let me know if you have comments and/or improvements :)

I am facing the issue ArrowNotImplementedError  but no solution is working. Please help me

Can you open an new issue and share the error message as well as the script you used ? We'd be happy to help :)