From 6831a80e8332820184622e2709ee9f3602583644 Mon Sep 17 00:00:00 2001 From: Quentin Lhoest Date: Thu, 29 Jan 2026 19:14:47 +0100 Subject: [PATCH 1/2] push_to_hub() for videos --- src/datasets/arrow_dataset.py | 7 ------ src/datasets/arrow_writer.py | 10 +++++++- src/datasets/features/video.py | 42 +++++++++++++++++++++++++++++++- src/datasets/iterable_dataset.py | 7 ------ 4 files changed, 50 insertions(+), 16 deletions(-) diff --git a/src/datasets/arrow_dataset.py b/src/datasets/arrow_dataset.py index ee8e8784a3e..bfffff9f238 100644 --- a/src/datasets/arrow_dataset.py +++ b/src/datasets/arrow_dataset.py @@ -5773,13 +5773,6 @@ def push_to_hub( >>> french_dataset = load_dataset("/", "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.") diff --git a/src/datasets/arrow_writer.py b/src/datasets/arrow_writer.py index e6735058d73..da6c41b55b3 100644 --- a/src/datasets/arrow_writer.py +++ b/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" if require_storage_embed(feature) else None for col, feature in self._features.items() + }, ) if self.use_content_defined_chunking is not False: self.pa_writer.add_key_value_metadata( diff --git a/src/datasets/features/video.py b/src/datasets/features/video.py index d65855a4dc8..cf1c19551ca 100644 --- a/src/datasets/features/video.py +++ b/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""" diff --git a/src/datasets/iterable_dataset.py b/src/datasets/iterable_dataset.py index b66c8977ade..6532c04dd57 100644 --- a/src/datasets/iterable_dataset.py +++ b/src/datasets/iterable_dataset.py @@ -4189,13 +4189,6 @@ def push_to_hub( >>> french_dataset = load_dataset("/", "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}). " From ddf50e608d43d35af2b39ce238dd1fccbf9c8779 Mon Sep 17 00:00:00 2001 From: Quentin Lhoest Date: Thu, 29 Jan 2026 19:45:12 +0100 Subject: [PATCH 2/2] fix tests --- src/datasets/arrow_writer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/datasets/arrow_writer.py b/src/datasets/arrow_writer.py index da6c41b55b3..1d2eaf5c9ac 100644 --- a/src/datasets/arrow_writer.py +++ b/src/datasets/arrow_writer.py @@ -721,7 +721,7 @@ def _build_writer(self, inferred_schema: pa.Schema): }, use_dictionary=[col for col, feature in self._features.items() if not require_storage_embed(feature)], column_encoding={ - col: "plain" if require_storage_embed(feature) else None for col, feature in self._features.items() + col: "PLAIN" for col, feature in self._features.items() if require_storage_embed(feature) }, ) if self.use_content_defined_chunking is not False: