add to_parquet and push_to_hub

Author: lhoestq

src/datasets/features/audio.py

@@ -241,7 +241,7 @@ def cast_storage(self, storage: Union[pa.StringArray, pa.StructArray]) -> pa.Str
             storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=storage.is_null())
         return array_cast(storage, self.pa_type)
 
-    def embed_storage(self, storage: pa.StructArray) -> pa.StructArray:
+    def embed_storage(self, storage: pa.StructArray, token_per_repo_id=None) -> pa.StructArray:
         """Embed audio files into the Arrow array.
 
         Args:
@@ -252,12 +252,24 @@ def embed_storage(self, storage: pa.StructArray) -> pa.StructArray:
             `pa.StructArray`: Array in the Audio 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):
-            with xopen(path, "rb") as f:
-                bytes_ = f.read()
-            return bytes_
+            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(
             [

src/datasets/features/image.py

@@ -250,7 +250,7 @@ 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) -> pa.StructArray:
+    def embed_storage(self, storage: pa.StructArray, token_per_repo_id=None) -> pa.StructArray:
         """Embed image files into the Arrow array.
 
         Args:
@@ -261,12 +261,24 @@ def embed_storage(self, storage: pa.StructArray) -> pa.StructArray:
             `pa.StructArray`: Array in the Image 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):
-            with xopen(path, "rb") as f:
-                bytes_ = f.read()
-            return bytes_
+            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(
             [

src/datasets/features/pdf.py

@@ -216,7 +216,7 @@ def cast_storage(self, storage: Union[pa.StringArray, pa.StructArray, pa.ListArr
             storage = pa.StructArray.from_arrays([bytes_array, path_array], ["bytes", "path"], mask=storage.is_null())
         return array_cast(storage, self.pa_type)
 
-    def embed_storage(self, storage: pa.StructArray) -> pa.StructArray:
+    def embed_storage(self, storage: pa.StructArray, token_per_repo_id=None) -> pa.StructArray:
         """Embed PDF files into the Arrow array.
 
         Args:
@@ -227,12 +227,24 @@ def embed_storage(self, storage: pa.StructArray) -> pa.StructArray:
             `pa.StructArray`: Array in the PDF 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):
-            with xopen(path, "rb") as f:
-                bytes_ = f.read()
-            return bytes_
+            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(
             [