Add support for Split.ALL

docs/source/package_reference/builder_classes.rst

@@ -23,6 +23,8 @@ Two main classes are mostly used during the dataset building process.
 
 .. autoclass:: datasets.NamedSplit
 
+.. autoclass:: datasets.NamedSplitAll
+
 .. autoclass:: datasets.ReadInstruction
 
 .. autoclass:: datasets.utils::DownloadConfig

src/datasets/__init__.py

@@ -62,7 +62,17 @@
 from .keyhash import KeyHasher
 from .load import import_main_class, load_dataset, load_from_disk, load_metric, prepare_module
 from .metric import Metric
-from .splits import NamedSplit, Split, SplitBase, SplitDict, SplitGenerator, SplitInfo, SubSplitInfo, percent
+from .splits import (
+    NamedSplit,
+    NamedSplitAll,
+    Split,
+    SplitBase,
+    SplitDict,
+    SplitGenerator,
+    SplitInfo,
+    SubSplitInfo,
+    percent,
+)
 from .utils import *
 from .utils.tqdm_utils import disable_progress_bar
 

src/datasets/arrow_dataset.py

@@ -58,7 +58,7 @@
 from .formatting import format_table, get_format_type_from_alias, get_formatter, query_table
 from .info import DatasetInfo
 from .search import IndexableMixin
-from .splits import NamedSplit
+from .splits import NamedSplit, Split
 from .table import (
     ConcatenationTable,
     InMemoryTable,
@@ -718,7 +718,7 @@ def load_from_disk(dataset_path: str, fs=None, keep_in_memory: Optional[bool] =
             indices_table = None
 
         split = state["_split"]
-        split = NamedSplit(split) if split is not None else split
+        split = Split(split) if split is not None else split
 
         return Dataset(
             arrow_table=arrow_table,

src/datasets/builder.py

@@ -761,7 +761,10 @@ def _build_single_dataset(
     ):
         """as_dataset for a single split."""
         verify_infos = not ignore_verifications
-        if isinstance(split, str):
+        if not isinstance(split, ReadInstruction):
+            split = str(split)
+            if split == "all":
+                split = "+".join(self.info.splits.keys())
             split = Split(split)
 
         # Build base dataset

src/datasets/splits.py

@@ -377,7 +377,7 @@ def __init__(self):
         super(NamedSplitAll, self).__init__("all")
 
     def __repr__(self):
-        return f"NamedSplitAll({self._name!r})"
+        return "NamedSplitAll()"
 
     def get_read_instruction(self, split_dict):
         # Merge all dataset split together
@@ -398,6 +398,7 @@ class Split:
       model architecture, etc.).
     - `TEST`: the testing data. This is the data to report metrics on. Typically
       you do not want to use this during model iteration as you may overfit to it.
+    - `ALL`: the union of all defined dataset splits.
 
     Note: All splits, including compositions inherit from `datasets.SplitBase`
 
@@ -407,10 +408,11 @@ class Split:
     TRAIN = NamedSplit("train")
     TEST = NamedSplit("test")
     VALIDATION = NamedSplit("validation")
+    ALL = NamedSplitAll()
 
     def __new__(cls, name):
         """Create a custom split with datasets.Split('custom_name')."""
-        return NamedSplit(name)
+        return NamedSplitAll() if name == "all" else NamedSplit(name)
 
 
 # Similar to SplitInfo, but contain an additional slice info

tests/test_builder.py

@@ -267,6 +267,14 @@ def _post_processing_resources(self, split):
             self.assertListEqual(dset.column_names, ["text", "tokens"])
             del dset
 
+            dset = dummy_builder.as_dataset("all")
+            self.assertIsInstance(dset, Dataset)
+            self.assertEqual(dset.split, "train+test")
+            self.assertEqual(len(dset), 20)
+            self.assertDictEqual(dset.features, Features({"text": Value("string"), "tokens": [Value("string")]}))
+            self.assertListEqual(dset.column_names, ["text", "tokens"])
+            del dset
+
         def _post_process(self, dataset, resources_paths):
             return dataset.select([0, 1], keep_in_memory=True)