Repository: huggingface/datasets
Pull request: Add Arrow type casting to struct for Image and Audio + Support nested casting
## 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

Problem statement:
Changing sampling rate in audio dataset and subsequently mapping with `num_proc > 1` leads to weird bug
Running: 

```python
from datasets import load_dataset, DatasetDict
import datasets
from transformers import AutoFeatureExtractor

raw_datasets = DatasetDict()

raw_datasets["train"] = load_dataset("common_voice", "ab", split="train")

feature_extractor = AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-base")

raw_datasets = raw_datasets.cast_column(
    "audio", datasets.features.Audio(sampling_rate=feature_extractor.sampling_rate)
)

num_workers = 16

def prepare_dataset(batch):
    sample = batch["audio"]

    inputs = feature_extractor(sample["array"], sampling_rate=sample["sampling_rate"])
    batch["input_values"] = inputs.input_values[0]
    batch["input_length"] = len(batch["input_values"])
    return batch

raw_datasets.map(
    prepare_dataset,
    remove_columns=next(iter(raw_datasets.values())).column_names,
    num_proc=16,
    desc="preprocess datasets",
)
```

gives

```bash
  File "/home/patrick/experiments/run_bug.py", line 25, in <module>
    raw_datasets.map(
  File "/home/patrick/python_bin/datasets/dataset_dict.py", line 492, in map
    {
  File "/home/patrick/python_bin/datasets/dataset_dict.py", line 493, in <dictcomp>
    k: dataset.map(
  File "/home/patrick/python_bin/datasets/arrow_dataset.py", line 2139, in map
    shards = [
  File "/home/patrick/python_bin/datasets/arrow_dataset.py", line 2140, in <listcomp>
    self.shard(num_shards=num_proc, index=rank, contiguous=True, keep_in_memory=keep_in_memory)
  File "/home/patrick/python_bin/datasets/arrow_dataset.py", line 3164, in shard
    return self.select(
  File "/home/patrick/python_bin/datasets/arrow_dataset.py", line 485, in wrapper
    out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
  File "/home/patrick/python_bin/datasets/fingerprint.py", line 411, in wrapper
    out = func(self, *args, **kwargs)
  File "/home/patrick/python_bin/datasets/arrow_dataset.py", line 2756, in select
    return self._new_dataset_with_indices(indices_buffer=buf_writer.getvalue(), fingerprint=new_fingerprint)
  File "/home/patrick/python_bin/datasets/arrow_dataset.py", line 2667, in _new_dataset_with_indices
    return Dataset(
  File "/home/patrick/python_bin/datasets/arrow_dataset.py", line 659, in __init__
    raise ValueError(
ValueError: External features info don't match the dataset:
Got
{'client_id': Value(dtype='string', id=None), 'path': Value(dtype='string', id=None), 'audio': Audio(sampling_rate=16000, mono=True, _storage_dtype='string', id=None), 'sentence': Value(dtype='string', id=None), 'up_votes': Value(dtype='int64', id=None), 'down_votes': Value(dtype='int64', id=None), 'age': Value(dtype='string', id=None), 'gender': Value(dtype='string', id=None), 'accent': Value(dtype='string', id=None), 'locale': Value(dtype='string', id=None), 'segment': Value(dtype='string', id=None)}
with type
struct<client_id: string, path: string, audio: string, sentence: string, up_votes: int64, down_votes: int64, age: string, gender: string, accent: string, locale: string, segment: string>

but expected something like
{'client_id': Value(dtype='string', id=None), 'path': Value(dtype='string', id=None), 'audio': {'path': Value(dtype='string', id=None), 'bytes': Value(dtype='binary', id=None)}, 'sentence': Value(dtype='string', id=None), 'up_votes': Value(dtype='int64', id=None), 'down_votes': Value(dtype='int64', id=None), 'age': Value(dtype='string', id=None), 'gender': Value(dtype='string', id=None), 'accent': Value(dtype='string', id=None), 'locale': Value(dtype='string', id=None), 'segment': Value(dtype='string', id=None)}
with type
struct<client_id: string, path: string, audio: struct<path: string, bytes: binary>, sentence: string, up_votes: int64, down_votes: int64, age: string, gender: string, accent: string, locale: string, segment: string>
```

Versions:

```python
- `datasets` version: 1.16.2.dev0
- Platform: Linux-5.15.8-76051508-generic-x86_64-with-glibc2.33
- Python version: 3.9.7
- PyArrow version: 6.0.1
```

and `transformers`:

```
- `transformers` version: 4.16.0.dev0
- Platform: Linux-5.15.8-76051508-generic-x86_64-with-glibc2.33
- Python version: 3.9.7
```

Hints from issue discussion:
Same error occures when using max samples with https://github.com/huggingface/transformers/blob/master/examples/pytorch/speech-recognition/run_speech_recognition_seq2seq.py
I'm seeing this too, when using preprocessing_num_workers with 
https://github.com/huggingface/transformers/blob/master/examples/pytorch/speech-recognition/run_speech_recognition_ctc.py

Commit message: whoops forgot one line

Patch to review:
diff --git a/.circleci/config.yml b/.circleci/config.yml
index 15735f0da93..f525d3f439f 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -11,6 +11,7 @@ jobs:
         resource_class: medium
         steps:
             - checkout
+            - run: sudo apt-get -y update && sudo apt-get install -y libsndfile1-dev
             - run: pip install --upgrade pip
             - run: python -m venv venv
             - run: source venv/bin/activate
@@ -26,6 +27,7 @@ jobs:
         resource_class: medium
         steps:
             - checkout
+            - run: sudo apt-get -y update && sudo apt-get install -y libsndfile1-dev
             - run: pip install --upgrade pip
             - run: python -m venv venv
             - run: source venv/bin/activate
diff --git a/datasets/wmt19/dummy/cs-en/1.0.0/dummy_data.zip b/datasets/wmt19/dummy/cs-en/1.0.0/dummy_data.zip
index eb4e883564b..d67563b4784 100644
Binary files a/datasets/wmt19/dummy/cs-en/1.0.0/dummy_data.zip and b/datasets/wmt19/dummy/cs-en/1.0.0/dummy_data.zip differ
diff --git a/setup.py b/setup.py
index 1c09beef8e6..d7a547ce84f 100644
--- a/setup.py
+++ b/setup.py
@@ -136,6 +136,7 @@
     "tensorflow>=2.3,!=2.6.0,!=2.6.1",
     "torch",
     "torchaudio",
+    "soundfile",
     "transformers",
     # datasets dependencies
     "bs4",
@@ -171,6 +172,7 @@
 ]
 
 TESTS_REQUIRE.extend(VISION_REQURE)
+TESTS_REQUIRE.extend(AUDIO_REQUIRE)
 
 if os.name != "nt":
     # dependencies of unbabel-comet
diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py
index ab5e8d39fe4..a5f0036b261 100644
--- a/src/datasets/arrow_dataset.py
+++ b/src/datasets/arrow_dataset.py
@@ -74,14 +74,7 @@
 from .info import DatasetInfo
 from .search import IndexableMixin
 from .splits import NamedSplit, Split, SplitInfo
-from .table import (
-    InMemoryTable,
-    MemoryMappedTable,
-    Table,
-    cast_with_sliced_list_support,
-    concat_tables,
-    list_table_cache_files,
-)
+from .table import InMemoryTable, MemoryMappedTable, Table, concat_tables, list_table_cache_files, table_cast
 from .tasks import TaskTemplate
 from .utils import logging
 from .utils.deprecation_utils import deprecated
@@ -636,8 +629,6 @@ def __init__(
 
         if self._data.schema.metadata is not None and "huggingface".encode("utf-8") in self._data.schema.metadata:
             metadata = json.loads(self._data.schema.metadata["huggingface".encode("utf-8")].decode())
-            if "info" in metadata and self.info.features is None:  # try to load features from the arrow file metadata
-                self._info.features = DatasetInfo.from_dict(metadata["info"]).features
             if (
                 "fingerprint" in metadata and self._fingerprint is None
             ):  # try to load fingerprint from the arrow file metadata
@@ -780,7 +771,7 @@ def from_pandas(
             info = DatasetInfo()
         info.features = features
         table = InMemoryTable.from_pandas(
-            df=df, preserve_index=preserve_index, schema=pa.schema(features.type) if features is not None else None
+            df=df, preserve_index=preserve_index, schema=features.arrow_schema if features is not None else None
         )
         return cls(table, info=info, split=split)
 
@@ -815,10 +806,12 @@ def from_dict(
         if features is not None:
             mapping = features.encode_batch(mapping)
         mapping = {
-            col: OptimizedTypedSequence(data, type=features.type[col].type if features is not None else None, col=col)
+            col: OptimizedTypedSequence(data, type=features[col] if features is not None else None, col=col)
             for col, data in mapping.items()
         }
         pa_table = InMemoryTable.from_pydict(mapping=mapping)
+        if info.features is None:
+            info.features = Features({col: ts.get_inferred_type() for col, ts in mapping.items()})
         return cls(pa_table, info=info, split=split)
 
     @staticmethod
@@ -1221,12 +1214,13 @@ def stringify_column(batch):
             dset = self
 
         # Create the new feature
-        class_names = sorted(sample for sample in dset.unique(column) if include_nulls or sample is not None)
+        class_names = sorted(str(sample) for sample in dset.unique(column) if include_nulls or sample is not None)
         dst_feat = ClassLabel(names=class_names)
 
         def cast_to_class_labels(batch):
             batch[column] = [
-                dst_feat.str2int(sample) if include_nulls or sample is not None else None for sample in batch[column]
+                dst_feat.str2int(str(sample)) if include_nulls or sample is not None else None
+                for sample in batch[column]
             ]
             return batch
 
@@ -1280,7 +1274,7 @@ def flatten_(self, max_depth=16):
                 self._data = self._data.flatten()
             else:
                 break
-        self.info.features = Features.from_arrow_schema(self._data.schema)
+        self.info.features = self.features.flatten(max_depth=max_depth)
         self._data = update_metadata_with_features(self._data, self.features)
         logger.info(f'Flattened dataset from depth {depth} to depth { 1 if depth + 1 < max_depth else "unknown"}.')
 
@@ -1299,7 +1293,7 @@ def flatten(self, new_fingerprint, max_depth=16) -> "Dataset":
                 dataset._data = dataset._data.flatten()
             else:
                 break
-        dataset.info.features = Features.from_arrow_schema(dataset._data.schema)
+        dataset.info.features = self.features.flatten(max_depth=max_depth)
         dataset._data = update_metadata_with_features(dataset._data, dataset.features)
         logger.info(f'Flattened dataset from depth {depth} to depth {1 if depth + 1 < max_depth else "unknown"}.')
         dataset._fingerprint = new_fingerprint
@@ -1348,10 +1342,8 @@ def cast_(
         type = features.type
         schema = pa.schema({col_name: type[col_name].type for col_name in self._data.column_names})
         dataset = self.with_format("arrow")
-        # capture the PyArrow version here to make the lambda serializable on Windows
-        is_pyarrow_at_least_4 = config.PYARROW_VERSION.major >= 4
         dataset = dataset.map(
-            lambda t: t.cast(schema) if is_pyarrow_at_least_4 else cast_with_sliced_list_support(t, schema),
+            partial(table_cast, schema=schema),
             batched=True,
             batch_size=batch_size,
             keep_in_memory=keep_in_memory,
@@ -1406,14 +1398,12 @@ def cast(
                 f"as the columns in the dataset: {self._data.column_names}"
             )
 
-        type = features.type
-        schema = pa.schema({col_name: type[col_name].type for col_name in self._data.column_names})
+        schema = features.arrow_schema
         format = self.format
         dataset = self.with_format("arrow")
         # capture the PyArrow version here to make the lambda serializable on Windows
-        is_pyarrow_at_least_4 = config.PYARROW_VERSION.major >= 4
         dataset = dataset.map(
-            lambda t: t.cast(schema) if is_pyarrow_at_least_4 else cast_with_sliced_list_support(t, schema),
+            partial(table_cast, schema=schema),
             batched=True,
             batch_size=batch_size,
             keep_in_memory=keep_in_memory,
@@ -1433,15 +1423,17 @@ def cast_column(self, column: str, feature: FeatureType, new_fingerprint: str) -
 
         Args:
             column (:obj:`str`): Column name.
-            feature (:class:`Feature`): Target feature.
+            feature (:class:`FeatureType`): Target feature.
 
         Returns:
             :class:`Dataset`
         """
-        if hasattr(feature, "decode_example"):
+        if hasattr(feature, "cast_storage"):
             dataset = copy.deepcopy(self)
             dataset.features[column] = feature
             dataset._fingerprint = new_fingerprint
+            dataset._data = dataset._data.cast(dataset.features.arrow_schema)
+            dataset._data = update_metadata_with_features(dataset._data, dataset.features)
             return dataset
         else:
             features = self.features.copy()
@@ -3860,8 +3852,8 @@ def add_item(self, item: dict, new_fingerprint: str):
         # Cast to align the schemas of the tables and concatenate the tables
         table = concat_tables(
             [
-                self._data.cast(pa.schema(dset_features.type)) if self.features != dset_features else self._data,
-                item_table.cast(pa.schema(item_features.type)),
+                self._data.cast(dset_features.arrow_schema) if self.features != dset_features else self._data,
+                item_table.cast(item_features.arrow_schema),
             ]
         )
         if self._indices is None:
diff --git a/src/datasets/arrow_writer.py b/src/datasets/arrow_writer.py
index 4bf1d4b4c54..dfd1896d253 100644
--- a/src/datasets/arrow_writer.py
+++ b/src/datasets/arrow_writer.py
@@ -19,23 +19,27 @@
 import socket
 import sys
 from dataclasses import asdict
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 import numpy as np
 import pyarrow as pa
 
+from datasets.features.features import FeatureType, Value
+
 from . import config, utils
 from .features import (
     Features,
-    ImageExtensionType,
+    Image,
     _ArrayXDExtensionType,
     cast_to_python_objects,
+    generate_from_arrow_type,
+    get_nested_type,
     list_of_np_array_to_pyarrow_listarray,
     numpy_to_pyarrow_listarray,
-    objects_to_list_of_image_dicts,
 )
 from .info import DatasetInfo
 from .keyhash import DuplicatedKeysError, KeyHasher
+from .table import array_cast, cast_array_to_feature, table_cast
 from .utils import logging
 from .utils.file_utils import hash_url_to_filename
 from .utils.py_utils import first_non_null_value
@@ -63,144 +67,193 @@ class TypedSequence:
 
     Example::
 
-        from datasets.features import Array2DExtensionType
+        from datasets.features import Array2D, Array2DExtensionType, Value
         from datasets.arrow_writer import TypedSequence
         import pyarrow as pa
 
-        arr = pa.array(TypedSequence([1, 2, 3], type=pa.int32()))
+        arr = pa.array(TypedSequence([1, 2, 3], type=Value("int32")))
         assert arr.type == pa.int32()
 
-        arr = pa.array(TypedSequence([1, 2, 3], try_type=pa.int32()))
+        arr = pa.array(TypedSequence([1, 2, 3], try_type=Value("int32")))
         assert arr.type == pa.int32()
 
-        arr = pa.array(TypedSequence(["foo", "bar"], try_type=pa.int32()))
+        arr = pa.array(TypedSequence(["foo", "bar"], try_type=Value("int32")))
         assert arr.type == pa.string()
 
-        arr = pa.array(TypedSequence([[[1, 2, 3]]], type=Array2DExtensionType((1, 3), "int64")))
+        arr = pa.array(TypedSequence([[[1, 2, 3]]], type=Array2D((1, 3), "int64")))
         assert arr.type == Array2DExtensionType((1, 3), "int64")
 
         table = pa.Table.from_pydict({
-            "image": TypedSequence([[[1, 2, 3]]], type=Array2DExtensionType((1, 3), "int64"))
+            "image": TypedSequence([[[1, 2, 3]]], type=Array2D((1, 3), "int64"))
         })
         assert table["image"].type == Array2DExtensionType((1, 3), "int64")
 
     """
 
-    def __init__(self, data, type=None, try_type=None, optimized_int_type=None):
+    def __init__(
+        self,
+        data: Iterable,
+        type: Optional[FeatureType] = None,
+        try_type: Optional[FeatureType] = None,
+        optimized_int_type: Optional[FeatureType] = None,
+    ):
         # assert type is None or try_type is None,
         if type is not None and try_type is not None:
             raise ValueError("You cannot specify both type and try_type")
+        # set attributes
         self.data = data
         self.type = type
         self.try_type = try_type  # is ignored if it doesn't match the data
         self.optimized_int_type = optimized_int_type
+        # when trying a type (is ignored if data is not compatible)
+        self.trying_type = self.try_type is not None
+        self.trying_int_optimization = optimized_int_type is not None and type is None and try_type is None
+        # used to get back the inferred type after __arrow_array__() is called once
+        self._inferred_type = None
+
+    def get_inferred_type(self) -> FeatureType:
+        """Return the inferred feature type.
+        This is done by converting the sequence to an Arrow array, and getting the corresponding
+        feature type.
+
+        Since building the Arrow array can be expensive, the value of the inferred type is cached
+        as soon as pa.array is called on the typed sequence.
+
+        Returns:
+            FeatureType: inferred feature type of the sequence.
+        """
+        if self._inferred_type is None:
+            self._inferred_type = generate_from_arrow_type(pa.array(self).type)
+        return self._inferred_type
 
-    def __arrow_array__(self, type=None):
-        """This function is called when calling pa.array(typed_sequence)"""
+    @staticmethod
+    def _infer_custom_type_and_encode(data: Iterable) -> Tuple[Iterable, Optional[FeatureType]]:
+        """Implement type inference for custom objects like PIL.Image.Image -> Image type.
 
+        This function is only used for custom python objects that can't be direclty passed to build
+        an Arrow array. In such cases is infers the feature type to use, and it encodes the data so
+        that they can be passed to an Arrow array.
+
+        Args:
+            data (Iterable): array of data to infer the type, e.g. a list of PIL images.
+
+        Returns:
+            Tuple[Iterable, Optional[FeatureType]]: a tuple with:
+                - the (possibly encoded) array, if the inferred feature type requires encoding
+                - the inferred feature type if the array is made of supported custom objects like
+                    PIL images, else None.
+        """
         if config.PIL_AVAILABLE and "PIL" in sys.modules:
             import PIL.Image
 
+            non_null_idx, non_null_value = first_non_null_value(data)
+            if isinstance(non_null_value, PIL.Image.Image):
+                return [Image().encode_example(value) for value in data], Image()
+        return data, None
+
+    def __arrow_array__(self, type: Optional[pa.DataType] = None):
+        """This function is called when calling pa.array(typed_sequence)"""
+
         if type is not None:
             raise ValueError("TypedSequence is supposed to be used with pa.array(typed_sequence, type=None)")
-        trying_type = False
-        if type is not None:  # user explicitly passed the feature
-            pass
-        elif type is None and self.try_type:
-            type = self.try_type
-            trying_type = True
+        del type  # make sure we don't use it
+        data = self.data
+        # automatic type inference for custom objects
+        if self.type is None and self.try_type is None:
+            data, self._inferred_type = self._infer_custom_type_and_encode(data)
+        if self._inferred_type is None:
+            type = self.try_type if self.trying_type else self.type
         else:
-            type = self.type
-        trying_int_optimization = False
-        non_null_idx, non_null_value = first_non_null_value(self.data)
-        if type is None:  # automatic type inference for custom objects
-            if config.PIL_AVAILABLE and "PIL" in sys.modules and isinstance(non_null_value, PIL.Image.Image):
-                type = ImageExtensionType()
+            type = self._inferred_type
+        pa_type = get_nested_type(type) if type is not None else None
+        optimized_int_pa_type = (
+            get_nested_type(self.optimized_int_type) if self.optimized_int_type is not None else None
+        )
         try:
-            if isinstance(type, _ArrayXDExtensionType):
-                if isinstance(self.data, np.ndarray):
-                    storage = numpy_to_pyarrow_listarray(self.data, type=type.value_type)
-                elif isinstance(self.data, list) and self.data and isinstance(non_null_value, np.ndarray):
-                    storage = list_of_np_array_to_pyarrow_listarray(self.data, type=type.value_type)
+            # custom pyarrow types
+            if isinstance(pa_type, _ArrayXDExtensionType):
+                if isinstance(data, np.ndarray):
+                    storage = numpy_to_pyarrow_listarray(data, type=pa_type.value_type)
+                elif isinstance(data, list) and data and isinstance(first_non_null_value(data)[1], np.ndarray):
+                    storage = list_of_np_array_to_pyarrow_listarray(data, type=pa_type.value_type)
                 else:
-                    storage = pa.array(self.data, type.storage_dtype)
-                out = pa.ExtensionArray.from_storage(type, storage)
-            elif isinstance(type, ImageExtensionType):
-                storage = pa.array(objects_to_list_of_image_dicts(self.data), type=type.storage_type)
-                out = pa.ExtensionArray.from_storage(type, storage)
-            elif isinstance(self.data, np.ndarray):
-                out = numpy_to_pyarrow_listarray(self.data)
-                if type is not None:
-                    out = out.cast(type)
-            elif isinstance(self.data, list) and self.data and isinstance(non_null_value, np.ndarray):
-                out = list_of_np_array_to_pyarrow_listarray(self.data)
-                if type is not None:
-                    out = out.cast(type)
+                    storage = pa.array(data, pa_type.storage_dtype)
+                return pa.ExtensionArray.from_storage(pa_type, storage)
+
+            # efficient np array to pyarrow array
+            if isinstance(data, np.ndarray):
+                out = numpy_to_pyarrow_listarray(data)
+            elif isinstance(data, list) and data and isinstance(first_non_null_value(data)[1], np.ndarray):
+                out = list_of_np_array_to_pyarrow_listarray(data)
             else:
-                out = pa.array(cast_to_python_objects(self.data, only_1d_for_numpy=True), type=type)
-            if trying_type and not isinstance(type, ImageExtensionType) and non_null_idx != -1:
-                is_equal = (
-                    np.array_equal(np.array(out[non_null_idx].as_py()), self.data[non_null_idx])
-                    if isinstance(self.data[non_null_idx], np.ndarray)
-                    else out[non_null_idx].as_py() == self.data[non_null_idx]
-                )
-                if not is_equal:
-                    raise TypeError(
-                        "Specified try_type alters data. Please check that the type/feature that you provided match the type/features of the data."
-                    )
-            if self.optimized_int_type and self.type is None and self.try_type is None:
-                trying_int_optimization = True
+                out = pa.array(cast_to_python_objects(data, only_1d_for_numpy=True))
+
+            # use smaller integer precisions if possible
+            if self.trying_int_optimization:
                 if pa.types.is_int64(out.type):
-                    out = out.cast(self.optimized_int_type)
+                    out = out.cast(optimized_int_pa_type)
                 elif pa.types.is_list(out.type):
                     if pa.types.is_int64(out.type.value_type):
-                        out = out.cast(pa.list_(self.optimized_int_type))
+                        out = array_cast(out, pa.list_(optimized_int_pa_type))
                     elif pa.types.is_list(out.type.value_type) and pa.types.is_int64(out.type.value_type.value_type):
-                        out = out.cast(pa.list_(pa.list_(self.optimized_int_type)))
+                        out = array_cast(out, pa.list_(pa.list_(optimized_int_pa_type)))
+            # otherwise we can finally use the user's type
+            elif type is not None:
+                # We use array_cast_to_feature to support casting to custom types like Audio and Image
+                # Also, when trying type "string", we don't want to convert integers or floats to "string".
+                # We only do it if trying_type is False - since this is what the user asks for.
+                out = cast_array_to_feature(out, type, allow_number_to_str=not self.trying_type)
             return out
         except (TypeError, pa.lib.ArrowInvalid) as e:  # handle type errors and overflows
-            if trying_type:
+            if self.trying_type:
                 try:  # second chance
-                    if isinstance(self.data, np.ndarray):
-                        return numpy_to_pyarrow_listarray(self.data)
-                    elif (
-                        isinstance(self.data, list)
-                        and self.data
-                        and any(isinstance(value, np.ndarray) for value in self.data)
-                    ):
-                        return list_of_np_array_to_pyarrow_listarray(self.data)
+                    if isinstance(data, np.ndarray):
+                        return numpy_to_pyarrow_listarray(data)
+                    elif isinstance(data, list) and data and any(isinstance(value, np.ndarray) for value in data):
+                        return list_of_np_array_to_pyarrow_listarray(data)
                     else:
-                        return pa.array(cast_to_python_objects(self.data, only_1d_for_numpy=True))
+                        return pa.array(cast_to_python_objects(data, only_1d_for_numpy=True))
                 except pa.lib.ArrowInvalid as e:
                     if "overflow" in str(e):
                         raise OverflowError(
-                            f"There was an overflow with type {type_(self.data)}. Try to reduce writer_batch_size to have batches smaller than 2GB.\n({e})"
+                            f"There was an overflow with type {type_(data)}. Try to reduce writer_batch_size to have batches smaller than 2GB.\n({e})"
                         ) from None
-                    elif trying_int_optimization and "not in range" in str(e):
-                        optimized_int_type_str = np.dtype(self.optimized_int_type.to_pandas_dtype()).name
-                        logger.info(f"Failed to cast a sequence to {optimized_int_type_str}. Falling back to int64.")
+                    elif self.trying_int_optimization and "not in range" in str(e):
+                        optimized_int_pa_type_str = np.dtype(optimized_int_pa_type.to_pandas_dtype()).name
+                        logger.info(
+                            f"Failed to cast a sequence to {optimized_int_pa_type_str}. Falling back to int64."
+                        )
                         return out
                     else:
                         raise
             elif "overflow" in str(e):
                 raise OverflowError(
-                    f"There was an overflow with type {type_(self.data)}. Try to reduce writer_batch_size to have batches smaller than 2GB.\n({e})"
+                    f"There was an overflow with type {type_(data)}. Try to reduce writer_batch_size to have batches smaller than 2GB.\n({e})"
                 ) from None
-            elif trying_int_optimization and "not in range" in str(e):
-                optimized_int_type_str = np.dtype(self.optimized_int_type.to_pandas_dtype()).name
-                logger.info(f"Failed to cast a sequence to {optimized_int_type_str}. Falling back to int64.")
+            elif self.trying_int_optimization and "not in range" in str(e):
+                optimized_int_pa_type_str = np.dtype(optimized_int_pa_type.to_pandas_dtype()).name
+                logger.info(f"Failed to cast a sequence to {optimized_int_pa_type_str}. Falling back to int64.")
                 return out
             else:
                 raise
 
 
 class OptimizedTypedSequence(TypedSequence):
-    def __init__(self, data, type=None, try_type=None, col=None, optimized_int_type=None):
+    def __init__(
+        self,
+        data,
+        type: Optional[FeatureType] = None,
+        try_type: Optional[FeatureType] = None,
+        col: Optional[str] = None,
+        optimized_int_type: Optional[FeatureType] = None,
+    ):
         optimized_int_type_by_col = {
-            "attention_mask": pa.int8(),  # binary tensor
-            "special_tokens_mask": pa.int8(),
-            "input_ids": pa.int32(),  # typical vocab size: 0-50k (max ~500k, never > 1M)
-            "token_type_ids": pa.int8(),  # binary mask; some (XLNetModel) use an additional token represented by a 2
+            "attention_mask": Value("int8"),  # binary tensor
+            "special_tokens_mask": Value("int8"),
+            "input_ids": Value("int32"),  # typical vocab size: 0-50k (max ~500k, never > 1M)
+            "token_type_ids": Value(
+                "int8"
+            ),  # binary mask; some (XLNetModel) use an additional token represented by a 2
         }
         if type is None and try_type is None:
             optimized_int_type = optimized_int_type_by_col.get(col, None)
@@ -337,38 +390,18 @@ def write_examples_on_file(self):
         if not self.current_examples:
             return
 
-        # Since current_examples contains (example, key) tuples
+        # order the columns properly
         cols = (
             [col for col in self.schema.names if col in self.current_examples[0][0]]
             + [col for col in self.current_examples[0][0].keys() if col not in self.schema.names]
             if self.schema
             else self.current_examples[0][0].keys()
         )
-
-        schema = None if self.pa_writer is None and self.update_features else self.schema
-        try_schema = self.schema if self.pa_writer is None and self.update_features else None
-        arrays = []
-        inferred_types = []
+        batch_examples = {}
         for col in cols:
-            col_type = schema.field(col).type if schema else None
-            col_try_type = try_schema.field(col).type if try_schema is not None and col in try_schema.names else None
-            typed_sequence = OptimizedTypedSequence(
-                [row[0][col] for row in self.current_examples], type=col_type, try_type=col_try_type, col=col
-            )
-            pa_array = pa.array(typed_sequence)
-            inferred_type = pa_array.type
-            first_example = pa.array(OptimizedTypedSequence(typed_sequence.data[:1], type=inferred_type))[0]
-            if pa_array[0] != first_example:  # Sanity check (check for overflow in StructArray or ListArray)
-                # This check fails with FloatArrays with nans, which is not what we want, so account for that:
-                if not isinstance(pa_array[0], pa.lib.FloatScalar):
-                    raise OverflowError(
-                        f"There was an overflow in the {type(pa_array)}. Try to reduce writer_batch_size to have batches smaller than 2GB"
-                    )
-            arrays.append(pa_array)
-            inferred_types.append(inferred_type)
-        schema = pa.schema(zip(cols, inferred_types)) if self.pa_writer is None else self.schema
-        table = pa.Table.from_arrays(arrays, schema=schema)
-        self.write_table(table)
+            # Since current_examples contains (example, key) tuples
+            batch_examples[col] = [row[0][col] for row in self.current_examples]
+        self.write_batch(batch_examples=batch_examples)
         self.current_examples = []
 
     def write_rows_on_file(self):
@@ -447,15 +480,24 @@ def write_batch(
         """
         if batch_examples and len(next(iter(batch_examples.values()))) == 0:
             return
-        schema = None if self.pa_writer is None and self.update_features else self.schema
-        try_schema = self.schema if self.pa_writer is None and self.update_features else None
-        typed_sequence_examples = {}
-        for col in sorted(batch_examples.keys()):
-            col_type = schema.field(col).type if schema else None
-            col_try_type = try_schema.field(col).type if try_schema is not None and col in try_schema.names else None
+        features = None if self.pa_writer is None and self.update_features else self._features
+        try_features = self._features if self.pa_writer is None and self.update_features else None
+        arrays = []
+        inferred_features = Features()
+        cols = (
+            [col for col in self.schema.names if col in batch_examples]
+            + [col for col in batch_examples.keys() if col not in self.schema.names]
+            if self.schema
+            else batch_examples.keys()
+        )
+        for col in cols:
+            col_type = features[col] if features else None
+            col_try_type = try_features[col] if try_features is not None and col in try_features else None
             typed_sequence = OptimizedTypedSequence(batch_examples[col], type=col_type, try_type=col_try_type, col=col)
-            typed_sequence_examples[col] = typed_sequence
-        pa_table = pa.Table.from_pydict(typed_sequence_examples)
+            arrays.append(pa.array(typed_sequence))
+            inferred_features[col] = typed_sequence.get_inferred_type()
+        schema = inferred_features.arrow_schema if self.pa_writer is None else self.schema
+        pa_table = pa.Table.from_arrays(arrays, schema=schema)
         self.write_table(pa_table, writer_batch_size)
 
     def write_table(self, pa_table: pa.Table, writer_batch_size: Optional[int] = None):
@@ -468,9 +510,7 @@ def write_table(self, pa_table: pa.Table, writer_batch_size: Optional[int] = Non
             writer_batch_size = self.writer_batch_size
         if self.pa_writer is None:
             self._build_writer(inferred_schema=pa_table.schema)
-        # reorder the arrays if necessary + cast to self._schema
-        # we can't simply use .cast here because we may need to change the order of the columns
-        pa_table = pa.Table.from_arrays([pa_table[name] for name in self._schema.names], schema=self._schema)
+        pa_table = table_cast(pa_table, self._schema)
         batches: List[pa.RecordBatch] = pa_table.to_batches(max_chunksize=writer_batch_size)
         self._num_bytes += sum(batch.nbytes for batch in batches)
         self._num_examples += pa_table.num_rows
@@ -520,7 +560,7 @@ def __init__(
 
         if features is not None:
             self._features: Features = features
-            self._schema: pa.Schema = pa.schema(features.type)
+            self._schema: pa.Schema = features.arrow_schema
         else:
             self._schema: pa.Schema = schema
             self._features: Features = Features.from_arrow_schema(schema)
diff --git a/src/datasets/features/__init__.py b/src/datasets/features/__init__.py
index febd96e3369..23dfee60170 100644
--- a/src/datasets/features/__init__.py
+++ b/src/datasets/features/__init__.py
@@ -8,5 +8,5 @@
     _cast_to_python_objects,
     _is_zero_copy_only,
 )
-from .image import Image, ImageExtensionType, objects_to_list_of_image_dicts
+from .image import Image, objects_to_list_of_image_dicts
 from .translation import Translation, TranslationVariableLanguages
diff --git a/src/datasets/features/audio.py b/src/datasets/features/audio.py
index 43410b5b4b5..2d1deb0d7e2 100644
--- a/src/datasets/features/audio.py
+++ b/src/datasets/features/audio.py
@@ -1,9 +1,10 @@
 from dataclasses import dataclass, field
 from io import BytesIO
-from typing import Any, ClassVar, Optional
+from typing import Any, ClassVar, Optional, Union
 
 import pyarrow as pa
 
+from ..table import array_cast
 from ..utils.streaming_download_manager import xopen
 
 
@@ -20,6 +21,14 @@ class Audio:
 
       This is useful for archived files with sequential access.
 
+    - A :obj:`dict` with the keys:
+
+        - path: String with relative path of the audio file to the archive file.
+        - array: Array containing the audio sample
+        - sampling_rate: Integer corresponding to the samping rate of the audio sample.
+
+      This is useful for archived files with sequential access.
+
     Args:
         sampling_rate (:obj:`int`, optional): Target sampling rate. If `None`, the native sampling rate is used.
         mono (:obj:`bool`, default ``True``): Whether to convert the audio signal to mono by averaging samples across
@@ -28,17 +37,14 @@ class Audio:
 
     sampling_rate: Optional[int] = None
     mono: bool = True
-    _storage_dtype: str = "string"
     id: Optional[str] = None
     # Automatically constructed
     dtype: ClassVar[str] = "dict"
-    pa_type: ClassVar[Any] = None
+    pa_type: ClassVar[Any] = pa.struct({"bytes": pa.binary(), "path": pa.string()})
     _type: str = field(default="Audio", init=False, repr=False)
 
     def __call__(self):
-        return (
-            pa.struct({"path": pa.string(), "bytes": pa.binary()}) if self._storage_dtype == "struct" else pa.string()
-        )
+        return self.pa_type
 
     def encode_example(self, value):
         """Encode example into a format for Arrow.
@@ -47,18 +53,30 @@ def encode_example(self, value):
             value (:obj:`str` or :obj:`dict`): Data passed as input to Audio feature.
 
         Returns:
-            :obj:`str` or :obj:`dict`
+            :obj:`dict`
         """
-        if isinstance(value, dict):
-            self._storage_dtype = "struct"
-        return value
+        try:
+            import soundfile as sf  # soundfile is a dependency of librosa, needed to decode audio files.
+        except ImportError as err:
+            raise ImportError("To support encoding audio data, please install 'soundfile'.") from err
+        if isinstance(value, str):
+            return {"bytes": None, "path": value}
+        elif isinstance(value, dict) and "array" in value:
+            buffer = BytesIO()
+            sf.write(buffer, value["array"], value["sampling_rate"])
+            return {"bytes": buffer.getvalue(), "path": value.get("path")}
+        elif value.get("bytes") is not None or value.get("path") is not None:
+            return {"bytes": value.get("bytes"), "path": value.get("path")}
+        else:
+            raise ValueError(
+                f"An audio sample should have one of 'path' or 'bytes' but they are missing or None in {value}."
+            )
 
     def decode_example(self, value):
         """Decode example audio file into audio data.
 
         Args:
-            value (obj:`str` or :obj:`dict`): Either a string with the absolute audio file path or a dictionary with
-                keys:
+            value (:obj:`dict`): a dictionary with keys:
 
                 - path: String with relative audio file path.
                 - bytes: Bytes of the audio file.
@@ -66,8 +84,10 @@ def decode_example(self, value):
         Returns:
             dict
         """
-        path, file = (value["path"], BytesIO(value["bytes"])) if isinstance(value, dict) else (value, None)
-        if path.endswith("mp3"):
+        path, file = (value["path"], BytesIO(value["bytes"])) if value["bytes"] is not None else (value["path"], None)
+        if path is None and file is None:
+            raise ValueError(f"An audio sample should have one of 'path' or 'bytes' but both are None in {value}.")
+        elif path is not None and path.endswith("mp3"):
             array, sampling_rate = self._decode_mp3(file if file else path)
         else:
             if file:
@@ -76,6 +96,39 @@ def decode_example(self, value):
                 array, sampling_rate = self._decode_non_mp3_path_like(path)
         return {"path": path, "array": array, "sampling_rate": sampling_rate}
 
+    def cast_storage(self, storage: Union[pa.StringArray, pa.StructArray]) -> pa.StructArray:
+        """Cast an Arrow array to the Audio arrow storage type.
+        The Arrow types that can be converted to the Audio pyarrow storage type are:
+
+        - pa.string() - it must contain the "path" data
+        - pa.struct({"bytes": pa.binary()})
+        - pa.struct({"path": pa.string()})
+        - pa.struct({"bytes": pa.binary(), "path": pa.string()})  - order doesn't matter
+
+        Args:
+            storage (Union[pa.StringArray, pa.StructArray]): [description]
+
+        Returns:
+            pa.StructArray: Array in the Audio arrow storage type, that is
+                pa.struct({"bytes": pa.binary(), "path": pa.string()})
+        """
+        if pa.types.is_string(storage.type):
+            bytes_array = pa.array([None] * len(storage), type=pa.binary())
+            storage = pa.StructArray.from_arrays([bytes_array, storage], ["bytes", "path"])
+        elif pa.types.is_struct(storage.type) and storage.type.get_all_field_indices("array"):
+            storage = pa.array([Audio().encode_example(x) for x in storage.to_pylist()])
+        elif pa.types.is_struct(storage.type):
+            if storage.type.get_field_index("bytes") >= 0:
+                bytes_array = storage.field("bytes")
+            else:
+                bytes_array = pa.array([None] * len(storage), type=pa.binary())
+            if storage.type.get_field_index("path") >= 0:
+                path_array = storage.field("path")
+            else:
+                path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"])
+        return array_cast(storage, self.pa_type)
+
     def _decode_non_mp3_path_like(self, path):
         try:
             import librosa
@@ -91,7 +144,7 @@ def _decode_non_mp3_file_like(self, file):
             import librosa
             import soundfile as sf
         except ImportError as err:
-            raise ImportError("To support decoding audio files, please install 'librosa'.") from err
+            raise ImportError("To support decoding audio files, please install 'librosa' and 'soundfile'.") from err
 
         array, sampling_rate = sf.read(file)
         array = array.T
diff --git a/src/datasets/features/features.py b/src/datasets/features/features.py
index 2d1fa9486df..674629635ef 100644
--- a/src/datasets/features/features.py
+++ b/src/datasets/features/features.py
@@ -16,10 +16,11 @@
 # Lint as: python3
 """ This class handle features definition in datasets and some utilities to display table type."""
 import copy
+import json
 import re
 import sys
 from collections.abc import Iterable
-from dataclasses import dataclass, field, fields
+from dataclasses import _asdict_inner, dataclass, field, fields
 from functools import reduce
 from operator import mul
 from typing import Any, ClassVar, Dict, List, Optional
@@ -36,7 +37,7 @@
 
 from datasets import config, utils
 from datasets.features.audio import Audio
-from datasets.features.image import Image, ImageExtensionType, PandasImageExtensionDtype
+from datasets.features.image import Image, encode_pil_image
 from datasets.features.translation import Translation, TranslationVariableLanguages
 from datasets.utils.logging import get_logger
 
@@ -201,10 +202,7 @@ def _cast_to_python_objects(obj: Any, only_1d_for_numpy: bool) -> Tuple[Any, boo
         else:
             return [_cast_to_python_objects(x, only_1d_for_numpy=only_1d_for_numpy)[0] for x in np.asarray(obj)], True
     elif config.PIL_AVAILABLE and "PIL" in sys.modules and isinstance(obj, PIL.Image.Image):
-        if not only_1d_for_numpy:
-            return obj, False
-        else:
-            return [_cast_to_python_objects(x, only_1d_for_numpy=only_1d_for_numpy)[0] for x in np.array(obj)], True
+        return encode_pil_image(obj), True
     elif isinstance(obj, pd.Series):
         return obj.values.tolist(), True
     elif isinstance(obj, pd.DataFrame):
@@ -626,8 +624,6 @@ def __eq__(self, other) -> np.ndarray:
 def pandas_types_mapper(dtype):
     if isinstance(dtype, _ArrayXDExtensionType):
         return PandasArrayExtensionDtype(dtype.value_type)
-    elif isinstance(dtype, ImageExtensionType):
-        return PandasImageExtensionDtype()
 
 
 @dataclass
@@ -839,7 +835,7 @@ def get_nested_type(schema: FeatureType) -> pa.DataType:
     elif isinstance(schema, Sequence):
         value_type = get_nested_type(schema.feature)
         # We allow to reverse list of dict => dict of list for compatibility with tfds
-        if isinstance(value_type, pa.StructType):
+        if isinstance(schema.feature, dict):
             return pa.struct({f.name: pa.list_(f.type, schema.length) for f in value_type})
         return pa.list_(value_type, schema.length)
 
@@ -908,6 +904,42 @@ def encode_nested_example(schema, obj):
     return obj
 
 
+def decode_nested_example(schema, obj):
+    """Decode a nested example.
+    This is used since some features (in particular Audio and Image) have some logic during decoding.
+
+    To avoid iterating over possibly long lists, it first checks (recursively) if the first element that is not None or empty (if it is a sequence) has to be decoded.
+    If the first element needs to be decoded, then all the elements of the list will be decoded, otherwise they'll stay the same.
+    """
+    # Nested structures: we allow dict, list/tuples, sequences
+    if isinstance(schema, dict):
+        return {
+            k: decode_nested_example(sub_schema, sub_obj) for k, (sub_schema, sub_obj) in utils.zip_dict(schema, obj)
+        }
+    elif isinstance(schema, (list, tuple)):
+        sub_schema = schema[0]
+        if obj is None:
+            return None
+        else:
+            if len(obj) > 0:
+                for first_elmt in obj:
+                    if _check_non_null_non_empty_recursive(first_elmt, sub_schema):
+                        break
+                if decode_nested_example(sub_schema, first_elmt) != first_elmt:
+                    return [decode_nested_example(sub_schema, o) for o in obj]
+            return list(obj)
+    elif isinstance(schema, Sequence):
+        # We allow to reverse list of dict => dict of list for compatiblity with tfds
+        if isinstance(schema.feature, dict):
+            return {k: decode_nested_example([schema.feature[k]], obj[k]) for k in schema.feature}
+        else:
+            return decode_nested_example([schema.feature], obj)
+    # Object with special decoding:
+    elif isinstance(schema, (Audio, Image)):
+        return schema.decode_example(obj) if obj is not None else None
+    return obj
+
+
 def generate_from_dict(obj: Any):
     """Regenerate the nested feature object from a deserialized dict.
     We use the '_type' fields to get the dataclass name to load.
@@ -956,8 +988,6 @@ def generate_from_arrow_type(pa_type: pa.DataType) -> FeatureType:
     elif isinstance(pa_type, _ArrayXDExtensionType):
         array_feature = [None, None, Array2D, Array3D, Array4D, Array5D][pa_type.ndims]
         return array_feature(shape=pa_type.shape, dtype=pa_type.value_type)
-    elif isinstance(pa_type, ImageExtensionType):
-        return Image()
     elif isinstance(pa_type, pa.DictionaryType):
         raise NotImplementedError  # TODO(thom) this will need access to the dictionary as well (for labels). I.e. to the py_table
     elif isinstance(pa_type, pa.DataType):
@@ -1000,6 +1030,17 @@ def list_of_np_array_to_pyarrow_listarray(l_arr: List[np.ndarray], type: pa.Data
         return pa.array([], type=type)
 
 
+def require_decoding(feature: FeatureType) -> bool:
+    if isinstance(feature, dict):
+        return any(require_decoding(f) for f in feature.values())
+    elif isinstance(feature, (list, tuple)):
+        return require_decoding(feature[0])
+    elif isinstance(feature, Sequence):
+        return require_decoding(feature.feature)
+    else:
+        return hasattr(feature, "decode_example")
+
+
 class Features(dict):
     """A special dictionary that defines the internal structure of a dataset.
 
@@ -1031,6 +1072,19 @@ class Features(dict):
         - :class:`datasets.Translation` and :class:`datasets.TranslationVariableLanguages`, the two features specific to Machine Translation
     """
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._column_requires_decoding: Dict[str, bool] = {
+            col: require_decoding(feature) for col, feature in self.items()
+        }
+
+    def __setitem__(self, column_name: str, feature: FeatureType):
+        super().__setitem__(column_name, feature)
+        self._column_requires_decoding[column_name] = require_decoding(feature)
+
+    def __reduce__(self):
+        return Features, (dict(self),)
+
     @property
     def type(self):
         """
@@ -1041,10 +1095,22 @@ def type(self):
         """
         return get_nested_type(self)
 
+    @property
+    def arrow_schema(self):
+        """
+        Features schema.
+
+        Returns:
+            :obj:`pyarrow.Schema`
+        """
+        hf_metadata = {"info": {"features": self.to_dict()}}
+        return pa.schema(self.type).with_metadata({"huggingface": json.dumps(hf_metadata)})
+
     @classmethod
     def from_arrow_schema(cls, pa_schema: pa.Schema) -> "Features":
         """
         Construct Features from Arrow Schema.
+        It also checks the schema metadata for Hugging Face Datasets features.
 
         Args:
             pa_schema (:obj:`pyarrow.Schema`): Arrow Schema.
@@ -1052,6 +1118,11 @@ def from_arrow_schema(cls, pa_schema: pa.Schema) -> "Features":
         Returns:
             :class:`Features`
         """
+        # try to load features from the arrow schema metadata
+        if pa_schema.metadata is not None and "huggingface".encode("utf-8") in pa_schema.metadata:
+            metadata = json.loads(pa_schema.metadata["huggingface".encode("utf-8")].decode())
+            if "info" in metadata and "features" in metadata["info"] and metadata["info"]["features"] is not None:
+                return Features.from_dict(metadata["info"]["features"])
         obj = {field.name: generate_from_arrow_type(field.type) for field in pa_schema}
         return cls(**obj)
 
@@ -1083,6 +1154,9 @@ def from_dict(cls, dic) -> "Features":
         obj = generate_from_dict(dic)
         return cls(**obj)
 
+    def to_dict(self):
+        return _asdict_inner(self, dict)
+
     def encode_example(self, example):
         """
         Encode example into a format for Arrow.
@@ -1123,11 +1197,12 @@ def decode_example(self, example: dict):
         Returns:
             :obj:`dict[str, Any]`
         """
+
         return {
-            column: feature.decode_example(value)
-            if hasattr(feature, "decode_example") and value is not None
+            column_name: decode_nested_example(feature, value)
+            if self._column_requires_decoding[column_name]
             else value
-            for column, (feature, value) in utils.zip_dict(
+            for column_name, (feature, value) in utils.zip_dict(
                 {key: value for key, value in self.items() if key in example}, example
             )
         }
@@ -1144,7 +1219,7 @@ def decode_column(self, column: list, column_name: str):
         """
         return (
             [self[column_name].decode_example(value) if value is not None else None for value in column]
-            if hasattr(self[column_name], "decode_example")
+            if self._column_requires_decoding[column_name]
             else column
         )
 
@@ -1161,7 +1236,7 @@ def decode_batch(self, batch: dict):
         for column_name, column in batch.items():
             decoded_batch[column_name] = (
                 [self[column_name].decode_example(value) if value is not None else None for value in column]
-                if hasattr(self[column_name], "decode_example")
+                if self._column_requires_decoding[column_name]
                 else column
             )
         return decoded_batch
@@ -1235,3 +1310,31 @@ def recursive_reorder(source, target, stack=""):
                 return source
 
         return Features(recursive_reorder(self, other))
+
+    def flatten(self, max_depth=16) -> "Features":
+        """Flatten the features. Every dictionary column is removed and is replaced by
+        all the subfields it contains. The new fields are named by concatenating the
+        name of the original column and the subfield name like this: "<original>.<subfield>".
+
+        If a column contains nested dictionaries, then all the lower-level subfields names are
+        also concatenated to form new columns: "<original>.<subfield>.<subsubfield>", etc.
+
+        Returns:
+            Features: the flattened features
+        """
+        for depth in range(1, max_depth):
+            no_change = True
+            flattened = self.copy()
+            for column_name, subfeature in self.items():
+                if isinstance(subfeature, dict):
+                    no_change = False
+                    flattened.update({f"{column_name}.{k}": v for k, v in subfeature.items()})
+                    del flattened[column_name]
+                elif isinstance(subfeature, Sequence) and isinstance(subfeature.feature, dict):
+                    no_change = False
+                    flattened.update({f"{column_name}.{k}": Sequence(v) for k, v in subfeature.feature.items()})
+                    del flattened[column_name]
+            self = flattened
+            if no_change:
+                break
+        return self
diff --git a/src/datasets/features/image.py b/src/datasets/features/image.py
index abdf50d9c80..2e6df779a24 100644
--- a/src/datasets/features/image.py
+++ b/src/datasets/features/image.py
@@ -1,14 +1,12 @@
 from dataclasses import dataclass, field
 from io import BytesIO
-from typing import TYPE_CHECKING, Any, ClassVar, List, Optional, Sequence, Union
+from typing import TYPE_CHECKING, Any, ClassVar, List, Optional, Union
 
 import numpy as np
-import pandas as pd
 import pyarrow as pa
-from pandas.api.extensions import ExtensionArray as PandasExtensionArray
-from pandas.api.extensions import ExtensionDtype as PandasExtensionDtype
 
 from .. import config
+from ..table import array_cast
 from ..utils.file_utils import is_local_path
 from ..utils.py_utils import first_non_null_value, no_op_if_value_is_null
 from ..utils.streaming_download_manager import xopen
@@ -21,140 +19,6 @@
 _IMAGE_COMPRESSION_FORMATS: Optional[List[str]] = None
 
 
-class ImageExtensionType(pa.PyExtensionType):
-    def __init__(self):
-        pa.PyExtensionType.__init__(self, pa.struct({"bytes": pa.binary(), "path": pa.string()}))
-
-    def __arrow_ext_class__(self):
-        return ImageExtensionArray
-
-    def __reduce__(self):
-        return self.__class__, ()
-
-    def to_pandas_dtype(self):
-        return PandasImageExtensionDtype()
-
-
-class ImageExtensionArray(pa.ExtensionArray):
-    def __array__(self):
-        return self.to_numpy(zero_copy_only=False)
-
-    def __getitem__(self, i):
-        return self.storage[i]
-
-    def to_pylist(self):
-        return self.to_numpy(zero_copy_only=False).tolist()
-
-
-class PandasImageExtensionDtype(PandasExtensionDtype):
-    def __from_arrow__(self, array: Union[pa.Array, pa.ChunkedArray]):
-        if isinstance(array, pa.ChunkedArray):
-            numpy_arr = np.hstack([chunk.to_numpy(zero_copy_only=False) for chunk in array.chunks])
-        else:
-            numpy_arr = array.to_numpy(zero_copy_only=False)
-        return PandasImageExtensionArray(numpy_arr)
-
-    @classmethod
-    def construct_array_type(cls):
-        return PandasImageExtensionArray
-
-    @property
-    def type(self) -> type:
-        # Expensive calls under the propery decorator are not a good practice, but it is what it is.
-        if config.PIL_AVAILABLE:
-            import PIL.Image
-        else:
-            raise ImportError("Pillow is not available.")
-        return PIL.Image.Image
-
-    @property
-    def kind(self) -> str:
-        return "O"
-
-    @property
-    def name(self) -> str:
-        return "image"
-
-
-class PandasImageExtensionArray(PandasExtensionArray):
-    na_value = None
-
-    def __init__(self, data: np.ndarray, copy: bool = False):
-        self._data = data if not copy else np.array(data)
-        self._dtype = PandasImageExtensionDtype()
-
-    def __array__(self):
-        return self._data
-
-    def copy(self, deep: bool = False) -> "PandasImageExtensionArray":
-        return PandasImageExtensionArray(self._data, copy=True)
-
-    @classmethod
-    def _from_sequence(
-        cls, scalars, dtype: Optional[PandasImageExtensionDtype] = None, copy: bool = False
-    ) -> "PandasImageExtensionArray":
-        data = np.array(scalars, dtype=np.object, copy=copy)
-        return cls(data, copy=copy)
-
-    @classmethod
-    def _concat_same_type(cls, to_concat: Sequence["PandasImageExtensionArray"]) -> "PandasImageExtensionArray":
-        data = np.hstack([va._data for va in to_concat])
-        return cls(data, copy=False)
-
-    @property
-    def dtype(self) -> PandasImageExtensionDtype:
-        return self._dtype
-
-    @property
-    def nbytes(self) -> int:
-        return self._data.nbytes
-
-    def isna(self) -> np.ndarray:
-        return np.array([pd.isna(arr).any() for arr in self._data])
-
-    def __setitem__(self, key: Union[int, slice, np.ndarray], value: Any) -> None:
-        raise NotImplementedError
-
-    def __getitem__(self, item: Union[int, slice, np.ndarray]) -> Union[np.ndarray, "PandasImageExtensionArray"]:
-        if isinstance(item, int):
-            return self._data[item]
-        return PandasImageExtensionArray(self._data[item], copy=False)
-
-    def take(
-        self, indices: Sequence[int], allow_fill: bool = False, fill_value: bool = None
-    ) -> "PandasImageExtensionArray":
-        indices: np.ndarray = np.asarray(indices, dtype=np.int)
-        if allow_fill:
-            fill_value = self.dtype.na_value if fill_value is None else np.asarray(fill_value, dtype=np.object)
-            mask = indices == -1
-            if (indices < -1).any():
-                raise ValueError("Invalid value in `indices`, must be all >= -1 for `allow_fill` is True")
-            elif len(self) > 0:
-                pass
-            elif not np.all(mask):
-                raise IndexError("Invalid take for empty PandasImageExtensionArray, must be all -1.")
-            else:
-                data = np.array([fill_value] * len(indices), dtype=np.object)
-                return PandasImageExtensionArray(data, copy=False)
-        took = self._data.take(indices)
-        if allow_fill and mask.any():
-            took[mask] = [fill_value] * np.sum(mask)
-        return PandasImageExtensionArray(took, copy=False)
-
-    def map(self, mapper):
-        # More info about this (undocumented) function can be found here:
-        # https://github.com/pandas-dev/pandas/issues/23179
-        return PandasImageExtensionArray(pd.Series(self._data).map(mapper).to_numpy())
-
-    def __len__(self) -> int:
-        return len(self._data)
-
-    def __eq__(self, other) -> np.ndarray:
-        if not isinstance(other, PandasImageExtensionArray):
-            raise NotImplementedError(f"Invalid type to compare to: {type(other)}")
-        return (self._data == other._data).all()
-
-
 @dataclass(unsafe_hash=True)
 class Image:
     """Image feature to read image data from an image file.
@@ -174,12 +38,13 @@ class Image:
 
     id: Optional[str] = None
     # Automatically constructed
-    dtype: ClassVar[str] = "dict"
-    pa_type: ClassVar[Any] = None
+    dtype: ClassVar[str] = "PIL.Image.Image"
+    # pa_type: ClassVar[Any] = ImageExtensionType
+    pa_type: ClassVar[Any] = pa.struct({"bytes": pa.binary(), "path": pa.string()})
     _type: str = field(default="Image", init=False, repr=False)
 
     def __call__(self):
-        return ImageExtensionType()
+        return self.pa_type
 
     def encode_example(self, value: Union[str, dict, np.ndarray, "PIL.Image.Image"]) -> dict:
         """Encode example into a format for Arrow.
@@ -188,7 +53,7 @@ def encode_example(self, value: Union[str, dict, np.ndarray, "PIL.Image.Image"])
             value (:obj:`str`, :obj:`np.ndarray`, :obj:`PIL.Image.Image` or :obj:`dict`): Data passed as input to Image feature.
 
         Returns:
-            :obj:`dict`
+            :obj:`dict` with "path" and "bytes" fields
         """
         if config.PIL_AVAILABLE:
             import PIL.Image
@@ -201,12 +66,13 @@ def encode_example(self, value: Union[str, dict, np.ndarray, "PIL.Image.Image"])
             image = PIL.Image.fromarray(value.astype(np.uint8))
             return {"path": None, "bytes": image_to_bytes(image)}
         elif isinstance(value, PIL.Image.Image):
-            if hasattr(value, "filename") and value.filename != "":
-                return {"path": value.filename, "bytes": None}
-            else:
-                return {"path": None, "bytes": image_to_bytes(value)}
+            return encode_pil_image(value)
+        elif value.get("bytes") is not None or value.get("path") is not None:
+            return {"bytes": value.get("bytes"), "path": value.get("path")}
         else:
-            return value
+            raise ValueError(
+                f"An audio sample should have one of 'path' or 'bytes' but they are missing or None in {value}."
+            )
 
     def decode_example(self, value):
         """Decode example image file into image data.
@@ -226,12 +92,14 @@ def decode_example(self, value):
             raise ImportError("To support decoding images, please install 'Pillow'.")
 
         if isinstance(value, str):
-            path, bytes_ = value, None
+            path, bytes_ = value, None  # TODO(QL, Mario): remove this case ?
         else:
             path, bytes_ = value["path"], value["bytes"]
 
         if bytes_ is None:
-            if isinstance(path, str):
+            if path is None:
+                raise ValueError(f"An image should have one of 'path' or 'bytes' but both are None in {value}.")
+            else:
                 if is_local_path(path):
                     image = PIL.Image.open(path)
                 else:
@@ -242,6 +110,50 @@ def decode_example(self, value):
             image = PIL.Image.open(BytesIO(bytes_))
         return image
 
+    def cast_storage(self, storage: Union[pa.StringArray, pa.StructArray, pa.ListArray]) -> pa.StructArray:
+        """Cast an Arrow array to the Image arrow storage type.
+        The Arrow types that can be converted to the Image pyarrow storage type are:
+
+        - pa.string() - it must contain the "path" data
+        - pa.struct({"bytes": pa.binary()})
+        - pa.struct({"path": pa.string()})
+        - pa.struct({"bytes": pa.binary(), "path": pa.string()})  - order doesn't matter
+        - pa.list(*) - it must contain the image array data
+
+        Args:
+            storage (Union[pa.StringArray, pa.StructArray, pa.ListArray]): [description]
+
+        Returns:
+            pa.StructArray: Array in the Image arrow storage type, that is
+                pa.struct({"bytes": pa.binary(), "path": pa.string()})
+        """
+        if config.PIL_AVAILABLE:
+            import PIL.Image
+        else:
+            raise ImportError("To support encoding images, please install 'Pillow'.")
+
+        if pa.types.is_string(storage.type):
+            bytes_array = pa.array([None] * len(storage), type=pa.binary())
+            storage = pa.StructArray.from_arrays([bytes_array, storage], ["bytes", "path"])
+        elif pa.types.is_struct(storage.type):
+            if storage.type.get_field_index("bytes") >= 0:
+                bytes_array = storage.field("bytes")
+            else:
+                bytes_array = pa.array([None] * len(storage), type=pa.binary())
+            if storage.type.get_field_index("path") >= 0:
+                path_array = storage.field("path")
+            else:
+                path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"])
+        elif pa.types.is_list(storage.type):
+            bytes_array = pa.array(
+                [image_to_bytes(PIL.Image.fromarray(np.array(arr, np.uint8))) for arr in storage.to_pylist()],
+                type=pa.binary(),
+            )
+            path_array = pa.array([None] * len(storage), type=pa.string())
+            storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"])
+        return array_cast(storage, self.pa_type)
+
 
 def list_image_compression_formats() -> List[str]:
     if config.PIL_AVAILABLE:
@@ -264,6 +176,13 @@ def image_to_bytes(image: "PIL.Image.Image") -> bytes:
     return buffer.getvalue()
 
 
+def encode_pil_image(image: "PIL.Image.Image") -> dict:
+    if hasattr(image, "filename") and image.filename != "":
+        return {"path": image.filename, "bytes": None}
+    else:
+        return {"path": None, "bytes": image_to_bytes(image)}
+
+
 def objects_to_list_of_image_dicts(
     objs: Union[List[str], List[dict], List[np.ndarray], List["PIL.Image.Image"]]
 ) -> List[dict]:
@@ -285,11 +204,7 @@ def objects_to_list_of_image_dicts(
                 for obj in objs
             ]
         elif isinstance(obj, PIL.Image.Image):
-            obj_to_image_dict_func = no_op_if_value_is_null(
-                lambda obj: {"path": obj.filename, "bytes": None}
-                if hasattr(obj, "filename") and obj.filename != ""
-                else {"path": None, "bytes": image_to_bytes(obj)}
-            )
+            obj_to_image_dict_func = no_op_if_value_is_null(encode_pil_image)
             return [obj_to_image_dict_func(obj) for obj in objs]
         else:
             return objs
diff --git a/src/datasets/iterable_dataset.py b/src/datasets/iterable_dataset.py
index 25144b515eb..d16f26d296e 100644
--- a/src/datasets/iterable_dataset.py
+++ b/src/datasets/iterable_dataset.py
@@ -11,14 +11,15 @@
 from .formatting import PythonFormatter
 from .info import DatasetInfo
 from .splits import NamedSplit
+from .table import table_cast
 
 
 def _infer_features_from_batch(batch: Dict[str, list], try_features: Optional[Features] = None) -> Features:
     pa_table = pa.Table.from_pydict(batch)
     if try_features is not None:
         try:
-            pa_table = pa_table.cast(pa.schema(try_features.type))
-        except (pa.ArrowInvalid, pa.ArrowNotImplementedError):
+            pa_table = table_cast(pa_table, pa.schema(try_features.type))
+        except (TypeError, pa.ArrowInvalid, pa.ArrowNotImplementedError):
             pass
     return Features.from_arrow_schema(pa_table.schema)
 
diff --git a/src/datasets/packaged_modules/json/json.py b/src/datasets/packaged_modules/json/json.py
index ba289463e71..369bf9c4e60 100644
--- a/src/datasets/packaged_modules/json/json.py
+++ b/src/datasets/packaged_modules/json/json.py
@@ -8,6 +8,7 @@
 import pyarrow.json as paj
 
 import datasets
+from datasets.table import table_cast
 from datasets.utils.file_utils import readline
 
 
@@ -27,7 +28,7 @@ class JsonConfig(datasets.BuilderConfig):
 
     @property
     def schema(self):
-        return pa.schema(self.features.type) if self.features is not None else None
+        return self.features.arrow_schema if self.features is not None else None
 
 
 class Json(datasets.ArrowBasedBuilder):
@@ -77,11 +78,8 @@ def _cast_classlabels(self, pa_table: pa.Table) -> pa.Table:
                         raise ValueError(
                             f"Field '{col}' from the JSON data of type {pa_table[col].type} is not compatible with ClassLabel. Compatible types are int64 and string."
                         )
-            # Cast allows str <-> int/float
-            # Before casting, rearrange JSON field names to match passed features schema field names order
-            pa_table = pa.Table.from_arrays(
-                [pa_table[name] for name in self.config.features], schema=self.config.schema
-            )
+            # Cast allows str <-> int/float or str to Audio for example
+            pa_table = table_cast(pa_table, self.config.schema)
         return pa_table
 
     def _generate_tables(self, files):
diff --git a/src/datasets/table.py b/src/datasets/table.py
index bc1ab5d4e25..c291d566468 100644
--- a/src/datasets/table.py
+++ b/src/datasets/table.py
@@ -1,20 +1,27 @@
 import copy
 import os
 import tempfile
-from functools import wraps
+from functools import partial, wraps
 from itertools import groupby
-from typing import List, Optional, Tuple, TypeVar, Union
+from typing import TYPE_CHECKING, List, Optional, Tuple, TypeVar, Union
 
 import numpy as np
 import pyarrow as pa
 import pyarrow.compute as pc
+from packaging import version
 
 from . import config
 from .utils.logging import get_logger
 
 
+if TYPE_CHECKING:
+    from .features import Features, FeatureType
+
+
 logger = get_logger(__name__)
 
+IS_PYARROW_AT_LEAST_4 = config.PYARROW_VERSION.major >= 4
+
 
 def inject_arrow_table_documentation(arrow_table_method):
     def wrapper(method):
@@ -318,6 +325,18 @@ def rename_columns(self, *args, **kwargs):
     def drop(self, *args, **kwargs):
         raise NotImplementedError()
 
+    # Additional methods that are based on the PyArrow Table methods
+
+    def select_columns(self, columns: List[int]) -> "Table":
+        """Return the table by keeping only the requested columns
+
+        Returns:
+            Table: table with only a subset of the columns
+        """
+        for column_to_remove in set(range(len(self.column_names))) - set(columns):
+            self = self.remove_column(column_to_remove)
+        return self
+
 
 class TableBlock(Table):
     """
@@ -393,7 +412,7 @@ def combine_chunks(self, *args, **kwargs):
 
     @inject_arrow_table_documentation(pa.Table.cast)
     def cast(self, *args, **kwargs):
-        return InMemoryTable(self.table.cast(*args, **kwargs))
+        return InMemoryTable(table_cast(self.table, *args, **kwargs))
 
     @inject_arrow_table_documentation(pa.Table.replace_schema_metadata)
     def replace_schema_metadata(self, *args, **kwargs):
@@ -474,7 +493,10 @@ def __setstate__(self, state):
     def _apply_replays(table: pa.Table, replays: Optional[List[Replay]] = None) -> pa.Table:
         if replays is not None:
             for name, args, kwargs in replays:
-                table = getattr(table, name)(*args, **kwargs)
+                if name == "cast":
+                    table = table_cast(table, *args, **kwargs)
+                else:
+                    table = getattr(table, name)(*args, **kwargs)
         return table
 
     def _append_replay(self, replay: Replay) -> List[Replay]:
@@ -511,7 +533,7 @@ def combine_chunks(self, *args, **kwargs):
     def cast(self, *args, **kwargs):
         replay = ("cast", copy.deepcopy(args), copy.deepcopy(kwargs))
         replays = self._append_replay(replay)
-        return MemoryMappedTable(self.table.cast(*args, **kwargs), self.path, replays)
+        return MemoryMappedTable(table_cast(self.table, *args, **kwargs), self.path, replays)
 
     @inject_arrow_table_documentation(pa.Table.replace_schema_metadata)
     def replace_schema_metadata(self, *args, **kwargs):
@@ -777,7 +799,7 @@ def combine_chunks(self, *args, **kwargs):
 
     @inject_arrow_table_documentation(pa.Table.cast)
     def cast(self, target_schema, *args, **kwargs):
-        table = self.table.cast(target_schema, *args, **kwargs)
+        table = table_cast(self.table, target_schema, *args, **kwargs)
         blocks = []
         for subtables in self.blocks:
             new_tables = []
@@ -888,23 +910,24 @@ def list_table_cache_files(table: Table) -> List[str]:
         return []
 
 
-def cast_with_sliced_list_support(pa_table: pa.Table, schema: pa.Schema) -> pa.Table:
-    """Same as pyarrow.Table.cast, except it works for sliced list arrays"""
+def _wrap_for_chunked_arrays(func):
+    """Apply the function on each chunk of a pyarrow.ChunkedArray, or on the array directly"""
 
-    def wrap_for_chunked_arrays(func):
-        """Apply the function on each chunk of a pyarrow.ChunkedArray, or on the array directly"""
+    def wrapper(array, *args, **kwargs):
+        if isinstance(array, pa.ChunkedArray):
+            return pa.chunked_array([func(chunk, *args, **kwargs) for chunk in array.chunks])
+        else:
+            return func(array, *args, **kwargs)
+
+    return wrapper
 
-        def wrapper(array):
-            if isinstance(array, pa.ChunkedArray):
-                return pa.chunked_array([func(chunk) for chunk in array.chunks])
-            else:
-                return func(array)
 
-        return wrapper
+def _sanitize_sliced_list_arrays_for_cast(func):
+    """Sanitize pyarrow.ListArray objects for pyarrow.Array.cast in case they are sliced list arrays"""
 
-    @wrap_for_chunked_arrays
-    def reset_sliced_list_offset(array: pa.ListArray):
-        """Return the same pyarrow.ListArray but with array.offset == 0 for compatibility with cast"""
+    @_wrap_for_chunked_arrays
+    def _sanitize(array: pa.ListArray) -> pa.ListArray:
+        """Return the same pyarrow.ListArray but with array.offset == 0 for compatibility with cast for pyarrow 3"""
         if array.offset == 0:
             return array
         elif len(array) == 0:
@@ -915,5 +938,188 @@ def reset_sliced_list_offset(array: pa.ListArray):
             new_offsets = pc.subtract(array.offsets, values_offset)  # update the offsets accordingly
             return pa.ListArray.from_arrays(new_offsets, new_values)
 
-    arrays = [reset_sliced_list_offset(array) if isinstance(array.type, pa.ListType) else array for array in pa_table]
-    return pa.Table.from_arrays(arrays, schema=schema)
+    def wrapper(array, *args, **kwargs):
+        if pa.types.is_list(array.type) and config.PYARROW_VERSION < version.parse("4.0.0"):
+            array = _sanitize(array)
+        return func(array, *args, **kwargs)
+
+    return wrapper
+
+
+@_sanitize_sliced_list_arrays_for_cast
+@_wrap_for_chunked_arrays
+def array_cast(array: pa.Array, pa_type: pa.DataType, allow_number_to_str=True):
+    """Improved version of pa.Array.cast
+
+    It supports casting pa.StructArray objects to re-order the fields.
+    It also let you control certain aspects of the casting, e.g. whether
+    to disable numbers (floats or ints) to strings.
+
+    Args:
+        array (pa.Array): PyArrow array to cast
+        pa_type (pa.DataType): target PyArrow type
+        allow_number_to_str (bool, optional): Weheter to allow casting numbers to strings.
+            Defaults to True.
+
+    Raises:
+        pa.ArrowInvalidError: if the arrow data casting fails
+        TypeError: if the target type is not supported according, e.g.
+
+            - if a field is missing
+            = if casting from numbers to strings and allow_number_to_str is False
+
+    Returns:
+        pa.Array: the casted array
+    """
+    _c = partial(array_cast, allow_number_to_str=allow_number_to_str)
+    if isinstance(array, pa.ExtensionArray):
+        array = array.storage
+    if isinstance(pa_type, pa.ExtensionType):
+        return pa_type.wrap_array(array)
+    elif pa.types.is_struct(array.type):
+        if pa.types.is_struct(pa_type) and (
+            set(field.name for field in pa_type) == set(field.name for field in array.type)
+        ):
+            arrays = [
+                _c(array.field(field.name), field.type, allow_number_to_str=allow_number_to_str) for field in pa_type
+            ]
+            return pa.StructArray.from_arrays(arrays, fields=list(pa_type))
+    elif pa.types.is_list(array.type):
+        if pa.types.is_fixed_size_list(pa_type):
+            if pa_type.list_size * len(array) == len(array.values):
+                return pa.FixedSizeListArray.from_arrays(
+                    _c(array.values, pa_type.value_type, allow_number_to_str=allow_number_to_str),
+                    pa_type.list_size,
+                )
+        elif pa.types.is_list(pa_type):
+            return pa.ListArray.from_arrays(
+                array.offsets, _c(array.values, pa_type.value_type, allow_number_to_str=allow_number_to_str)
+            )
+    elif pa.types.is_fixed_size_list(array.type):
+        if pa.types.is_fixed_size_list(pa_type):
+            return pa.FixedSizeListArray.from_arrays(
+                _c(array.values, pa_type.value_type, allow_number_to_str=allow_number_to_str),
+                pa_type.list_size,
+            )
+        elif pa.types.is_list(pa_type):
+            offsets_arr = pa.array(range(len(array) + 1), pa.int32())
+            return pa.ListArray.from_arrays(
+                offsets_arr, _c(array.values, pa_type.value_type, allow_number_to_str=allow_number_to_str)
+            )
+    else:
+        if (
+            not allow_number_to_str
+            and pa.types.is_string(pa_type)
+            and (pa.types.is_floating(array.type) or pa.types.is_integer(array.type))
+        ):
+            raise TypeError(
+                f"Couldn't cast array of type {array.type} to {pa_type} since allow_number_to_str is set to {allow_number_to_str}"
+            )
+        return array.cast(pa_type)
+    raise TypeError(f"Couldn't cast array of type\n{array.type}\nto\n{pa_type}")
+
+
+@_sanitize_sliced_list_arrays_for_cast
+@_wrap_for_chunked_arrays
+def cast_array_to_feature(array: pa.Array, feature: "FeatureType", allow_number_to_str=True):
+    """Cast an array to the arrow type that corresponds to the requested feature type.
+    For custom features like Audio or Image, it takes into account the "cast_storage" methods
+    they defined to enable casting from other arrow types.
+
+    Args:
+        array (pa.Array): the PyArrow array to cast
+        feature (FeatureType): the target feature type
+        allow_number_to_str (bool, optional): Weheter to allow casting numbers to strings.
+            Defaults to True.
+
+    Raises:
+        pa.ArrowInvalidError: if the arrow data casting fails
+        TypeError: if the target type is not supported according, e.g.
+
+            - if a field is missing
+            = if casting from numbers to strings and allow_number_to_str is False
+
+    Returns:
+        pa.Array: the casted array
+    """
+    from .features import Sequence, get_nested_type
+
+    _c = partial(cast_array_to_feature, allow_number_to_str=allow_number_to_str)
+
+    if isinstance(array, pa.ExtensionArray):
+        array = array.storage
+    if hasattr(feature, "cast_storage"):
+        return feature.cast_storage(array)
+    elif pa.types.is_struct(array.type):
+        # feature must be a dict or Sequence(subfeatures_dict)
+        if isinstance(feature, Sequence) and isinstance(feature.feature, dict):
+            feature = {
+                name: Sequence(subfeature, length=feature.length) for name, subfeature in feature.feature.items()
+            }
+        if isinstance(feature, dict) and set(field.name for field in array.type) == set(feature):
+            arrays = [_c(array.field(name), subfeature) for name, subfeature in feature.items()]
+            return pa.StructArray.from_arrays(arrays, names=list(feature))
+    elif pa.types.is_list(array.type):
+        # feature must be either [subfeature] or Sequence(subfeature)
+        if isinstance(feature, list):
+            return pa.ListArray.from_arrays(array.offsets, _c(array.values, feature[0]))
+        elif isinstance(feature, Sequence):
+            if feature.length > -1:
+                if feature.length * len(array) == len(array.values):
+                    return pa.FixedSizeListArray.from_arrays(_c(array.values, feature.feature), feature.length)
+            else:
+                return pa.ListArray.from_arrays(array.offsets, _c(array.values, feature.feature))
+    elif pa.types.is_fixed_size_list(array.type):
+        # feature must be either [subfeature] or Sequence(subfeature)
+        if isinstance(feature, list):
+            return pa.ListArray.from_arrays(array.offsets, _c(array.values, feature[0]))
+        elif isinstance(feature, Sequence):
+            if feature.length > -1:
+                if feature.length * len(array) == len(array.values):
+                    return pa.FixedSizeListArray.from_arrays(_c(array.values, feature.feature), feature.length)
+            else:
+                offsets_arr = pa.array(range(len(array) + 1), pa.int32())
+                return pa.ListArray.from_arrays(offsets_arr, _c(array.values, feature.feature))
+    if pa.types.is_null(array.type):
+        return array_cast(array, get_nested_type(feature), allow_number_to_str=allow_number_to_str)
+    elif not isinstance(feature, (Sequence, dict, list, tuple)):
+        return array_cast(array, feature(), allow_number_to_str=allow_number_to_str)
+    raise TypeError(f"Couldn't cast array of type\n{array.type}\nto\n{feature}")
+
+
+def cast_table_to_features(table: pa.Table, features: "Features"):
+    """Cast an table to the arrow schema that corresponds to the requested features.
+
+    Args:
+        table (pa.Table): PyArrow table to cast
+        features (Features): target features.
+
+    Returns:
+        pa.Table: the casted table
+    """
+    if sorted(table.column_names) != sorted(features):
+        raise ValueError(f"Couldn't cast\n{table.schema}\nto\n{features}\nbecause column names don't match")
+    arrays = [cast_array_to_feature(table[name], feature) for name, feature in features.items()]
+    return pa.Table.from_arrays(arrays, schema=features.arrow_schema)
+
+
+def table_cast(table: pa.Table, schema: pa.Schema):
+    """Improved version of pa.Table.cast
+
+    It supports casting to feature types stored in the schema metadata.
+
+    Args:
+        table (pa.Table): PyArrow table to cast
+        schema (pa.Schema): target PyArrow schema.
+
+    Returns:
+        pa.Table: the casted table
+    """
+    if table.schema != schema:
+        from .features import Features
+
+        return cast_table_to_features(table, Features.from_arrow_schema(schema))
+    elif table.schema.metadata != schema.metadata:
+        return table.replace_schema_metadata(schema.metadata)
+    else:
+        return table
diff --git a/tests/features/test_audio.py b/tests/features/test_audio.py
index 2c45eebda04..793f7328158 100644
--- a/tests/features/test_audio.py
+++ b/tests/features/test_audio.py
@@ -4,10 +4,11 @@
 from ctypes.util import find_library
 from importlib.util import find_spec
 
+import pyarrow as pa
 import pytest
 
-from datasets import Dataset, load_dataset
-from datasets.features import Audio, Features, Value
+from datasets import Dataset, concatenate_datasets, load_dataset
+from datasets.features import Audio, Features, Sequence, Value
 
 
 # pytestmark = pytest.mark.audio
@@ -60,16 +61,46 @@ def test_audio_instantiation():
     assert audio.mono is True
     assert audio.id is None
     assert audio.dtype == "dict"
-    assert audio.pa_type is None
+    assert audio.pa_type == pa.struct({"bytes": pa.binary(), "path": pa.string()})
     assert audio._type == "Audio"
-    assert audio._storage_dtype == "string"
+
+
+def test_audio_feature_type_to_arrow():
+    features = Features({"audio": Audio()})
+    assert features.arrow_schema == pa.schema({"audio": Audio().pa_type})
+    features = Features({"struct_containing_an_audio": {"audio": Audio()}})
+    assert features.arrow_schema == pa.schema({"struct_containing_an_audio": pa.struct({"audio": Audio().pa_type})})
+    features = Features({"sequence_of_audios": Sequence(Audio())})
+    assert features.arrow_schema == pa.schema({"sequence_of_audios": pa.list_(Audio().pa_type)})
+
+
+@pytest.mark.parametrize(
+    "build_example",
+    [
+        lambda audio_path: audio_path,
+        lambda audio_path: {"path": audio_path},
+        lambda audio_path: {"path": audio_path, "bytes": None},
+        lambda audio_path: {"path": audio_path, "bytes": open(audio_path, "rb").read()},
+        lambda audio_path: {"path": None, "bytes": open(audio_path, "rb").read()},
+        lambda audio_path: {"bytes": open(audio_path, "rb").read()},
+    ],
+)
+def test_audio_feature_encode_example(shared_datadir, build_example):
+    audio_path = str(shared_datadir / "test_audio_44100.wav")
+    audio = Audio()
+    encoded_example = audio.encode_example(build_example(audio_path))
+    assert isinstance(encoded_example, dict)
+    assert encoded_example.keys() == {"bytes", "path"}
+    assert encoded_example["bytes"] is not None or encoded_example["path"] is not None
+    decoded_example = audio.decode_example(encoded_example)
+    assert decoded_example.keys() == {"path", "array", "sampling_rate"}
 
 
 @require_sndfile
 def test_audio_decode_example(shared_datadir):
     audio_path = str(shared_datadir / "test_audio_44100.wav")
     audio = Audio()
-    decoded_example = audio.decode_example(audio_path)
+    decoded_example = audio.decode_example(audio.encode_example(audio_path))
     assert decoded_example.keys() == {"path", "array", "sampling_rate"}
     assert decoded_example["path"] == audio_path
     assert decoded_example["array"].shape == (202311,)
@@ -80,7 +111,7 @@ def test_audio_decode_example(shared_datadir):
 def test_audio_resampling(shared_datadir):
     audio_path = str(shared_datadir / "test_audio_44100.wav")
     audio = Audio(sampling_rate=16000)
-    decoded_example = audio.decode_example(audio_path)
+    decoded_example = audio.decode_example(audio.encode_example(audio_path))
     assert decoded_example.keys() == {"path", "array", "sampling_rate"}
     assert decoded_example["path"] == audio_path
     assert decoded_example["array"].shape == (73401,)
@@ -92,7 +123,7 @@ def test_audio_resampling(shared_datadir):
 def test_audio_decode_example_mp3(shared_datadir):
     audio_path = str(shared_datadir / "test_audio_44100.mp3")
     audio = Audio()
-    decoded_example = audio.decode_example(audio_path)
+    decoded_example = audio.decode_example(audio.encode_example(audio_path))
     assert decoded_example.keys() == {"path", "array", "sampling_rate"}
     assert decoded_example["path"] == audio_path
     assert decoded_example["array"].shape == (109440,)
@@ -303,6 +334,62 @@ def test_resampling_after_loading_dataset_with_audio_feature_mp3(shared_datadir)
     assert column[0]["sampling_rate"] == 16000
 
 
+@pytest.mark.parametrize(
+    "build_data",
+    [
+        lambda audio_path: {"audio": [audio_path]},
+        lambda audio_path: {"audio": [{"path": audio_path}]},
+        lambda audio_path: {"audio": [{"path": audio_path, "bytes": None}]},
+        lambda audio_path: {"audio": [{"path": audio_path, "bytes": open(audio_path, "rb").read()}]},
+        lambda audio_path: {"audio": [{"path": None, "bytes": open(audio_path, "rb").read()}]},
+        lambda audio_path: {"audio": [{"bytes": open(audio_path, "rb").read()}]},
+    ],
+)
+def test_dataset_cast_to_audio_features(shared_datadir, build_data):
+    audio_path = str(shared_datadir / "test_audio_44100.wav")
+    data = build_data(audio_path)
+    dset = Dataset.from_dict(data)
+    item = dset.cast(Features({"audio": Audio()}))[0]
+    assert item.keys() == {"audio"}
+    assert item["audio"].keys() == {"path", "array", "sampling_rate"}
+    item = dset.cast_column("audio", Audio())[0]
+    assert item.keys() == {"audio"}
+    assert item["audio"].keys() == {"path", "array", "sampling_rate"}
+
+
+def test_dataset_concatenate_audio_features(shared_datadir):
+    # we use a different data structure between 1 and 2 to make sure they are compatible with each other
+    audio_path = str(shared_datadir / "test_audio_44100.wav")
+    data1 = {"audio": [audio_path]}
+    dset1 = Dataset.from_dict(data1, features=Features({"audio": Audio()}))
+    data2 = {"audio": [{"bytes": open(audio_path, "rb").read()}]}
+    dset2 = Dataset.from_dict(data2, features=Features({"audio": Audio()}))
+    concatenated_dataset = concatenate_datasets([dset1, dset2])
+    assert len(concatenated_dataset) == len(dset1) + len(dset2)
+    assert concatenated_dataset[0]["audio"]["array"].shape == dset1[0]["audio"]["array"].shape
+    assert concatenated_dataset[1]["audio"]["array"].shape == dset2[0]["audio"]["array"].shape
+
+
+def test_dataset_concatenate_nested_audio_features(shared_datadir):
+    # we use a different data structure between 1 and 2 to make sure they are compatible with each other
+    audio_path = str(shared_datadir / "test_audio_44100.wav")
+    features = Features({"list_of_structs_of_audios": [{"audio": Audio()}]})
+    data1 = {"list_of_structs_of_audios": [[{"audio": audio_path}]]}
+    dset1 = Dataset.from_dict(data1, features=features)
+    data2 = {"list_of_structs_of_audios": [[{"audio": {"bytes": open(audio_path, "rb").read()}}]]}
+    dset2 = Dataset.from_dict(data2, features=features)
+    concatenated_dataset = concatenate_datasets([dset1, dset2])
+    assert len(concatenated_dataset) == len(dset1) + len(dset2)
+    assert (
+        concatenated_dataset[0]["list_of_structs_of_audios"][0]["audio"]["array"].shape
+        == dset1[0]["list_of_structs_of_audios"][0]["audio"]["array"].shape
+    )
+    assert (
+        concatenated_dataset[1]["list_of_structs_of_audios"][0]["audio"]["array"].shape
+        == dset2[0]["list_of_structs_of_audios"][0]["audio"]["array"].shape
+    )
+
+
 @require_sndfile
 def test_dataset_with_audio_feature_map_is_not_decoded(shared_datadir):
     audio_path = str(shared_datadir / "test_audio_44100.wav")
@@ -310,9 +397,10 @@ def test_dataset_with_audio_feature_map_is_not_decoded(shared_datadir):
     features = Features({"audio": Audio(), "text": Value("string")})
     dset = Dataset.from_dict(data, features=features)
 
+    expected_audio = features.encode_batch(data)["audio"][0]
     for item in dset:
         assert item.keys() == {"audio", "text"}
-        assert item == {"audio": audio_path, "text": "Hello"}
+        assert item == {"audio": expected_audio, "text": "Hello"}
 
     def process_text(example):
         example["text"] = example["text"] + " World!"
@@ -321,7 +409,7 @@ def process_text(example):
     processed_dset = dset.map(process_text)
     for item in processed_dset:
         assert item.keys() == {"audio", "text"}
-        assert item == {"audio": audio_path, "text": "Hello World!"}
+        assert item == {"audio": expected_audio, "text": "Hello World!"}
 
 
 @require_sndfile
diff --git a/tests/features/test_features.py b/tests/features/test_features.py
index e1916dec36c..51c03ee6e1f 100644
--- a/tests/features/test_features.py
+++ b/tests/features/test_features.py
@@ -197,6 +197,13 @@ def test_reorder_fields_as(self):
         self.assertEqual(reordered_features.type, expected.type)
         self.assertNotEqual(reordered_features.type, features.type)
 
+    def test_flatten(self):
+        features = Features({"foo": {"bar1": Value("int32"), "bar2": {"foobar": Value("string")}}})
+        _features = features.copy()
+        flattened_features = features.flatten()
+        assert flattened_features == {"foo.bar1": Value("int32"), "foo.bar2.foobar": Value("string")}
+        assert features == _features, "calling flatten shouldn't alter the current features"
+
 
 def test_classlabel_init(tmp_path_factory):
     names = ["negative", "positive"]
diff --git a/tests/features/test_image.py b/tests/features/test_image.py
index e2959a5e5b4..ec084149d04 100644
--- a/tests/features/test_image.py
+++ b/tests/features/test_image.py
@@ -3,9 +3,11 @@
 
 import numpy as np
 import pandas as pd
+import pyarrow as pa
 import pytest
 
-from datasets import Dataset, Features, Image, Value, load_dataset
+from datasets import Dataset, Features, Image, Sequence, Value, load_dataset
+from datasets.arrow_dataset import concatenate_datasets
 from datasets.features.image import image_to_bytes
 
 from ..utils import require_pil
@@ -31,11 +33,45 @@ def iter_archive(archive_path):
 def test_image_instantiation():
     image = Image()
     assert image.id is None
-    assert image.dtype == "dict"
-    assert image.pa_type is None
+    assert image.dtype == "PIL.Image.Image"
+    assert image.pa_type == pa.struct({"bytes": pa.binary(), "path": pa.string()})
     assert image._type == "Image"
 
 
+def test_image_feature_type_to_arrow():
+    features = Features({"image": Image()})
+    assert features.arrow_schema == pa.schema({"image": Image().pa_type})
+    features = Features({"struct_containing_an_image": {"image": Image()}})
+    assert features.arrow_schema == pa.schema({"struct_containing_an_image": pa.struct({"image": Image().pa_type})})
+    features = Features({"sequence_of_images": Sequence(Image())})
+    assert features.arrow_schema == pa.schema({"sequence_of_images": pa.list_(Image().pa_type)})
+
+
+@require_pil
+@pytest.mark.parametrize(
+    "build_example",
+    [
+        lambda image_path: image_path,
+        lambda image_path: {"path": image_path},
+        lambda image_path: {"path": image_path, "bytes": None},
+        lambda image_path: {"path": image_path, "bytes": open(image_path, "rb").read()},
+        lambda image_path: {"path": None, "bytes": open(image_path, "rb").read()},
+        lambda image_path: {"bytes": open(image_path, "rb").read()},
+    ],
+)
+def test_image_feature_encode_example(shared_datadir, build_example):
+    import PIL.Image
+
+    image_path = str(shared_datadir / "test_image_rgb.jpg")
+    image = Image()
+    encoded_example = image.encode_example(build_example(image_path))
+    assert isinstance(encoded_example, dict)
+    assert encoded_example.keys() == {"bytes", "path"}
+    assert encoded_example["bytes"] is not None or encoded_example["path"] is not None
+    decoded_example = image.decode_example(encoded_example)
+    assert isinstance(decoded_example, PIL.Image.Image)
+
+
 @require_pil
 def test_image_decode_example(shared_datadir):
     import PIL.Image
@@ -147,6 +183,7 @@ def test_dataset_with_image_feature_from_np_array():
     assert column[0].size == (640, 480)
 
 
+@require_pil
 def test_dataset_with_image_feature_tar_jpg(tar_jpg_path):
     import PIL.Image
 
@@ -181,6 +218,67 @@ def test_dataset_with_image_feature_tar_jpg(tar_jpg_path):
     assert column[0].mode == "RGB"
 
 
+@require_pil
+@pytest.mark.parametrize(
+    "build_data",
+    [
+        lambda image_path: {"image": [image_path]},
+        lambda image_path: {"image": [{"path": image_path}]},
+        lambda image_path: {"image": [{"path": image_path, "bytes": None}]},
+        lambda image_path: {"image": [{"path": image_path, "bytes": open(image_path, "rb").read()}]},
+        lambda image_path: {"image": [{"path": None, "bytes": open(image_path, "rb").read()}]},
+        lambda image_path: {"image": [{"bytes": open(image_path, "rb").read()}]},
+    ],
+)
+def test_dataset_cast_to_image_features(shared_datadir, build_data):
+    import PIL.Image
+
+    image_path = str(shared_datadir / "test_image_rgb.jpg")
+    data = build_data(image_path)
+    dset = Dataset.from_dict(data)
+    item = dset.cast(Features({"image": Image()}))[0]
+    assert item.keys() == {"image"}
+    assert isinstance(item["image"], PIL.Image.Image)
+    item = dset.cast_column("image", Image())[0]
+    assert item.keys() == {"image"}
+    assert isinstance(item["image"], PIL.Image.Image)
+
+
+@require_pil
+def test_dataset_concatenate_image_features(shared_datadir):
+    # we use a different data structure between 1 and 2 to make sure they are compatible with each other
+    image_path = str(shared_datadir / "test_image_rgb.jpg")
+    data1 = {"image": [image_path]}
+    dset1 = Dataset.from_dict(data1, features=Features({"image": Image()}))
+    data2 = {"image": [{"bytes": open(image_path, "rb").read()}]}
+    dset2 = Dataset.from_dict(data2, features=Features({"image": Image()}))
+    concatenated_dataset = concatenate_datasets([dset1, dset2])
+    assert len(concatenated_dataset) == len(dset1) + len(dset2)
+    assert concatenated_dataset[0]["image"] == dset1[0]["image"]
+    assert concatenated_dataset[1]["image"] == dset2[0]["image"]
+
+
+@require_pil
+def test_dataset_concatenate_nested_image_features(shared_datadir):
+    # we use a different data structure between 1 and 2 to make sure they are compatible with each other
+    image_path = str(shared_datadir / "test_image_rgb.jpg")
+    features = Features({"list_of_structs_of_images": [{"image": Image()}]})
+    data1 = {"list_of_structs_of_images": [[{"image": image_path}]]}
+    dset1 = Dataset.from_dict(data1, features=features)
+    data2 = {"list_of_structs_of_images": [[{"image": {"bytes": open(image_path, "rb").read()}}]]}
+    dset2 = Dataset.from_dict(data2, features=features)
+    concatenated_dataset = concatenate_datasets([dset1, dset2])
+    assert len(concatenated_dataset) == len(dset1) + len(dset2)
+    assert (
+        concatenated_dataset[0]["list_of_structs_of_images"][0]["image"]
+        == dset1[0]["list_of_structs_of_images"][0]["image"]
+    )
+    assert (
+        concatenated_dataset[1]["list_of_structs_of_images"][0]["image"]
+        == dset2[0]["list_of_structs_of_images"][0]["image"]
+    )
+
+
 @require_pil
 def test_dataset_with_image_feature_map(shared_datadir):
     image_path = str(shared_datadir / "test_image_rgb.jpg")
@@ -242,7 +340,12 @@ def test_dataset_with_image_feature_map_change_image(shared_datadir):
 
     for item in dset:
         assert item.keys() == {"image"}
-        assert item == {"image": {"path": image_path, "bytes": None}}
+        assert item == {
+            "image": {
+                "bytes": None,
+                "path": image_path,
+            }
+        }
 
     # return pil image
 
@@ -253,7 +356,7 @@ def process_image_resize_by_example(example):
     decoded_dset = dset.map(process_image_resize_by_example)
     for item in decoded_dset:
         assert item.keys() == {"image"}
-        assert item == {"image": {"path": None, "bytes": image_to_bytes(pil_image.resize((100, 100)))}}
+        assert item == {"image": {"bytes": image_to_bytes(pil_image.resize((100, 100))), "path": None}}
 
     def process_image_resize_by_batch(batch):
         batch["image"] = [image.resize((100, 100)) for image in batch["image"]]
@@ -262,7 +365,7 @@ def process_image_resize_by_batch(batch):
     decoded_dset = dset.map(process_image_resize_by_batch, batched=True)
     for item in decoded_dset:
         assert item.keys() == {"image"}
-        assert item == {"image": {"path": None, "bytes": image_to_bytes(pil_image.resize((100, 100)))}}
+        assert item == {"image": {"bytes": image_to_bytes(pil_image.resize((100, 100))), "path": None}}
 
     # return np.ndarray (e.g. when using albumentations)
 
@@ -275,8 +378,8 @@ def process_image_resize_by_example_return_np_array(example):
         assert item.keys() == {"image"}
         assert item == {
             "image": {
-                "path": None,
                 "bytes": image_to_bytes(PIL.Image.fromarray(np.array(pil_image.resize((100, 100))))),
+                "path": None,
             }
         }
 
@@ -289,8 +392,8 @@ def process_image_resize_by_batch_return_np_array(batch):
         assert item.keys() == {"image"}
         assert item == {
             "image": {
-                "path": None,
                 "bytes": image_to_bytes(PIL.Image.fromarray(np.array(pil_image.resize((100, 100))))),
+                "path": None,
             }
         }
 
diff --git a/tests/test_arrow_dataset.py b/tests/test_arrow_dataset.py
index 531045fd376..a0d9dbd9d3b 100644
--- a/tests/test_arrow_dataset.py
+++ b/tests/test_arrow_dataset.py
@@ -80,7 +80,7 @@ def assert_arrow_metadata_are_synced_with_dataset_features(dataset: Dataset):
     features = DatasetInfo.from_dict(metadata["info"]).features
     assert features is not None
     assert dataset.features is not None
-    assert features.type == dataset.features.type
+    assert sorted(features) == sorted(field.name for field in dataset.data.schema)
 
 
 IN_MEMORY_PARAMETERS = [
@@ -764,8 +764,8 @@ def test_flatten(self, in_memory):
                 with self._to(in_memory, tmp_dir, dset) as dset:
                     fingerprint = dset._fingerprint
                     dset.flatten_()
-                    self.assertListEqual(dset.column_names, ["a.b.c", "foo"])
-                    self.assertListEqual(list(dset.features.keys()), ["a.b.c", "foo"])
+                    self.assertListEqual(sorted(dset.column_names), ["a.b.c", "foo"])
+                    self.assertListEqual(sorted(dset.features.keys()), ["a.b.c", "foo"])
                     self.assertDictEqual(
                         dset.features, Features({"a.b.c": Sequence(Value("string")), "foo": Value("int64")})
                     )
diff --git a/tests/test_arrow_writer.py b/tests/test_arrow_writer.py
index e5f94892cb5..3fda2be4fe3 100644
--- a/tests/test_arrow_writer.py
+++ b/tests/test_arrow_writer.py
@@ -8,7 +8,7 @@
 import pytest
 
 from datasets.arrow_writer import ArrowWriter, OptimizedTypedSequence, TypedSequence
-from datasets.features import Array2DExtensionType
+from datasets.features import Array2D, Array2DExtensionType, ClassLabel, Features, Value
 from datasets.keyhash import DuplicatedKeysError, InvalidKeyError
 
 
@@ -23,38 +23,38 @@ def test_array_type_forbidden(self):
 
     def test_try_type_and_type_forbidden(self):
         with self.assertRaises(ValueError):
-            _ = pa.array(TypedSequence([1, 2, 3], try_type=pa.bool_(), type=pa.int64()))
+            _ = pa.array(TypedSequence([1, 2, 3], try_type=Value("bool"), type=Value("int64")))
 
     def test_compatible_type(self):
-        arr = pa.array(TypedSequence([1, 2, 3], type=pa.int32()))
+        arr = pa.array(TypedSequence([1, 2, 3], type=Value("int32")))
         self.assertEqual(arr.type, pa.int32())
 
     def test_incompatible_type(self):
         with self.assertRaises((TypeError, pa.lib.ArrowInvalid)):
-            _ = pa.array(TypedSequence(["foo", "bar"], type=pa.int64()))
+            _ = pa.array(TypedSequence(["foo", "bar"], type=Value("int64")))
 
     def test_try_compatible_type(self):
-        arr = pa.array(TypedSequence([1, 2, 3], try_type=pa.int32()))
+        arr = pa.array(TypedSequence([1, 2, 3], try_type=Value("int32")))
         self.assertEqual(arr.type, pa.int32())
 
     def test_try_incompatible_type(self):
-        arr = pa.array(TypedSequence(["foo", "bar"], try_type=pa.int64()))
+        arr = pa.array(TypedSequence(["foo", "bar"], try_type=Value("int64")))
         self.assertEqual(arr.type, pa.string())
 
     def test_compatible_extension_type(self):
-        arr = pa.array(TypedSequence([[[1, 2, 3]]], type=Array2DExtensionType((1, 3), "int64")))
+        arr = pa.array(TypedSequence([[[1, 2, 3]]], type=Array2D((1, 3), "int64")))
         self.assertEqual(arr.type, Array2DExtensionType((1, 3), "int64"))
 
     def test_incompatible_extension_type(self):
         with self.assertRaises((TypeError, pa.lib.ArrowInvalid)):
-            _ = pa.array(TypedSequence(["foo", "bar"], type=Array2DExtensionType((1, 3), "int64")))
+            _ = pa.array(TypedSequence(["foo", "bar"], type=Array2D((1, 3), "int64")))
 
     def test_try_compatible_extension_type(self):
-        arr = pa.array(TypedSequence([[[1, 2, 3]]], try_type=Array2DExtensionType((1, 3), "int64")))
+        arr = pa.array(TypedSequence([[[1, 2, 3]]], try_type=Array2D((1, 3), "int64")))
         self.assertEqual(arr.type, Array2DExtensionType((1, 3), "int64"))
 
     def test_try_incompatible_extension_type(self):
-        arr = pa.array(TypedSequence(["foo", "bar"], try_type=Array2DExtensionType((1, 3), "int64")))
+        arr = pa.array(TypedSequence(["foo", "bar"], try_type=Array2D((1, 3), "int64")))
         self.assertEqual(arr.type, pa.string())
 
 
@@ -86,6 +86,27 @@ def test_write(fields, writer_batch_size):
     _check_output(output.getvalue(), expected_num_chunks=num_examples if writer_batch_size == 1 else 1)
 
 
+def test_write_with_features():
+    output = pa.BufferOutputStream()
+    features = Features({"labels": ClassLabel(names=["neg", "pos"])})
+    with ArrowWriter(stream=output, features=features) as writer:
+        writer.write({"labels": 0})
+        writer.write({"labels": 1})
+        num_examples, num_bytes = writer.finalize()
+    assert num_examples == 2
+    assert num_bytes > 0
+    assert writer._schema == features.arrow_schema
+    assert writer._schema.metadata == features.arrow_schema.metadata
+    stream = pa.BufferReader(output.getvalue())
+    f = pa.ipc.open_stream(stream)
+    pa_table: pa.Table = f.read_all()
+    schema = pa_table.schema
+    assert pa_table.num_rows == 2
+    assert schema == features.arrow_schema
+    assert schema.metadata == features.arrow_schema.metadata
+    assert features == Features.from_arrow_schema(schema)
+
+
 @pytest.mark.parametrize("writer_batch_size", [None, 1, 10])
 def test_key_datatype(writer_batch_size):
     output = pa.BufferOutputStream()
@@ -214,7 +235,7 @@ def change_first_primitive_element_in_list(lst, value):
         lst[0] = value
 
 
-@pytest.mark.parametrize("optimized_int_type, expected_dtype", [(None, pa.int64()), (pa.int32(), pa.int32())])
+@pytest.mark.parametrize("optimized_int_type, expected_dtype", [(None, pa.int64()), (Value("int32"), pa.int32())])
 @pytest.mark.parametrize("sequence", [[1, 2, 3], [[1, 2, 3]], [[[1, 2, 3]]]])
 def test_optimized_int_type_for_typed_sequence(sequence, optimized_int_type, expected_dtype):
     arr = pa.array(TypedSequence(sequence, optimized_int_type=optimized_int_type))
diff --git a/tests/test_dataset_cards.py b/tests/test_dataset_cards.py
index be23cce3b57..96cb2d628d6 100644
--- a/tests/test_dataset_cards.py
+++ b/tests/test_dataset_cards.py
@@ -41,7 +41,7 @@ def get_changed_datasets(repo_path: Path) -> List[Path]:
     changed_datasets = set(
         f.resolve().relative_to(datasets_dir_path).parts[0]
         for f in changed_files
-        if f.exists() and str(f.resolve()).startswith(str(datasets_dir_path))
+        if f.exists() and str(f.resolve()).startswith(str(datasets_dir_path)) and not f.name == "dummy_data.zip"
     )
 
     return sorted(dataset_name for dataset_name in changed_datasets if dataset_name not in _PACKAGED_DATASETS_MODULES)
diff --git a/tests/test_dataset_dict.py b/tests/test_dataset_dict.py
index 82d3fc172f7..61524f08826 100644
--- a/tests/test_dataset_dict.py
+++ b/tests/test_dataset_dict.py
@@ -49,7 +49,7 @@ def test_flatten_in_place(self):
         dset = DatasetDict({"train": dset_split, "test": dset_split})
         dset.flatten_()
         self.assertDictEqual(dset.column_names, {"train": ["a.b.c", "foo"], "test": ["a.b.c", "foo"]})
-        self.assertListEqual(list(dset["train"].features.keys()), ["a.b.c", "foo"])
+        self.assertListEqual(sorted(dset["train"].features.keys()), ["a.b.c", "foo"])
         self.assertDictEqual(
             dset["train"].features, Features({"a.b.c": Sequence(Value("string")), "foo": Value("int64")})
         )
@@ -63,7 +63,7 @@ def test_flatten(self):
         dset = DatasetDict({"train": dset_split, "test": dset_split})
         dset = dset.flatten()
         self.assertDictEqual(dset.column_names, {"train": ["a.b.c", "foo"], "test": ["a.b.c", "foo"]})
-        self.assertListEqual(list(dset["train"].features.keys()), ["a.b.c", "foo"])
+        self.assertListEqual(sorted(dset["train"].features.keys()), ["a.b.c", "foo"])
         self.assertDictEqual(
             dset["train"].features, Features({"a.b.c": Sequence(Value("string")), "foo": Value("int64")})
         )
diff --git a/tests/test_table.py b/tests/test_table.py
index 54d2e2ab480..e54f8529960 100644
--- a/tests/test_table.py
+++ b/tests/test_table.py
@@ -6,7 +6,8 @@
 import pyarrow as pa
 import pytest
 
-from datasets import config
+from datasets import Sequence, Value
+from datasets.features.features import ClassLabel, Features
 from datasets.table import (
     ConcatenationTable,
     InMemoryTable,
@@ -17,6 +18,7 @@
     _in_memory_arrow_table_from_file,
     _interpolation_search,
     _memory_mapped_arrow_table_from_file,
+    cast_array_to_feature,
     concat_tables,
     inject_arrow_table_documentation,
 )
@@ -322,6 +324,31 @@ def test_in_memory_table_cast(in_memory_pa_table):
     assert isinstance(table, InMemoryTable)
 
 
+def test_in_memory_table_cast_reorder_struct():
+    table = InMemoryTable(
+        pa.Table.from_pydict(
+            {
+                "top": [
+                    {
+                        "foo": "a",
+                        "bar": "b",
+                    }
+                ]
+            }
+        )
+    )
+    schema = pa.schema({"top": pa.struct({"bar": pa.string(), "foo": pa.string()})})
+    assert table.cast(schema).schema == schema
+
+
+def test_in_memory_table_cast_with_hf_features():
+    table = InMemoryTable(pa.Table.from_pydict({"labels": [0, 1]}))
+    features = Features({"labels": ClassLabel(names=["neg", "pos"])})
+    schema = features.arrow_schema
+    assert table.cast(schema).schema == schema
+    assert Features.from_arrow_schema(table.cast(schema).schema) == features
+
+
 def test_in_memory_table_replace_schema_metadata(in_memory_pa_table):
     metadata = {"huggingface": "{}"}
     table = InMemoryTable(in_memory_pa_table).replace_schema_metadata(metadata)
@@ -769,13 +796,9 @@ def test_concatenation_table_cast(
             for k, v in zip(in_memory_pa_table.schema.names, in_memory_pa_table.schema.types)
         }
     )
-    if config.PYARROW_VERSION.major < 4:
-        with pytest.raises(pa.ArrowNotImplementedError):
-            ConcatenationTable.from_blocks(blocks).cast(schema)
-    else:
-        table = ConcatenationTable.from_blocks(blocks).cast(schema)
-        assert table.table == in_memory_pa_table.cast(schema)
-        assert isinstance(table, ConcatenationTable)
+    table = ConcatenationTable.from_blocks(blocks).cast(schema)
+    assert table.table == in_memory_pa_table.cast(schema)
+    assert isinstance(table, ConcatenationTable)
     schema = pa.schema(
         {
             k: v if v != pa.int64() else pa.int32()
@@ -980,3 +1003,17 @@ def test_indexed_table_mixin():
     assert all(table._offsets.tolist() == np.cumsum([0] + [n_rows_per_chunk] * n_chunks))
     assert table.fast_slice(5) == pa_table.slice(5)
     assert table.fast_slice(2, 13) == pa_table.slice(2, 13)
+
+
+def test_cast_array_to_features():
+    arr = pa.array([[0, 1]])
+    assert cast_array_to_feature(arr, Sequence(Value("string"))).type == pa.list_(pa.string())
+    with pytest.raises(TypeError):
+        cast_array_to_feature(arr, Sequence(Value("string")), allow_number_to_str=False)
+
+
+def test_cast_array_to_features_nested():
+    arr = pa.array([[{"foo": [0]}]])
+    assert cast_array_to_feature(arr, [{"foo": Sequence(Value("string"))}]).type == pa.list_(
+        pa.struct({"foo": pa.list_(pa.string())})
+    )


Code review comments:
diff_hunk: @@ -226,12 +92,14 @@ def decode_example(self, value):
             raise ImportError("To support decoding images, please install 'Pillow'.")
 
         if isinstance(value, str):
-            path, bytes_ = value, None
+            path, bytes_ = value, None  # TODO(QL, Mario): remove this case ?
path: src/datasets/features/image.py
text: @user1:
Yes, let's remove this because it's no longer needed. Previously, we would need this for the hacky `cast_column` cast from `Value("string")` to `Image()`:
```python
ds = Dataset.from_dict({"image": ["path/to/image1", "path/to/image2"]})
ds = ds.cast_column("image", Image())
```