Embed image /audio data in dl_and_prepare parquet

src/datasets/arrow_writer.py

@@ -38,7 +38,7 @@
 from .filesystems import is_remote_filesystem
 from .info import DatasetInfo
 from .keyhash import DuplicatedKeysError, KeyHasher
-from .table import array_cast, cast_array_to_feature, table_cast
+from .table import array_cast, cast_array_to_feature, embed_table_storage, table_cast
 from .utils import logging
 from .utils.file_utils import hash_url_to_filename
 from .utils.py_utils import asdict, first_non_null_value
@@ -287,6 +287,7 @@ def __init__(
         update_features: bool = False,
         with_metadata: bool = True,
         unit: str = "examples",
+        embed_local_files: bool = False,
         storage_options: Optional[dict] = None,
     ):
         if path is None and stream is None:
@@ -332,6 +333,7 @@ def __init__(
         self.update_features = update_features
         self.with_metadata = with_metadata
         self.unit = unit
+        self.embed_local_files = embed_local_files
 
         self._num_examples = 0
         self._num_bytes = 0
@@ -536,6 +538,8 @@ def write_table(self, pa_table: pa.Table, writer_batch_size: Optional[int] = Non
         if self.pa_writer is None:
             self._build_writer(inferred_schema=pa_table.schema)
         pa_table = table_cast(pa_table, self._schema)
+        if self.embed_local_files:
+            pa_table = embed_table_storage(pa_table)
         self._num_bytes += pa_table.nbytes
         self._num_examples += pa_table.num_rows
         self.pa_writer.write_table(pa_table, writer_batch_size)

src/datasets/builder.py

@@ -619,6 +619,7 @@ def download_and_prepare(
                 If True, will get token from ~/.huggingface.
             file_format (:obj:`str`, optional): format of the data files in which the dataset will be written.
                 Supported formats: "arrow", "parquet". Default to "arrow" format.
+                If the format is "parquet", then image and audio data are embedded into the Parquet files instead of pointing to local files.
 
                 <Added version="2.5.0"/>
             max_shard_size (:obj:`Union[str, int]`, optional): Maximum number of bytes written per shard.
@@ -1348,6 +1349,7 @@ def _prepare_split(
         generator = self._generate_examples(**split_generator.gen_kwargs)
 
         writer_class = ParquetWriter if file_format == "parquet" else ArrowWriter
+        embed_local_files = file_format == "parquet"
 
         shard_id = 0
         # TODO: embed the images/audio files inside parquet files.
@@ -1358,6 +1360,7 @@ def _prepare_split(
             hash_salt=split_info.name,
             check_duplicates=check_duplicate_keys,
             storage_options=self._fs.storage_options,
+            embed_local_files=embed_local_files,
         )
         total_num_examples, total_num_bytes = 0, 0
         try:
@@ -1381,6 +1384,7 @@ def _prepare_split(
                         hash_salt=split_info.name,
                         check_duplicates=check_duplicate_keys,
                         storage_options=self._fs.storage_options,
+                        embed_local_files=embed_local_files,
                     )
                 example = self.info.features.encode_example(record)
                 writer.write(example, key)
@@ -1474,13 +1478,15 @@ def _prepare_split(
         generator = self._generate_tables(**split_generator.gen_kwargs)
 
         writer_class = ParquetWriter if file_format == "parquet" else ArrowWriter
+        embed_local_files = file_format == "parquet"
 
         shard_id = 0
         # TODO: embed the images/audio files inside parquet files.
         writer = writer_class(
             features=self.info.features,
             path=fpath.replace("SSSSS", f"{shard_id:05d}"),
             storage_options=self._fs.storage_options,
+            embed_local_files=embed_local_files,
         )
         total_num_examples, total_num_bytes = 0, 0
         try:
@@ -1499,6 +1505,7 @@ def _prepare_split(
                         features=writer._features,
                         path=fpath.replace("SSSSS", f"{shard_id:05d}"),
                         storage_options=self._fs.storage_options,
+                        embed_local_files=embed_local_files,
                     )
                 writer.write_table(table)
         finally:

tests/test_arrow_writer.py

@@ -326,3 +326,28 @@ def test_parquet_writer_write():
     stream = pa.BufferReader(output.getvalue())
     pa_table: pa.Table = pq.read_table(stream)
     assert pa_table.to_pydict() == {"col_1": ["foo", "bar"], "col_2": [1, 2]}
+
+
+@require_pil
+@pytest.mark.parametrize("embed_local_files", [False, True])
+def test_writer_embed_local_files(tmp_path, embed_local_files):
+    import PIL.Image
+
+    image_path = str(tmp_path / "test_image_rgb.jpg")
+    PIL.Image.fromarray(np.zeros((5, 5), dtype=np.uint8)).save(image_path, format="png")
+    output = pa.BufferOutputStream()
+    with ParquetWriter(
+        stream=output, features=Features({"image": Image()}), embed_local_files=embed_local_files
+    ) as writer:
+        writer.write({"image": image_path})
+        writer.finalize()
+    stream = pa.BufferReader(output.getvalue())
+    pa_table: pa.Table = pq.read_table(stream)
+    out = pa_table.to_pydict()
+    if embed_local_files:
+        assert out["image"][0]["path"] is None
+        with open(image_path, "rb") as f:
+            assert out["image"][0]["bytes"] == f.read()
+    else:
+        assert out["image"][0]["path"] == image_path
+        assert out["image"][0]["bytes"] is None