Fix regex get_data_files formatting for base paths (#6322)

* Fix regex from formatting url base_path

* Test test_get_data_patterns from Hub

* simply match basename instead

* more tests

* minor

* remove comment

---------

Co-authored-by: Albert Villanova del Moral <[email protected]>
Co-authored-by: Quentin Lhoest <[email protected]>
Co-authored-by: Quentin Lhoest <[email protected]>

src/datasets/data_files.py

@@ -244,8 +244,10 @@ def _get_data_files_patterns(
  except FileNotFoundError:
  continue
  if len(data_files) > 0:
- pattern = base_path + ("/" if base_path else "") + split_pattern
- splits: Set[str] = {string_to_dict(p, glob_pattern_to_regex(pattern))["split"] for p in data_files}
+ splits: Set[str] = {
+ string_to_dict(xbasename(p), glob_pattern_to_regex(xbasename(split_pattern)))["split"]
+ for p in data_files
+ }
  sorted_splits = [str(split) for split in DEFAULT_SPLITS if split in splits] + sorted(
  splits - set(DEFAULT_SPLITS)
  )

tests/test_data_files.py

@@ -502,8 +502,10 @@ def mock_fs(file_paths: List[str]):
  ["data", "data/train.txt", "data.test.txt"]
  ```
  """
-
- dir_paths = {file_path.rsplit("/", 1)[0] for file_path in file_paths if "/" in file_path}
+ file_paths = [file_path.split("://")[-1] for file_path in file_paths]
+ dir_paths = {
+ "/".join(file_path.split("/")[: i + 1]) for file_path in file_paths for i in range(file_path.count("/"))
+ }
  fs_contents = [{"name": dir_path, "type": "directory"} for dir_path in dir_paths] + [
  {"name": file_path, "type": "file", "size": 10} for file_path in file_paths
  ]
@@ -529,6 +531,7 @@ def ls(self, path, detail=True, refresh=True, **kwargs):
  return DummyTestFS
 
 
+@pytest.mark.parametrize("base_path", ["", "mock://", "my_dir"])
 @pytest.mark.parametrize(
  "data_file_per_split",
  [
@@ -598,20 +601,36 @@ def ls(self, path, detail=True, refresh=True, **kwargs):
  {"test": "test00001.txt"},
  ],
 )
-def test_get_data_files_patterns(data_file_per_split):
+def test_get_data_files_patterns(base_path, data_file_per_split):
  data_file_per_split = {k: v if isinstance(v, list) else [v] for k, v in data_file_per_split.items()}
- file_paths = [file_path for split_file_paths in data_file_per_split.values() for file_path in split_file_paths]
+ data_file_per_split = {
+ split: [
+ base_path + ("/" if base_path and base_path[-1] != "/" else "") + file_path
+ for file_path in data_file_per_split[split]
+ ]
+ for split in data_file_per_split
+ }
+ file_paths = sum(data_file_per_split.values(), [])
  DummyTestFS = mock_fs(file_paths)
  fs = DummyTestFS()
 
  def resolver(pattern):
- return [file_path for file_path in fs.glob(pattern) if fs.isfile(file_path)]
+ pattern = base_path + ("/" if base_path and base_path[-1] != "/" else "") + pattern
+ return [
+ file_path[len(fs._strip_protocol(base_path)) :].lstrip("/")
+ for file_path in fs.glob(pattern)
+ if fs.isfile(file_path)
+ ]
 
- patterns_per_split = _get_data_files_patterns(resolver)
+ patterns_per_split = _get_data_files_patterns(resolver, base_path=base_path)
  assert list(patterns_per_split.keys()) == list(data_file_per_split.keys()) # Test split order with list()
  for split, patterns in patterns_per_split.items():
  matched = [file_path for pattern in patterns for file_path in resolver(pattern)]
- assert matched == data_file_per_split[split]
+ expected = [
+ fs._strip_protocol(file_path)[len(fs._strip_protocol(base_path)) :].lstrip("/")
+ for file_path in data_file_per_split[split]
+ ]
+ assert matched == expected
 
 
 @pytest.mark.parametrize(

tests/test_upstream_hub.py

@@ -27,6 +27,7 @@
  load_dataset_builder,
 )
 from datasets.config import METADATA_CONFIGS_FIELD
+from datasets.data_files import get_data_patterns
 from datasets.packaged_modules.folder_based_builder.folder_based_builder import (
  FolderBasedBuilder,
  FolderBasedBuilderConfig,
@@ -884,3 +885,22 @@ def test_load_dataset_with_metadata_file(self, temporary_repo, text_file_with_me
  generator = builder._generate_examples(**gen_kwargs)
  result = [example for _, example in generator]
  assert len(result) == 1
+
+ def test_get_data_patterns(self, temporary_repo, tmp_path):
+ repo_dir = tmp_path / "test_get_data_patterns"
+ data_dir = repo_dir / "data"
+ data_dir.mkdir(parents=True)
+ data_file = data_dir / "train-00001-of-00009.parquet"
+ data_file.touch()
+ with temporary_repo() as repo_id:
+ self._api.create_repo(repo_id, token=self._token, repo_type="dataset")
+ self._api.upload_folder(
+ folder_path=str(repo_dir),
+ repo_id=repo_id,
+ repo_type="dataset",
+ token=self._token,
+ )
+ data_file_patterns = get_data_patterns(f"hf://datasets/{repo_id}")
+ assert data_file_patterns == {
+ "train": ["data/train-[0-9][0-9][0-9][0-9][0-9]-of-[0-9][0-9][0-9][0-9][0-9]*.*"]
+ }