Add Dataset.from_generator to materialize a stream or a Julia iterable

AGENTS.md

@@ -33,8 +33,10 @@ that lazily converts observations to Julia types).
   formatter: this is how `datasets`' spawned `num_proc` workers (Windows) get the format
   without Julia. A plain Python process without that `sys.path` entry doesn't know
   `"julia"`, so `save_to_disk` (src/serialization.jl) saves `"julia"` as `"numpy"`. It also
-  defines `JuliaFunction`, the wrapper that pickles a Julia `map`/`filter` callback as its
-  content fingerprint (see `src/fingerprint.jl`).
+  defines `JuliaFunction`, the wrapper that pickles a Julia `map`/`filter` callback, or the
+  generator function of `from_generator`, as its content fingerprint (see `src/fingerprint.jl`),
+  and the `examples`/`batch_rows` generators that pull a Julia iterable's examples (or split
+  its batches of columns into rows) for `from_generator`.
 - `src/dataset.jl` — `Dataset`, the wrapper over `datasets.Dataset`. 1-based
   indexing, format/transform machinery (`with_format`, `set_format!`,
   `with_jltransform`, `set_jltransform!`, `reset_format!`).
@@ -62,8 +64,12 @@ that lazily converts observations to Julia types).
   `streaming=true`, an `IterableDataset`/`IterableDatasetDict`).
 - `src/toplevel.jl` — Julia wrappers for module-level `datasets` functions
   (`concatenate_datasets`, `interleave_datasets`, `load_from_disk`, and
-  `from_csv`/`from_json`/`from_parquet`) that re-wrap results in the default
-  `"julia"` format.
+  `from_csv`/`from_json`/`from_parquet`/`from_generator`) that re-wrap results in the default
+  `"julia"` format. `from_generator` (also `Dataset(::IterableDataset)`) materializes a stream
+  (iterated unformatted, with its `features`/`split`) or a Julia iterable into a cached `Dataset`;
+  the cache is keyed by `_content_hash` of a Julia source, by the user's `fingerprint`, or else
+  by a random one in `datasets`' temporary directory. It rethrows the cause of `datasets`'
+  `DatasetGenerationError`.
 - `src/transforms.jl` — `py2jl` / `numpy2jl` / `jl2numpy` / `jl2py`. `py2jl`
   recursively converts Python containers, numpy arrays (copyless, zero-copy), and
   PIL images into Julia types; `jl2py` is the write-path dual. The `"julia"`
@@ -108,8 +114,10 @@ that lazily converts observations to Julia types).
 Managed by [CondaPkg.jl](https://github.com/JuliaPy/CondaPkg.jl) via
 `CondaPkg.toml` (conda-forge channel): `datasets`, `numpy`, `pillow`.
 
-- `datasets` is pinned to `>=4.0, <6`. Do **not** lower the floor below 4.0
-  (the lazy `Column` column-access path depends on datasets ≥ 4).
+- `datasets` is pinned to `>=5.1, <6`. Do **not** lower the floor below 5.1
+  (the lazy `Column` column-access path depends on datasets ≥ 4, and a batched
+  `IterableDataset.map` under the `"julia"` format relies on 5.1 stacking the numpy scalars of
+  re-batched rows into a typed array instead of an `object` one).
 - `python-gil` keeps the env on a GIL-enabled CPython, excluding conda-forge's
   free-threaded (no-GIL) builds (e.g. `cp314t`): PythonCall's free-threading support
   is experimental, and `datasets`' C extensions re-enable the GIL there anyway.

CHANGELOG.md

@@ -8,6 +8,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [0.4.4]
 
 ### Changed
+- The Python dependencies require `datasets` 5.1 (`>=5.1, <6`, previously `>=4.0`); an existing
+  environment with an older `datasets` is upgraded the next time CondaPkg resolves it.
 - `ds.iter(; batch_size, drop_last_batch = false)` on a `Dataset` or `IterableDataset` returns
   a `BatchIterator` whose batches are in the dataset's format, as indexing with a range returns
   them: a `Dict{String}` of Julia columns in the `"julia"` format (e.g. a `Vector{String}` for a
@@ -26,6 +28,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   Python functions given to `ds.py.map` still see an `object` array of 1-D arrays.
 
 ### Added
+- `Dataset.from_generator(source)` (also the public `HuggingFaceDatasets.from_generator`) builds a
+  cached, memory-mapped `Dataset` from the examples of an `IterableDataset`, of a Julia iterable
+  of `Dict`s or `NamedTuple`s (or a function returning one), or of a Python generator function;
+  `Dataset(stream)` does the same for a stream (#109). This materializes a streamed pipeline,
+  e.g. the token ids of a corpus streamed from the Hub, without keeping the raw data. With
+  `batched = true` a Julia iterable yields batches of columns, which convert faster. A Julia
+  source is cached under a fingerprint of its content, as for `map`; a stream is rebuilt on every
+  call, in a temporary directory, unless given a `fingerprint`. An error raised while generating
+  the examples is thrown as it is, instead of `datasets`' uninformative `DatasetGenerationError`.
 - `py2jl` converts PythonCall's `PyDict`, `PyList`, `PySet` and `PyArray` like the Python object
   they wrap, instead of returning them unchanged.
 

CondaPkg.toml

@@ -1,7 +1,7 @@
 channels = ["conda-forge"]
 
 [deps]
-datasets = ">=4.0, <6"
+datasets = ">=5.1, <6"
 numpy = ""
 pillow = ""
 # Exclude conda-forge's free-threaded (no-GIL) CPython builds, e.g. python 3.14 `cp314t`:

docs/src/api.md

@@ -41,6 +41,7 @@ load_from_disk
 from_csv
 from_json
 from_parquet
+from_generator
 ```
 
 ## Combining

docs/src/guide.md

@@ -92,6 +92,45 @@ Indexing with a single integer returns one observation; indexing with a range or
 returns a batch, represented as a dictionary mapping each column name to a vector of
 values (see below).
 
+### Materializing a stream or a Julia generator
+
+With `streaming=true`, `load_dataset` returns an [`IterableDataset`](@ref), whose `map` and
+`filter` are lazy: they run again on every pass and are never written to disk.
+[`from_generator`](@ref) (`Dataset.from_generator`, or `Dataset(stream)` for a stream) writes the
+examples of a stream, or of a Julia iterable, to an Arrow file in the cache and returns the
+memory-mapped [`Dataset`](@ref). This keeps only the result of a pipeline: tokenizing a corpus
+streamed from the Hub stores the token ids, without the parquet files and the Arrow copy of the
+text that a non-streaming `load_dataset` followed by `map` keeps.
+
+```julia-repl
+julia> stream = load_dataset("roneneldan/TinyStories"; split = "validation", streaming = true);
+
+julia> blocks = map(stream; batched = true, remove_columns = ["text"]) do b
+           ids = tokenize(b["text"])                                # a Vector{UInt16}
+           Dict("ids" => reshape(ids[1:end ÷ 513 * 513], 513, :))   # blocks of 513 ids
+       end;
+
+julia> ds = Dataset(blocks; fingerprint = "tinystories-validation-gpt2-513");
+
+julia> ds[1:32]["ids"];   # a 513×32 Matrix{UInt16}
+```
+
+A stream can't be fingerprinted, so without `fingerprint` it is rebuilt on every call, in a
+temporary directory. With one, a later call, also in a new session, loads the cached dataset
+without streaming. A Julia iterable of examples (`Dict`s or `NamedTuple`s), or a function
+returning one, is cached under a fingerprint of its content, as the function of `map` is. With
+`batched = true` it yields batches of columns, which convert faster:
+
+```julia-repl
+julia> ds = Dataset.from_generator((; x = i, y = Float32(i) / 2) for i in 1:3);
+
+julia> ds[:]["y"]
+3-element Vector{Float32}:
+ 0.5
+ 1.0
+ 1.5
+```
+
 ## Method forwarding
 
 Any method or property of the wrapped Python object is available directly on the Julia

src/HuggingFaceDatasets.jl

@@ -74,7 +74,7 @@ include("serialization.jl")
 include("fingerprint.jl")
 
 # `public` is a Julia 1.11+ keyword; `@compat` makes it a no-op on the supported 1.10.
-@compat public from_csv, from_json, from_parquet
+@compat public from_csv, from_json, from_parquet, from_generator
 
 # Julia < 1.12 bundles libssh2 1.11.0 and loads it before Python starts (at startup on 1.10, with
 # `Pkg` on 1.11). On Windows, the conda-forge `libcurl.dll` that `pyarrow` links then binds to that

src/dataset.jl

@@ -252,9 +252,9 @@ end
 # Expose the `datasets.Dataset` classmethods under their Python names, so e.g.
 # `Dataset.from_dict(data; ...)` / `Dataset.from_csv(path; ...)` / `Dataset.load_from_disk(path)`
 # mirror `datasets.Dataset.from_dict` / `.from_csv` / `.load_from_disk`. `getproperty` is
-# overloaded on the *type* object itself; the file constructors and `load_from_disk` route to
-# the top-level wrappers in `toplevel.jl` (`from_csv`/`from_json`/`from_parquet` and
-# `_wrap_toplevel`). Every other name falls back to normal `DataType` field access
+# overloaded on the *type* object itself; the file and generator constructors and
+# `load_from_disk` route to the top-level wrappers in `toplevel.jl`
+# (`from_csv`/`from_json`/`from_parquet`/`from_generator` and `_wrap_toplevel`). Every other name falls back to normal `DataType` field access
 # (`Dataset.name`, `Dataset.parameters`, ...), so type introspection is unaffected.
 function Base.getproperty(::Type{Dataset}, s::Symbol)
     if s === :from_dict
@@ -265,6 +265,8 @@ function Base.getproperty(::Type{Dataset}, s::Symbol)
         return from_json
     elseif s === :from_parquet
         return from_parquet
+    elseif s === :from_generator
+        return from_generator
     elseif s === :load_from_disk
         return (path; kws...) -> _wrap_toplevel(datasets.Dataset.load_from_disk(jl2py(path); kws...))
     else

src/python/huggingfacedatasets_jl.py

@@ -72,17 +72,6 @@ def __init__(self, features=None, token_per_repo_id=None, **np_array_kwargs):
         np_array_kwargs.setdefault("dtype", None)
         super().__init__(features=features, token_per_repo_id=token_per_repo_id, **np_array_kwargs)
 
-    def _consolidate(self, column):
-        # A batched `IterableDataset.map` re-batches rows this formatter already split into numpy
-        # scalars. Before 5.1, `datasets` consolidates those into an `object` array; stack them.
-        if (
-            isinstance(column, list)
-            and column
-            and all(isinstance(x, (np.number, np.bool_)) and x.dtype == column[0].dtype for x in column)
-        ):
-            return np.stack(column)
-        return super()._consolidate(column)
-
     def _extractors(self, pa_table, batch):
         # The columns read straight from Arrow, with their extractor: top-level
         # `Value("string")` columns, and in a batch, list columns of numbers. Neither needs
@@ -143,11 +132,13 @@ class JuliaFormatType(str):
 
 
 class JuliaFunction:
-    """A Julia function given to `Dataset.map`/`filter`, with a fingerprint of its code and data.
+    """A Julia function given to `Dataset.map`/`filter`, with a fingerprint of its code and data,
+    or the generator function given to `Dataset.from_generator`, with a fingerprint of its source.
 
-    `datasets` fingerprints a transform by pickling it, and a Julia function would pickle by
-    name. This wrapper pickles as its `fingerprint`, computed by HuggingFaceDatasets.jl, so the
-    cache follows what the function does. It does not unpickle into a working function.
+    `datasets` fingerprints a transform, or a generator, by pickling it, and a Julia function
+    would pickle by name. This wrapper pickles as its `fingerprint`, computed by
+    HuggingFaceDatasets.jl, so the cache follows what the function does. It does not unpickle
+    into a working function.
     """
 
     __slots__ = ("function", "fingerprint")
@@ -172,7 +163,29 @@ def __deepcopy__(self, memo):
 def _unpickled_julia_function(fingerprint):
     def function(*args, **kwargs):
         raise RuntimeError(
-            "A Julia function can only be called in the Julia session that passed it to `map`/`filter`."
+            "A Julia function can only be called in the Julia session that passed it to "
+            "`map`/`filter`/`from_generator`."
         )
 
     return function
+
+
+END = object()  # what the `next_item` functions from Julia return past the last item
+
+
+def examples(next_item):
+    """The examples of a Julia iterable, for `Dataset.from_generator`: what `next_item()`
+    returns until it returns `END`."""
+    while (example := next_item()) is not END:
+        yield example
+
+
+def batch_rows(next_item):
+    """The rows of the batches of a Julia iterable, for `Dataset.from_generator`: of each
+    mapping of column names to columns that `next_item()` returns until it returns `END`."""
+    while (batch := next_item()) is not END:
+        lengths = {name: len(column) for name, column in batch.items()}
+        if len(set(lengths.values())) > 1:
+            raise ValueError(f"The columns of a batch have different lengths: {lengths}")
+        for i in range(next(iter(lengths.values()), 0)):
+            yield {name: column[i] for name, column in batch.items()}

src/ragged.jl

@@ -10,7 +10,7 @@ different lengths, e.g. the token ids of a tokenized text column, so that it is
 column's flat Arrow buffer in one go rather than one array per row. Rows of equal length are
 stacked into a matrix instead, as for any array column. Batches that `datasets` assembles from
 already formatted rows rather than from Arrow are still vectors of vectors: those of an
-`IterableDataset` with `datasets` < 4.8.3, after a batched `map`, or from a generator.
+`IterableDataset` after a batched `map`, or from a generator.
 
 `r.values` and `r.offsets` are the underlying buffers, e.g. to write out all the tokens of a
 batch at once.