GitHub issue #3575: Add Arrow type casting to struct for Image and Audio + Support nested casting
State: closed | Created: 2022-01-13T15:36:59Z

## 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