Repository: huggingface/datasets
Pull request: Reload features from Parquet metadata
Resolves #5482. 

Attaches feature metadata to parquet files serialised using `Dataset.to_parquet`. 

This allows retrieving data with "rich" feature types (e.g., `datasets.features.image.Image` or `datasets.features.audio.Audio`) from parquet files without cumbersome casting (for an example, see #5482).

@lhoestq It seems that it is sufficient to attach metadata to the schema prior to serialising and features are loaded back with correct types afterwards automatically. 

I used the following script to test the implementation:

```python
from pathlib import Path

import datasets

dataset_name = "Maysee/tiny-imagenet"
ds = datasets.load_dataset(dataset_name, split=datasets.Split.TRAIN)
output_directory_path = Path(__file__).parent.joinpath("example_test_outputs", dataset_name.replace("/", "_"))
output_directory_path.mkdir(exist_ok=True, parents=True)
output_filepath = output_directory_path.joinpath("ds.parquet")
ds.to_parquet(str(output_filepath))

reloaded_ds = datasets.load_dataset(str(output_directory_path), split=datasets.Split.TRAIN)

assert ds.features == reloaded_ds.features
```

Prior to the change in this PR this script raises an `AssertionError` and the `Image` features lose their type after serialisation. After the change in this PR, the assertion does not raise an error and manual inspection of the features shows type `Image` for the respective columns of `reloaded_ds `.

Some open questions:
* How/where can I best add new unit tests for this implementation?
* What dataset would I best use in the tests? I chose `Maysee/tiny-imagenet` mainly because it is small and contains an ?Image` feature that can be used to test, but I'd be happy for suggestions on a suitable data source to use.
* Currently I'm calling `datasets.arrow_writer.ArrowWriter._build_metadata` as I need the same logic. However, I'm not happy with the coupling between `datasets.io.parquet` and `datasets.arrow_writer` it leaves me with. Suggest to factor this common logic out into a helper function and reuse it from both of these. Do you agree and if yes, could you please guide me where I would best place this function?

Many thanks in advance and kind regards,
MFreidank


  

Problem statement:
Reload features from Parquet metadata
The idea would be to allow this :
```python 
ds.to_parquet("my_dataset/ds.parquet")
reloaded = load_dataset("my_dataset")
assert ds.features == reloaded.features 
```
And it should also work with Image and Audio types (right now they're reloaded as a dict type)

This can be implemented by storing and reading the feature types in the parquet metadata, as we do for arrow files.

Hints from issue discussion:
I'd be happy to have a look, if nobody else has started working on this yet @lhoestq. 

It seems to me that for the `arrow` format features are currently attached as metadata [in `datasets.arrow_writer`](https://github.com/huggingface/datasets/blob/5f810b7011a8a4ab077a1847c024d2d9e267b065/src/datasets/arrow_writer.py#L412)  and retrieved from the metadata at `load_dataset` time using [`datasets.features.features.from_arrow_schema`](https://github.com/huggingface/datasets/blob/5f810b7011a8a4ab077a1847c024d2d9e267b065/src/datasets/features/features.py#L1602). 

This will need to be replicated for `parquet` via calls to [this api](https://arrow.apache.org/docs/python/generated/pyarrow.parquet.write_metadata.html) from `io.parquet.ParquetWriter` and `io.parquet.ParquetReader` [respectively](https://github.com/huggingface/datasets/blob/5f810b7011a8a4ab077a1847c024d2d9e267b065/src/datasets/io/parquet.py#L104).

Any other important considerations?

Thanks @MFreidank ! That's correct :)

Reading the metadata to infer the features can be ideally done in the `parquet.py` file in `packaged_builder` when a parquet file is read. You can cast the arrow table to the schema you get from the features.arrow_schema
#self-assign

Commit message: removed unused import

Patch to review:
diff --git a/src/datasets/io/parquet.py b/src/datasets/io/parquet.py
index 357c77e9105..b5e62a09b8c 100644
--- a/src/datasets/io/parquet.py
+++ b/src/datasets/io/parquet.py
@@ -1,7 +1,6 @@
 import os
 from typing import BinaryIO, Optional, Union
 
-import pyarrow as pa
 import pyarrow.parquet as pq
 
 from .. import Dataset, Features, NamedSplit, config
@@ -100,7 +99,8 @@ def _write(self, file_obj: BinaryIO, batch_size: int, **parquet_writer_kwargs) -
         """
         written = 0
         _ = parquet_writer_kwargs.pop("path_or_buf", None)
-        schema = pa.schema(self.dataset.features.type)
+        schema = self.dataset.features.arrow_schema
+
         writer = pq.ParquetWriter(file_obj, schema=schema, **parquet_writer_kwargs)
 
         for offset in logging.tqdm(
diff --git a/tests/io/test_parquet.py b/tests/io/test_parquet.py
index 1c95f00180d..0c317bb68cb 100644
--- a/tests/io/test_parquet.py
+++ b/tests/io/test_parquet.py
@@ -1,10 +1,12 @@
+import numpy as np
 import pyarrow.parquet as pq
 import pytest
 
 from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
+from datasets.features.image import Image
 from datasets.io.parquet import ParquetDatasetReader, ParquetDatasetWriter
 
-from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases
+from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases, require_pil
 
 
 def _check_parquet_dataset(dataset, expected_features):
@@ -130,3 +132,20 @@ def test_parquer_write(dataset, tmp_path):
     pf = pq.ParquetFile(tmp_path / "foo.parquet")
     output_table = pf.read()
     assert dataset.data.table == output_table
+
+
+@require_pil
+def test_dataset_to_parquet_keeps_features(shared_datadir, tmp_path):
+    import PIL.Image
+
+    image_path = str(shared_datadir / "test_image_rgb.jpg")
+    PIL.Image.fromarray(np.zeros((5, 5), dtype=np.uint8)).save(image_path, format="png")
+    data = {"image": [image_path]}
+    features = Features({"image": Image()})
+    dataset = Dataset.from_dict(data, features=features)
+    writer = ParquetDatasetWriter(dataset, tmp_path / "foo.parquet")
+    assert writer.write() > 0
+
+    reloaded_dataset = Dataset.from_parquet(str(tmp_path / "foo.parquet"))
+
+    assert dataset.features == reloaded_dataset.features


Code review comments:
diff_hunk: @@ -130,3 +132,20 @@ def test_parquer_write(dataset, tmp_path):
     pf = pq.ParquetFile(tmp_path / "foo.parquet")
     output_table = pf.read()
     assert dataset.data.table == output_table
+
+
+@require_pil
+def test_dataset_to_parquet_keeps_features(shared_datadir, tmp_path):
+    import PIL.Image
+
+    image_path = str(shared_datadir / "test_image_rgb.jpg")
+    PIL.Image.fromarray(np.zeros((5, 5), dtype=np.uint8)).save(image_path, format="png")
path: tests/io/test_parquet.py
text: @user1:
Oh actually you can also remove this line, the file already exists and is a dummy image we use for tests 

@author:
Done! The test actually does not use `PIL` anymore after that so I also dropped the `require_pil`.

@user1:
`require_pil` is actually required to use image datasets, so you can re-add it ^^

@author:
re-added ^^ sorry about that