Add skip_origin_metadata flag to optimize streaming startup time

src/datasets/builder.py

@@ -200,10 +200,17 @@ def create_config_id(
         else:
             return self.name
 
-    def _resolve_data_files(self, base_path: str, download_config: DownloadConfig) -> None:
+    def _resolve_data_files(
+        self,
+        base_path: str,
+        download_config: DownloadConfig,
+        skip_origin_metadata: bool = False,
+    ) -> None:
         if isinstance(self.data_files, DataFilesPatternsDict):
             base_path = xjoin(base_path, self.data_dir) if self.data_dir else base_path
-            self.data_files = self.data_files.resolve(base_path, download_config)
+            self.data_files = self.data_files.resolve(
+                base_path, download_config, skip_origin_metadata=skip_origin_metadata
+            )
 
 
 class DatasetBuilder:
@@ -274,6 +281,9 @@ class DatasetBuilder:
             It defines the number of samples that are kept in memory before writing them
             and also the length of the arrow chunks.
             None means that the ArrowWriter will use its default value.
+        skip_origin_metadata (`bool`, *optional*):
+            Whether to skip fetching origin metadata (ETag/mtime) for the data files.
+            It should only be true when streaming is true.
         **config_kwargs (additional keyword arguments): Keyword arguments to be passed to the corresponding builder
             configuration class, set on the class attribute [`DatasetBuilder.BUILDER_CONFIG_CLASS`]. The builder
             configuration class is [`BuilderConfig`] or a subclass of it.
@@ -313,6 +323,7 @@ def __init__(
         storage_options: Optional[dict] = None,
         writer_batch_size: Optional[int] = None,
         config_id: Optional[str] = None,
+        skip_origin_metadata: bool = False,
         **config_kwargs,
     ):
         # DatasetBuilder name
@@ -324,12 +335,14 @@ def __init__(
         self.storage_options = storage_options or {}
         self.dataset_name = camelcase_to_snakecase(dataset_name) if dataset_name else self.name
         self._writer_batch_size = writer_batch_size or self.DEFAULT_WRITER_BATCH_SIZE
+        self._skip_origin_metadata = skip_origin_metadata
 
         if data_files is not None and not isinstance(data_files, DataFilesDict):
             data_files = DataFilesDict.from_patterns(
                 sanitize_patterns(data_files),
                 base_path=base_path,
                 download_config=DownloadConfig(token=token, storage_options=self.storage_options),
+                skip_origin_metadata=skip_origin_metadata,
             )
 
         # Prepare config: DatasetConfig contains name, version and description but can be extended by each dataset
@@ -558,6 +571,7 @@ def _create_builder_config(
         builder_config._resolve_data_files(
             base_path=self.base_path,
             download_config=DownloadConfig(token=self.token, storage_options=self.storage_options),
+            skip_origin_metadata=self._skip_origin_metadata,
         )
 
         # compute the config id that is going to be used for caching
@@ -765,6 +779,9 @@ def download_and_prepare(
         >>> builder.download_and_prepare("s3://my-bucket/my_rotten_tomatoes", storage_options=storage_options, file_format="parquet")
         ```
         """
+        if getattr(self, "_skip_origin_metadata", False):
+            raise ValueError("This function is not intended for streaming.")
+
         output_dir = output_dir if output_dir is not None else self._cache_dir
         # output_dir can be a remote bucket on GCS or S3
         fs, output_dir = url_to_fs(output_dir, **(storage_options or {}))

src/datasets/data_files.py

@@ -581,10 +581,11 @@ def from_hf_repo(
         base_path: Optional[str] = None,
         allowed_extensions: Optional[list[str]] = None,
         download_config: Optional[DownloadConfig] = None,
+        skip_origin_metadata: bool = False,
     ) -> "DataFilesList":
         base_path = f"hf://datasets/{dataset_info.id}@{dataset_info.sha}/{base_path or ''}".rstrip("/")
         return cls.from_patterns(
-            patterns, base_path=base_path, allowed_extensions=allowed_extensions, download_config=download_config
+            patterns, base_path=base_path, allowed_extensions=allowed_extensions, download_config=download_config, skip_origin_metadata=skip_origin_metadata
         )
 
     @classmethod
@@ -594,10 +595,11 @@ def from_local_or_remote(
         base_path: Optional[str] = None,
         allowed_extensions: Optional[list[str]] = None,
         download_config: Optional[DownloadConfig] = None,
+        skip_origin_metadata: bool = False,
     ) -> "DataFilesList":
         base_path = base_path if base_path is not None else Path().resolve().as_posix()
         return cls.from_patterns(
-            patterns, base_path=base_path, allowed_extensions=allowed_extensions, download_config=download_config
+            patterns, base_path=base_path, allowed_extensions=allowed_extensions, download_config=download_config, skip_origin_metadata=skip_origin_metadata
         )
 
     @classmethod
@@ -607,6 +609,7 @@ def from_patterns(
         base_path: Optional[str] = None,
         allowed_extensions: Optional[list[str]] = None,
         download_config: Optional[DownloadConfig] = None,
+        skip_origin_metadata: bool = False,
     ) -> "DataFilesList":
         base_path = base_path if base_path is not None else Path().resolve().as_posix()
         data_files = []
@@ -623,7 +626,10 @@ def from_patterns(
             except FileNotFoundError:
                 if not has_magic(pattern):
                     raise
-        origin_metadata = _get_origin_metadata(data_files, download_config=download_config)
+        if not skip_origin_metadata:
+            origin_metadata = _get_origin_metadata(data_files, download_config=download_config)
+        else:
+            origin_metadata = [()] * len(data_files)
         return cls(data_files, origin_metadata)
 
     def filter(
@@ -668,6 +674,7 @@ def from_local_or_remote(
         base_path: Optional[str] = None,
         allowed_extensions: Optional[list[str]] = None,
         download_config: Optional[DownloadConfig] = None,
+        skip_origin_metadata: bool = False,
     ) -> "DataFilesDict":
         out = cls()
         for key, patterns_for_key in patterns.items():
@@ -679,6 +686,7 @@ def from_local_or_remote(
                     base_path=base_path,
                     allowed_extensions=allowed_extensions,
                     download_config=download_config,
+                    skip_origin_metadata=skip_origin_metadata,
                 )
             )
         return out
@@ -691,6 +699,7 @@ def from_hf_repo(
         base_path: Optional[str] = None,
         allowed_extensions: Optional[list[str]] = None,
         download_config: Optional[DownloadConfig] = None,
+        skip_origin_metadata: bool = False,
     ) -> "DataFilesDict":
         out = cls()
         for key, patterns_for_key in patterns.items():
@@ -703,6 +712,7 @@ def from_hf_repo(
                     base_path=base_path,
                     allowed_extensions=allowed_extensions,
                     download_config=download_config,
+                    skip_origin_metadata=skip_origin_metadata,
                 )
             )
         return out
@@ -714,6 +724,7 @@ def from_patterns(
         base_path: Optional[str] = None,
         allowed_extensions: Optional[list[str]] = None,
         download_config: Optional[DownloadConfig] = None,
+        skip_origin_metadata: bool = False,
     ) -> "DataFilesDict":
         out = cls()
         for key, patterns_for_key in patterns.items():
@@ -725,6 +736,7 @@ def from_patterns(
                     base_path=base_path,
                     allowed_extensions=allowed_extensions,
                     download_config=download_config,
+                    skip_origin_metadata=skip_origin_metadata,
                 )
             )
         return out
@@ -766,6 +778,7 @@ def resolve(
         self,
         base_path: str,
         download_config: Optional[DownloadConfig] = None,
+        skip_origin_metadata: bool = False,
     ) -> "DataFilesList":
         base_path = base_path if base_path is not None else Path().resolve().as_posix()
         data_files = []
@@ -782,7 +795,10 @@ def resolve(
             except FileNotFoundError:
                 if not has_magic(pattern):
                     raise
-        origin_metadata = _get_origin_metadata(data_files, download_config=download_config)
+        if not skip_origin_metadata:
+            origin_metadata = _get_origin_metadata(data_files, download_config=download_config)
+        else:
+            origin_metadata = [()] * len(data_files)
         return DataFilesList(data_files, origin_metadata)
 
     def filter_extensions(self, extensions: list[str]) -> "DataFilesPatternsList":
@@ -816,10 +832,11 @@ def resolve(
         self,
         base_path: str,
         download_config: Optional[DownloadConfig] = None,
+        skip_origin_metadata: bool = False,
     ) -> "DataFilesDict":
         out = DataFilesDict()
         for key, data_files_patterns_list in self.items():
-            out[key] = data_files_patterns_list.resolve(base_path, download_config)
+            out[key] = data_files_patterns_list.resolve(base_path, download_config, skip_origin_metadata=skip_origin_metadata)
         return out
 
     def filter_extensions(self, extensions: list[str]) -> "DataFilesPatternsDict":

src/datasets/load.py

@@ -406,6 +406,7 @@ def __init__(
         data_dir: Optional[str] = None,
         data_files: Optional[Union[str, list, dict]] = None,
         download_mode: Optional[Union[DownloadMode, str]] = None,
+        skip_origin_metadata: bool = False,
     ):
         if data_dir and os.path.isabs(data_dir):
             raise ValueError(f"`data_dir` must be relative to a dataset directory's root: {path}")
@@ -415,6 +416,7 @@ def __init__(
         self.data_files = data_files
         self.data_dir = data_dir
         self.download_mode = download_mode
+        self.skip_origin_metadata = skip_origin_metadata
 
     def get_module(self) -> DatasetModule:
         readme_path = os.path.join(self.path, config.REPOCARD_FILENAME)
@@ -442,6 +444,7 @@ def get_module(self) -> DatasetModule:
             patterns,
             base_path=base_path,
             allowed_extensions=_ALL_ALLOWED_EXTENSIONS,
+            skip_origin_metadata=self.skip_origin_metadata,
         )
         module_name, default_builder_kwargs = infer_module_for_data_files(
             data_files=data_files,
@@ -515,12 +518,14 @@ def __init__(
         data_files: Optional[Union[str, list, dict]] = None,
         download_config: Optional[DownloadConfig] = None,
         download_mode: Optional[Union[DownloadMode, str]] = None,
+        skip_origin_metadata: bool = False,
     ):
         self.name = name
         self.data_files = data_files
         self.data_dir = data_dir
         self.download_config = download_config
         self.download_mode = download_mode
+        self.skip_origin_metadata = skip_origin_metadata
         increase_load_count(name)
 
     def get_module(self) -> DatasetModule:
@@ -534,6 +539,7 @@ def get_module(self) -> DatasetModule:
             patterns,
             download_config=self.download_config,
             base_path=base_path,
+            skip_origin_metadata=self.skip_origin_metadata,
         )
 
         module_path, hash = _PACKAGED_DATASETS_MODULES[self.name]
@@ -561,6 +567,7 @@ def __init__(
         download_config: Optional[DownloadConfig] = None,
         download_mode: Optional[Union[DownloadMode, str]] = None,
         use_exported_dataset_infos: bool = False,
+        skip_origin_metadata: bool = False,
     ):
         self.name = name
         self.commit_hash = commit_hash
@@ -569,6 +576,7 @@ def __init__(
         self.download_config = download_config or DownloadConfig()
         self.download_mode = download_mode
         self.use_exported_dataset_infos = use_exported_dataset_infos
+        self.skip_origin_metadata = skip_origin_metadata
         increase_load_count(name)
 
     def get_module(self) -> DatasetModule:
@@ -641,6 +649,7 @@ def get_module(self) -> DatasetModule:
             base_path=base_path,
             allowed_extensions=_ALL_ALLOWED_EXTENSIONS,
             download_config=self.download_config,
+            skip_origin_metadata=self.skip_origin_metadata,
         )
         module_name, default_builder_kwargs = infer_module_for_data_files(
             data_files=data_files,
@@ -844,13 +853,15 @@ def __init__(
         data_files: Optional[Union[str, list, dict]] = None,
         download_config: Optional[DownloadConfig] = None,
         download_mode: Optional[Union[DownloadMode, str]] = None,
+        skip_origin_metadata: bool = False,
     ):
         self.path = Path(path).as_posix()
         self.name = Path(path).stem
         self.data_files = data_files
         self.data_dir = data_dir
         self.download_config = download_config
         self.download_mode = download_mode
+        self.skip_origin_metadata = skip_origin_metadata
 
     def get_module(self) -> DatasetModule:
         hffs = HfFileSystem(
@@ -889,6 +900,7 @@ def get_module(self) -> DatasetModule:
             patterns,
             base_path=base_path,
             allowed_extensions=_ALL_ALLOWED_EXTENSIONS,
+            skip_origin_metadata=self.skip_origin_metadata,
         )
         module_name, default_builder_kwargs = infer_module_for_data_files(
             data_files=data_files,
@@ -960,6 +972,7 @@ def dataset_module_factory(
     data_dir: Optional[str] = None,
     data_files: Optional[Union[dict, list, str, DataFilesDict]] = None,
     cache_dir: Optional[str] = None,
+    skip_origin_metadata: bool = False,
     **download_kwargs,
 ) -> DatasetModule:
     """
@@ -1052,6 +1065,7 @@ def dataset_module_factory(
             data_files=data_files,
             download_config=download_config,
             download_mode=download_mode,
+            skip_origin_metadata=skip_origin_metadata,
         ).get_module()
     # Try locally
     elif path.endswith(filename):
@@ -1060,7 +1074,11 @@ def dataset_module_factory(
         raise RuntimeError(f"Dataset scripts are no longer supported, but found {filename}")
     elif os.path.isdir(path) and not remote_only:
         return LocalDatasetModuleFactory(
-            path, data_dir=data_dir, data_files=data_files, download_mode=download_mode
+            path,
+            data_dir=data_dir,
+            data_files=data_files,
+            download_mode=download_mode,
+            skip_origin_metadata=skip_origin_metadata,
         ).get_module()
     # Try remotely
     elif path.startswith("buckets/"):
@@ -1097,6 +1115,7 @@ def dataset_module_factory(
             data_files=data_files,
             download_config=download_config,
             download_mode=download_mode,
+            skip_origin_metadata=skip_origin_metadata,
         ).get_module()
     elif is_relative_path(path) and path.count("/") <= 1:
         try:
@@ -1179,6 +1198,7 @@ def dataset_module_factory(
                     download_config=download_config,
                     download_mode=download_mode,
                     use_exported_dataset_infos=use_exported_dataset_infos,
+                    skip_origin_metadata=skip_origin_metadata,
                 ).get_module()
             except GatedRepoError as e:
                 message = f"Dataset '{path}' is a gated dataset on the Hub."
@@ -1312,6 +1332,7 @@ def load_dataset_builder(
         download_config.storage_options.update(storage_options)
     if features is not None:
         features = _fix_for_backward_compatible_features(features)
+    skip_origin_metadata = config_kwargs.pop("skip_origin_metadata", False)
     dataset_module = dataset_module_factory(
         path,
         revision=revision,
@@ -1320,6 +1341,7 @@ def load_dataset_builder(
         data_dir=data_dir,
         data_files=data_files,
         cache_dir=cache_dir,
+        skip_origin_metadata=skip_origin_metadata,
     )
     # Get dataset builder class
     builder_kwargs = dataset_module.builder_kwargs
@@ -1679,6 +1701,12 @@ def load_dataset(
             "To parallelize streaming, you can wrap the dataset with a PyTorch DataLoader using `num_workers` > 1 instead."
         )
 
+    if not streaming and config_kwargs.get("skip_origin_metadata", False):
+        logger.warning(
+            "`skip_origin_metadata` is only supported when `streaming=True`. Forcing `skip_origin_metadata=False`."
+        )
+        config_kwargs["skip_origin_metadata"] = False
+
     download_mode = DownloadMode(download_mode or DownloadMode.REUSE_DATASET_IF_EXISTS)
     verification_mode = VerificationMode(
         (verification_mode or VerificationMode.BASIC_CHECKS) if not save_infos else VerificationMode.ALL_CHECKS

tests/test_data_files.py

@@ -400,6 +400,17 @@ def test_DataFilesList_from_patterns_raises_FileNotFoundError(complex_data_dir):
         DataFilesList.from_patterns(["file_that_doesnt_exist.txt"], complex_data_dir)
 
 
+def test_DataFilesList_from_patterns_skip_origin_metadata(complex_data_dir, text_file):
+    data_files_list = DataFilesList.from_patterns([_TEST_URL, text_file.as_posix()], complex_data_dir, skip_origin_metadata=True)
+    assert list(data_files_list) == [_TEST_URL, text_file.as_posix()]
+    assert data_files_list.origin_metadata == [()] * 2
+
+
+def test_DataFilesList_from_patterns_skip_origin_metadata_raises_FileNotFoundError(complex_data_dir):
+    with pytest.raises(FileNotFoundError):
+        DataFilesList.from_patterns(["file_that_doesnt_exist.txt"], complex_data_dir, skip_origin_metadata=True)
+
+
 class TestDataFilesDict:
     def test_key_order_after_copy(self):
         data_files = DataFilesDict({"train": "train.csv", "test": "test.csv"})

tests/test_load.py

@@ -1125,6 +1125,11 @@ def test_load_streaming_private_dataset_with_zipped_data(hf_token, hf_private_da
     assert next(iter(ds)) is not None
 
 
+def test_load_dataset_skip_origin_metadata_warning(caplog, csv_path):
+    load_dataset("csv", data_files={"train": csv_path}, skip_origin_metadata=True)
+    assert "`skip_origin_metadata` is only supported when `streaming=True`" in caplog.text
+
+
 @require_pil
 @pytest.mark.integration
 def test_load_dataset_config_kwargs_passed_as_arguments():