Add Dataset.from_generator (#4957)

* Add `Dataset.from_generator`

* Add tests

* Docs

* Doc typo

* Add Returns to docstring

* Docstring for some params

* Remove docs changes to test CI

* Add from_generator to package reference

* Return doc

* Fix docstring

* Update src/datasets/arrow_dataset.py

Co-authored-by: Quentin Lhoest <[email protected]>

* Use for loop

* Call close on writer instances

Co-authored-by: Quentin Lhoest <[email protected]>

Author: mariosasko

docs/source/loading.mdx

@@ -220,6 +220,21 @@ Load a list of Python dictionaries with [`~Dataset.from_list`]:
 >>> dataset = Dataset.from_list(my_list)
 ```
 
+### Python generator
+
+Create a dataset from a Python generator with [`~Dataset.from_generator`]:
+
+```py
+>>> from datasets import Dataset
+>>> def my_gen():
+...     for i in range(1, 4):
+...         yield {"a": i}
+...
+>>> dataset = Dataset.from_generator(my_dict)
+```
+
+This approach supports loading data larger than available memory.
+
 ### Pandas DataFrame
 
 Load Pandas DataFrames with [`~Dataset.from_pandas`]:

docs/source/package_reference/main_classes.mdx

@@ -16,6 +16,7 @@ The base class [`Dataset`] implements a Dataset backed by an Apache Arrow table.
     - from_buffer
     - from_pandas
     - from_dict
+    - from_generator
     - data
     - cache_files
     - num_columns

src/datasets/arrow_dataset.py

@@ -928,6 +928,46 @@ def from_csv(
             path_or_paths, split=split, features=features, cache_dir=cache_dir, keep_in_memory=keep_in_memory, **kwargs
         ).read()
 
+    @staticmethod
+    def from_generator(
+        generator: Callable,
+        features: Optional[Features] = None,
+        cache_dir: str = None,
+        keep_in_memory: bool = False,
+        gen_kwargs: Optional[dict] = None,
+    ):
+        """Create a Dataset from a generator.
+
+        Args:
+            generator (:obj:`Callable`): A generator function that `yields` examples.
+            features (:class:`Features`, optional): Dataset features.
+            cache_dir (:obj:`str`, optional, default ``"~/.cache/huggingface/datasets"``): Directory to cache data.
+            keep_in_memory (:obj:`bool`, default ``False``): Whether to copy the data in-memory.
+            gen_kwargs(:obj:`dict`, optional): Keyword arguments to be passed to the `generator` callable.
+
+        Returns:
+            :class:`Dataset`
+
+        Example:
+
+        ```py
+        >>> def gen():
+        ...     yield {"text": "Good", "label": 0}
+        ...     yield {"text": "Bad", "label": 1}
+        ...
+        >>> ds = Dataset.from_generator(gen)
+        ```
+        """
+        from .io.generator import GeneratorDatasetInputStream
+
+        return GeneratorDatasetInputStream(
+            generator=generator,
+            features=features,
+            cache_dir=cache_dir,
+            keep_in_memory=keep_in_memory,
+            gen_kwargs=gen_kwargs,
+        ).read()
+
     @staticmethod
     def from_json(
         path_or_paths: Union[PathLike, List[PathLike]],

src/datasets/builder.py

@@ -1371,6 +1371,7 @@ def _prepare_split(
             ):
                 if max_shard_size is not None and writer._num_bytes > max_shard_size:
                     num_examples, num_bytes = writer.finalize()
+                    writer.close()
                     total_num_examples += num_examples
                     total_num_bytes += num_bytes
                     shard_id += 1
@@ -1382,11 +1383,12 @@ def _prepare_split(
                         check_duplicates=check_duplicate_keys,
                         storage_options=self._fs.storage_options,
                     )
-                example = self.info.features.encode_example(record)
+                example = self.info.features.encode_example(record) if self.info.features is not None else record
                 writer.write(example, key)
         finally:
             num_shards = shard_id + 1
             num_examples, num_bytes = writer.finalize()
+            writer.close()
             total_num_examples += num_examples
             total_num_bytes += num_bytes
 
@@ -1492,6 +1494,7 @@ def _prepare_split(
             ):
                 if max_shard_size is not None and writer._num_bytes > max_shard_size:
                     num_examples, num_bytes = writer.finalize()
+                    writer.close()
                     total_num_examples += num_examples
                     total_num_bytes += num_bytes
                     shard_id += 1
@@ -1504,6 +1507,7 @@ def _prepare_split(
         finally:
             num_shards = shard_id + 1
             num_examples, num_bytes = writer.finalize()
+            writer.close()
             total_num_examples += num_examples
             total_num_bytes += num_bytes
 

src/datasets/io/abc.py

@@ -26,3 +26,21 @@ def __init__(
     @abstractmethod
     def read(self) -> Union[Dataset, DatasetDict]:
         pass
+
+
+class AbstractDatasetInputStream(ABC):
+    def __init__(
+        self,
+        features: Optional[Features] = None,
+        cache_dir: str = None,
+        keep_in_memory: bool = False,
+        **kwargs,
+    ):
+        self.features = features
+        self.cache_dir = cache_dir
+        self.keep_in_memory = keep_in_memory
+        self.kwargs = kwargs
+
+    @abstractmethod
+    def read(self) -> Dataset:
+        pass

src/datasets/io/generator.py

@@ -0,0 +1,47 @@
+from typing import Callable, Optional
+
+from .. import Features
+from ..packaged_modules.generator.generator import Generator
+from .abc import AbstractDatasetInputStream
+
+
+class GeneratorDatasetInputStream(AbstractDatasetInputStream):
+    def __init__(
+        self,
+        generator: Callable,
+        features: Optional[Features] = None,
+        cache_dir: str = None,
+        keep_in_memory: bool = False,
+        gen_kwargs: Optional[dict] = None,
+        **kwargs,
+    ):
+        super().__init__(features=features, cache_dir=cache_dir, keep_in_memory=keep_in_memory, **kwargs)
+        self.builder = Generator(
+            cache_dir=cache_dir,
+            features=features,
+            generator=generator,
+            gen_kwargs=gen_kwargs,
+            **kwargs,
+        )
+
+    def read(self):
+        download_config = None
+        download_mode = None
+        ignore_verifications = False
+        use_auth_token = None
+        base_path = None
+
+        self.builder.download_and_prepare(
+            download_config=download_config,
+            download_mode=download_mode,
+            ignore_verifications=ignore_verifications,
+            # try_from_hf_gcs=try_from_hf_gcs,
+            base_path=base_path,
+            use_auth_token=use_auth_token,
+        )
+
+        # Build dataset for splits
+        dataset = self.builder.as_dataset(
+            split="train", ignore_verifications=ignore_verifications, in_memory=self.keep_in_memory
+        )
+        return dataset

src/datasets/packaged_modules/generator/__init__.py

@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+import datasets
+
+
+@dataclass
+class GeneratorConfig(datasets.BuilderConfig):
+    generator: Optional[Callable] = None
+    gen_kwargs: Optional[dict] = None
+    features: Optional[datasets.Features] = None
+
+    def __post_init__(self):
+        assert self.generator is not None, "generator must be specified"
+
+        if self.gen_kwargs is None:
+            self.gen_kwargs = {}
+
+
+class Generator(datasets.GeneratorBasedBuilder):
+    BUILDER_CONFIG_CLASS = GeneratorConfig
+
+    def _info(self):
+        return datasets.DatasetInfo(features=self.config.features)
+
+    def _split_generators(self, dl_manager):
+        return [datasets.SplitGenerator(name=datasets.Split.TRAIN, gen_kwargs={})]
+
+    def _generate_examples(self):
+        for idx, ex in enumerate(self.config.generator(**self.config.gen_kwargs)):
+            yield idx, ex

src/datasets/packaged_modules/generator/generator.py

@@ -3178,6 +3178,60 @@ def test_dataset_from_text_path_type(path_type, text_path, tmp_path):
     _check_text_dataset(dataset, expected_features)
 
 
+@pytest.fixture
+def data_generator():
+    def _gen():
+        data = [
+            {"col_1": "0", "col_2": 0, "col_3": 0.0},
+            {"col_1": "1", "col_2": 1, "col_3": 1.0},
+            {"col_1": "2", "col_2": 2, "col_3": 2.0},
+            {"col_1": "3", "col_2": 3, "col_3": 3.0},
+        ]
+        for item in data:
+            yield item
+
+    return _gen
+
+
+def _check_generator_dataset(dataset, expected_features):
+    assert isinstance(dataset, Dataset)
+    assert dataset.num_rows == 4
+    assert dataset.num_columns == 3
+    assert dataset.column_names == ["col_1", "col_2", "col_3"]
+    for feature, expected_dtype in expected_features.items():
+        assert dataset.features[feature].dtype == expected_dtype
+
+
+@pytest.mark.parametrize("keep_in_memory", [False, True])
+def test_dataset_from_generator_keep_in_memory(keep_in_memory, data_generator, tmp_path):
+    cache_dir = tmp_path / "cache"
+    expected_features = {"col_1": "string", "col_2": "int64", "col_3": "float64"}
+    with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
+        dataset = Dataset.from_generator(data_generator, cache_dir=cache_dir, keep_in_memory=keep_in_memory)
+    _check_generator_dataset(dataset, expected_features)
+
+
+@pytest.mark.parametrize(
+    "features",
+    [
+        None,
+        {"col_1": "string", "col_2": "int64", "col_3": "float64"},
+        {"col_1": "string", "col_2": "string", "col_3": "string"},
+        {"col_1": "int32", "col_2": "int32", "col_3": "int32"},
+        {"col_1": "float32", "col_2": "float32", "col_3": "float32"},
+    ],
+)
+def test_dataset_from_generator_features(features, data_generator, tmp_path):
+    cache_dir = tmp_path / "cache"
+    default_expected_features = {"col_1": "string", "col_2": "int64", "col_3": "float64"}
+    expected_features = features.copy() if features else default_expected_features
+    features = (
+        Features({feature: Value(dtype) for feature, dtype in features.items()}) if features is not None else None
+    )
+    dataset = Dataset.from_generator(data_generator, features=features, cache_dir=cache_dir)
+    _check_generator_dataset(dataset, expected_features)
+
+
 def test_dataset_to_json(dataset, tmp_path):
     file_path = tmp_path / "test_path.jsonl"
     bytes_written = dataset.to_json(path_or_buf=file_path)