Reimplemented partial split download support (revival of huggingface#6832)

Author: ArjunJagdale

src/datasets/arrow_reader.py

@@ -120,7 +120,7 @@ def make_file_instructions(
             dataset_name=name,
             split=info.name,
             filetype_suffix=filetype_suffix,
-            shard_lengths=name2shard_lengths[info.name],
+            num_shards=len(name2shard_lengths[info.name] or ()),
         )
         for info in split_infos
     }

src/datasets/builder.py

@@ -176,7 +176,7 @@ def __post_init__(self):
             else:
                 self.version = Version.from_dict(self.version)
         if self.splits is not None and not isinstance(self.splits, SplitDict):
-            self.splits = SplitDict.from_split_dict(self.splits)
+            self.splits = SplitDict.from_split_dict(self.splits, self.dataset_name)
         if self.supervised_keys is not None and not isinstance(self.supervised_keys, SupervisedKeysData):
             if isinstance(self.supervised_keys, (tuple, list)):
                 self.supervised_keys = SupervisedKeysData(*self.supervised_keys)

src/datasets/info.py

@@ -1410,6 +1410,7 @@ def load_dataset(
 
     # Download and prepare data
     builder_instance.download_and_prepare(
+        split = split,
         download_config=download_config,
         download_mode=download_mode,
         verification_mode=verification_mode,

src/datasets/load.py

@@ -17,6 +17,7 @@
 
 import itertools
 import os
+import posixpath
 import re
 
 
@@ -46,33 +47,33 @@ def snakecase_to_camelcase(name):
 
 
 def filename_prefix_for_name(name):
-    if os.path.basename(name) != name:
+    if posixpath.basename(name) != name:
         raise ValueError(f"Should be a dataset name, not a path: {name}")
     return camelcase_to_snakecase(name)
 
 
 def filename_prefix_for_split(name, split):
-    if os.path.basename(name) != name:
+    if posixpath.basename(name) != name:
         raise ValueError(f"Should be a dataset name, not a path: {name}")
     if not re.match(_split_re, split):
         raise ValueError(f"Split name should match '{_split_re}'' but got '{split}'.")
     return f"{filename_prefix_for_name(name)}-{split}"
 
 
-def filepattern_for_dataset_split(dataset_name, split, data_dir, filetype_suffix=None):
+def filepattern_for_dataset_split(path, dataset_name, split, filetype_suffix=None):
     prefix = filename_prefix_for_split(dataset_name, split)
+    filepath = posixpath.join(path, prefix)
+    filepath = f"{filepath}*"
     if filetype_suffix:
-        prefix += f".{filetype_suffix}"
-    filepath = os.path.join(data_dir, prefix)
-    return f"{filepath}*"
+        filepath += f".{filetype_suffix}"
+    return filepath
 
 
-def filenames_for_dataset_split(path, dataset_name, split, filetype_suffix=None, shard_lengths=None):
+def filenames_for_dataset_split(path, dataset_name, split, filetype_suffix=None, num_shards=1):
     prefix = filename_prefix_for_split(dataset_name, split)
-    prefix = os.path.join(path, prefix)
+    prefix = posixpath.join(path, prefix)
 
-    if shard_lengths:
-        num_shards = len(shard_lengths)
+    if num_shards > 1:
         filenames = [f"{prefix}-{shard_id:05d}-of-{num_shards:05d}" for shard_id in range(num_shards)]
         if filetype_suffix:
             filenames = [filename + f".{filetype_suffix}" for filename in filenames]

src/datasets/naming.py

@@ -1,6 +1,6 @@
 import itertools
 from dataclasses import dataclass
-from typing import Optional
+from typing import List, Optional
 
 import pyarrow as pa
 
@@ -27,12 +27,18 @@ class Arrow(datasets.ArrowBasedBuilder):
     def _info(self):
         return datasets.DatasetInfo(features=self.config.features)
 
-    def _split_generators(self, dl_manager):
+    def _available_splits(self) -> Optional[List[str]]:
+        return [str(split) for split in self.config.data_files] if isinstance(self.config.data_files, dict) else None
+
+    def _split_generators(self, dl_manager, splits: Optional[List[str]] = None):
         """We handle string, list and dicts in datafiles"""
         if not self.config.data_files:
             raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}")
         dl_manager.download_config.extract_on_the_fly = True
-        data_files = dl_manager.download_and_extract(self.config.data_files)
+        data_files = self.config.data_files
+        if splits and isinstance(data_files, dict):
+            data_files = {split: data_files[split] for split in splits}
+        data_files = dl_manager.download_and_extract(data_files)
         splits = []
         for split_name, files in data_files.items():
             if isinstance(files, str):

src/datasets/packaged_modules/arrow/arrow.py

@@ -153,10 +153,15 @@ def download_and_prepare(self, output_dir: Optional[str] = None, *args, **kwargs
         if output_dir is not None and output_dir != self.cache_dir:
             shutil.copytree(self.cache_dir, output_dir)
 
-    def _split_generators(self, dl_manager):
+    def _available_splits(self) -> Optional[List[str]]:
+        return [str(split) for split in self.info.splits]
+
+    def _split_generators(self, dl_manager, splits: Optional[List[str]] = None):
         # used to stream from cache
         if isinstance(self.info.splits, datasets.SplitDict):
             split_infos: list[datasets.SplitInfo] = list(self.info.splits.values())
+            if splits:
+                split_infos = [split_info for split_info in split_infos if split_info.name in splits]
         else:
             raise ValueError(f"Missing splits info for {self.dataset_name} in cache directory {self.cache_dir}")
         return [
@@ -168,7 +173,7 @@ def _split_generators(self, dl_manager):
                         dataset_name=self.dataset_name,
                         split=split_info.name,
                         filetype_suffix="arrow",
-                        shard_lengths=split_info.shard_lengths,
+                        num_shards=len(split_info.shard_lengths or ()),
                     )
                 },
             )

src/datasets/packaged_modules/cache/cache.py

@@ -148,12 +148,18 @@ class Csv(datasets.ArrowBasedBuilder):
     def _info(self):
         return datasets.DatasetInfo(features=self.config.features)
 
-    def _split_generators(self, dl_manager):
+    def _available_splits(self) -> Optional[List[str]]:
+        return [str(split) for split in self.config.data_files] if isinstance(self.config.data_files, dict) else None
+
+    def _split_generators(self, dl_manager, splits: Optional[List[str]] = None):
         """We handle string, list and dicts in datafiles"""
         if not self.config.data_files:
             raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}")
         dl_manager.download_config.extract_on_the_fly = True
-        data_files = dl_manager.download_and_extract(self.config.data_files)
+        data_files = self.config.data_files
+        if splits and isinstance(data_files, dict):
+            data_files = {split: data_files[split] for split in splits}
+        data_files = dl_manager.download_and_extract(data_files)
         splits = []
         for split_name, files in data_files.items():
             if isinstance(files, str):

src/datasets/packaged_modules/csv/csv.py

@@ -67,7 +67,10 @@ def _info(self):
 
         return datasets.DatasetInfo(features=self.config.features)
 
-    def _split_generators(self, dl_manager):
+    def _available_splits(self) -> Optional[List[str]]:
+        return [str(split) for split in self.config.data_files] if isinstance(self.config.data_files, dict) else None
+
+    def _split_generators(self, dl_manager, splits: Optional[List[str]] = None):
         if not self.config.data_files:
             raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}")
         dl_manager.download_config.extract_on_the_fly = True
@@ -120,6 +123,8 @@ def analyze(files_or_archives, downloaded_files_or_dirs, split):
                             )
 
         data_files = self.config.data_files
+        if splits and isinstance(data_files, dict):
+            data_files = {split: data_files[split] for split in splits}
         splits = []
         for split_name, files in data_files.items():
             if isinstance(files, str):

src/datasets/packaged_modules/folder_based_builder/folder_based_builder.py

@@ -1,7 +1,7 @@
 import io
 import itertools
 from dataclasses import dataclass
-from typing import Optional
+from typing import List, Optional
 
 import pandas as pd
 import pyarrow as pa
@@ -70,12 +70,18 @@ def _info(self):
             raise ValueError("The JSON loader parameter `newlines_in_values` is no longer supported")
         return datasets.DatasetInfo(features=self.config.features)
 
-    def _split_generators(self, dl_manager):
+    def _available_splits(self) -> Optional[List[str]]:
+        return [str(split) for split in self.config.data_files] if isinstance(self.config.data_files, dict) else None
+
+    def _split_generators(self, dl_manager, splits: Optional[List[str]] = None):
         """We handle string, list and dicts in datafiles"""
         if not self.config.data_files:
             raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}")
         dl_manager.download_config.extract_on_the_fly = True
-        data_files = dl_manager.download_and_extract(self.config.data_files)
+        data_files = self.config.data_files
+        if splits and isinstance(data_files, dict):
+            data_files = {split: data_files[split] for split in splits}
+        data_files = dl_manager.download_and_extract(data_files)
         splits = []
         for split_name, files in data_files.items():
             if isinstance(files, str):