diff --git a/src/datasets/builder.py b/src/datasets/builder.py index 604cb9cd208..0930c86e54e 100644 --- a/src/datasets/builder.py +++ b/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 self._skip_origin_metadata: + 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 {})) diff --git a/src/datasets/data_files.py b/src/datasets/data_files.py index 3e63c3ac65f..4b00dbfa95d 100644 --- a/src/datasets/data_files.py +++ b/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": diff --git a/src/datasets/load.py b/src/datasets/load.py index e938cb35e6d..1822574cf81 100644 --- a/src/datasets/load.py +++ b/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,6 +853,7 @@ 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 @@ -851,6 +861,7 @@ def __init__( 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." @@ -1320,6 +1340,7 @@ def load_dataset_builder( data_dir=data_dir, data_files=data_files, cache_dir=cache_dir, + skip_origin_metadata=config_kwargs.get("skip_origin_metadata", False), ) # Get dataset builder class builder_kwargs = dataset_module.builder_kwargs @@ -1361,6 +1382,7 @@ def load_dataset_builder( features=features, token=token, storage_options=storage_options, + skip_origin_metadata=config_kwargs.pop("skip_origin_metadata", False), **builder_kwargs, **config_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 diff --git a/tests/test_data_files.py b/tests/test_data_files.py index 77c2974a6ca..776745aa6a2 100644 --- a/tests/test_data_files.py +++ b/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"}) diff --git a/tests/test_load.py b/tests/test_load.py index 051f94a7dcc..621abbc92be 100644 --- a/tests/test_load.py +++ b/tests/test_load.py @@ -1125,6 +1125,18 @@ 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_builder_skip_origin_metadata(csv_path): + builder = load_dataset_builder("csv", data_files={"train": csv_path}, skip_origin_metadata=True) + assert builder._skip_origin_metadata is True + with pytest.raises(ValueError, match="This function is not intended for streaming."): + builder.download_and_prepare() + + +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():