push_to_hub() for videos

src/datasets/arrow_dataset.py

@@ -5773,13 +5773,6 @@ def push_to_hub(
         >>> french_dataset = load_dataset("<organization>/<dataset_id>", "fr")
         ```
         """
-        if "Video(" in str(self.features):
-            raise NotImplementedError(
-                "push_to_hub is not implemented for video datasets, instead you should upload the video files "
-                "using e.g. the huggingface_hub library and optionally upload a metadata.csv or metadata.jsonl "
-                "file containing other information like video captions, features or labels. More information "
-                "at https://huggingface.co/docs/datasets/main/en/video_load#videofolder"
-            )
         if config_name == "data":
             raise ValueError("`config_name` cannot be 'data'. Please, choose another name for configuration.")
 

src/datasets/arrow_writer.py

@@ -36,6 +36,7 @@
     get_nested_type,
     list_of_np_array_to_pyarrow_listarray,
     numpy_to_pyarrow_listarray,
+    require_storage_embed,
     to_pyarrow_listarray,
 )
 from .filesystems import is_remote_filesystem
@@ -150,7 +151,7 @@ def get_writer_batch_size_from_data_size(num_rows: int, num_bytes: int) -> int:
         writer_batch_size (`Optional[int]`):
             Writer batch size to pass to a parquet writer.
     """
-    return max(10, num_rows * convert_file_size_to_int(config.MAX_ROW_GROUP_SIZE) // num_bytes) if num_bytes > 0 else 1
+    return max(1, num_rows * convert_file_size_to_int(config.MAX_ROW_GROUP_SIZE) // num_bytes) if num_bytes > 0 else 1
 
 
 class SchemaInferenceError(ValueError):
@@ -715,6 +716,13 @@ def _build_writer(self, inferred_schema: pa.Schema):
             self._schema,
             use_content_defined_chunking=self.use_content_defined_chunking,
             write_page_index=self.write_page_index,
+            compression={
+                col: "none" if require_storage_embed(feature) else "snappy" for col, feature in self._features.items()
+            },
+            use_dictionary=[col for col, feature in self._features.items() if not require_storage_embed(feature)],
+            column_encoding={
+                col: "PLAIN" for col, feature in self._features.items() if require_storage_embed(feature)
+            },
         )
         if self.use_content_defined_chunking is not False:
             self.pa_writer.add_key_value_metadata(

src/datasets/features/video.py

@@ -10,7 +10,7 @@
 from ..download.download_config import DownloadConfig
 from ..table import array_cast
 from ..utils.file_utils import is_local_path, xopen
-from ..utils.py_utils import string_to_dict
+from ..utils.py_utils import no_op_if_value_is_null, string_to_dict
 
 
 if TYPE_CHECKING:
@@ -288,6 +288,46 @@ def cast_storage(self, storage: Union[pa.StringArray, pa.StructArray, pa.ListArr
             )
         return array_cast(storage, self.pa_type)
 
+    def embed_storage(self, storage: pa.StructArray, token_per_repo_id=None) -> pa.StructArray:
+        """Embed image files into the Arrow array.
+
+        Args:
+            storage (`pa.StructArray`):
+                PyArrow array to embed.
+
+        Returns:
+            `pa.StructArray`: Array in the Video arrow storage type, that is
+                `pa.struct({"bytes": pa.binary(), "path": pa.string()})`.
+        """
+        if token_per_repo_id is None:
+            token_per_repo_id = {}
+
+        @no_op_if_value_is_null
+        def path_to_bytes(path):
+            source_url = path.split("::")[-1]
+            pattern = (
+                config.HUB_DATASETS_URL if source_url.startswith(config.HF_ENDPOINT) else config.HUB_DATASETS_HFFS_URL
+            )
+            source_url_fields = string_to_dict(source_url, pattern)
+            token = token_per_repo_id.get(source_url_fields["repo_id"]) if source_url_fields is not None else None
+            download_config = DownloadConfig(token=token)
+            with xopen(path, "rb", download_config=download_config) as f:
+                return f.read()
+
+        bytes_array = pa.array(
+            [
+                (path_to_bytes(x["path"]) if x["bytes"] is None else x["bytes"]) if x is not None else None
+                for x in storage.to_pylist()
+            ],
+            type=pa.binary(),
+        )
+        path_array = pa.array(
+            [os.path.basename(path) if path is not None else None for path in storage.field("path").to_pylist()],
+            type=pa.string(),
+        )
+        storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=bytes_array.is_null())
+        return array_cast(storage, self.pa_type)
+
 
 def video_to_bytes(video: "VideoDecoder") -> bytes:
     """Convert a torchcodec Video object to bytes using native compression if possible"""

src/datasets/iterable_dataset.py

@@ -4189,13 +4189,6 @@ def push_to_hub(
         >>> french_dataset = load_dataset("<organization>/<dataset_id>", "fr")
         ```
         """
-        if "Video(" in str(self.features):
-            raise NotImplementedError(
-                "push_to_hub is not implemented for video datasets, instead you should upload the video files "
-                "using e.g. the huggingface_hub library and optionally upload a metadata.csv or metadata.jsonl "
-                "file containing other information like video captions, features or labels. More information "
-                "at https://huggingface.co/docs/datasets/main/en/video_load#videofolder"
-            )
         if num_proc is not None and num_proc > self.num_shards:
             logger.warning(
                 f"Too many num_proc: {num_proc} (max is dataset.num_shards={self.num_shards}). "