diff --git a/AGENTS.md b/AGENTS.md index cac24e7..a30e1db 100644 --- a/AGENTS.md +++ b/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. diff --git a/CHANGELOG.md b/CHANGELOG.md index 90522b8..e174192 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/CondaPkg.toml b/CondaPkg.toml index 19d58d2..546498b 100644 --- a/CondaPkg.toml +++ b/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`: diff --git a/docs/src/api.md b/docs/src/api.md index abe0a73..d9518e1 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -41,6 +41,7 @@ load_from_disk from_csv from_json from_parquet +from_generator ``` ## Combining diff --git a/docs/src/guide.md b/docs/src/guide.md index 0e36aa2..aef33e1 100644 --- a/docs/src/guide.md +++ b/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 diff --git a/src/HuggingFaceDatasets.jl b/src/HuggingFaceDatasets.jl index dfcecb1..c328052 100644 --- a/src/HuggingFaceDatasets.jl +++ b/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 diff --git a/src/dataset.jl b/src/dataset.jl index 5ed7794..8b240cd 100644 --- a/src/dataset.jl +++ b/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 diff --git a/src/python/huggingfacedatasets_jl.py b/src/python/huggingfacedatasets_jl.py index c6992a9..0aed3dc 100644 --- a/src/python/huggingfacedatasets_jl.py +++ b/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()} diff --git a/src/ragged.jl b/src/ragged.jl index f49bbee..11eee19 100644 --- a/src/ragged.jl +++ b/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. diff --git a/src/toplevel.jl b/src/toplevel.jl index 41f4606..29210c7 100644 --- a/src/toplevel.jl +++ b/src/toplevel.jl @@ -171,3 +171,185 @@ See also [`from_csv`](@ref) and [`from_json`](@ref). """ from_parquet(path_or_paths; kws...) = _wrap_toplevel(datasets.Dataset.from_parquet(jl2py(path_or_paths); kws...)) + +""" + from_generator(source; batched = false, fingerprint = nothing, features = nothing, + split = nothing, cache_dir = nothing, kws...) + Dataset.from_generator(source; kws...) + Dataset(stream::IterableDataset; kws...) + +Build a [`Dataset`](@ref) from the examples of `source`, forwarding to +[`datasets.Dataset.from_generator`](https://huggingface.co/docs/datasets/package_reference/main_classes#datasets.Dataset.from_generator). +The examples are written to an Arrow file in the cache, which the returned dataset +memory-maps. `source` can be + +- an [`IterableDataset`](@ref), e.g. a stream from `load_dataset(...; streaming=true)` with + `map`/`filter` steps, iterated in this process. This keeps only the result of a streamed + pipeline, e.g. the token ids and not the text. The stream's `features` and `split` carry over. + `Dataset(stream; kws...)` is the same. +- a Julia iterable of examples (`AbstractDict`s or `NamedTuple`s mapping column names to + values), converted with [`jl2py`](@ref), or a function of no arguments returning one. With + `batched = true`, the iterable yields batches instead: `AbstractDict`s or `NamedTuple`s mapping + column names to columns, which are vectors or arrays whose last axis indexes the observations, + as for [`Dataset(::AbstractDict)`](@ref). That is faster, since each column is converted at + once. +- a Python generator function, handed to `datasets` as is, along with the other keyword + arguments (`gen_kwargs`, `num_proc`, ...). + +`features` sets the schema, otherwise inferred from the examples, and `split` the split name +(`"train"` by default). Other keyword arguments (`keep_in_memory`, ...) are forwarded to Python. +The result is returned in the default `"julia"` format. + +!!! note "Caching" + The dataset is cached under a fingerprint of `source`. A later call with the same + fingerprint, `features` and `split`, also in a later session, loads it from the cache + without iterating `source`. For a Julia source the fingerprint hashes what the source is and + does, as for the function of [`map`](@ref): its values (the examples of a `Vector`, ...), the + code of its functions and of the functions they use from your own modules, the values they + capture and the globals they read, and `batched`. Files or other inputs from outside Julia, + e.g. the lines of a file the source reads, are not tracked. + + An `IterableDataset`, or a source with something that can't be hashed, has no fingerprint: + its dataset is rebuilt on every call, in a temporary directory deleted at the end of the + session (or in `cache_dir`, if given). Pass `fingerprint`, a string naming the examples, to + cache them under it instead, and change it when the examples change. + +!!! warning "No multiprocessing (`num_proc`)" + The examples of a Julia source or of an `IterableDataset` are generated in this process, so + `num_proc > 1` throws an `ArgumentError`. To parallelize a stream's Julia `map`, use + `batched=true` and parallelize over each batch in the function, e.g. with `Threads.@threads`. + +See also [`load_dataset`](@ref) and [`IterableDataset`](@ref). + +# Examples + +```julia +julia> ds = Dataset.from_generator((; x = i, y = Float32(i) / 2) for i in 1:3) +Dataset({ + features: ['x', 'y'], + num_rows: 3 +}) + +julia> ds[:]["y"] +3-element Vector{Float32}: + 0.5 + 1.0 + 1.5 + +julia> ds = Dataset.from_generator(; batched = true) do # 2 batches of 3 rows + (Dict("x" => reshape(Int32.(1:6) .+ Int32(6k), 2, 3)) for k in 0:1) + end; + +julia> ds[:]["x"] +2×6 Matrix{Int32}: + 1 3 5 7 9 11 + 2 4 6 8 10 12 +``` + +Tokenize a streamed corpus into blocks of 513 token ids, and keep only those in the cache: + +```julia +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, :)) + end; + +julia> ds = Dataset(blocks; fingerprint = "tinystories-validation-gpt2-513"); + +julia> ds[1:32]["ids"] # a 513×32 Matrix{UInt16} +``` +""" +function from_generator(source; batched::Bool = false, kws...) + examples, convert = batched ? (pyhfds.batch_rows, _batch_to_py) : (pyhfds.examples, _example_to_py) + generator = pyfunc(() -> examples(_pynext(convert, source isa Function ? source() : source))) + return _from_generator(generator, source, (; batched); kws...) +end + +function from_generator(stream::IterableDataset; features = nothing, split = nothing, kws...) + py = getfield(stream, :py) + features = something(features, py.features) + split = something(split, pyis(py.split, pybuiltins.None) ? datasets.Split.TRAIN : py.split) + # Unformatted: Python values, which `features` encodes as they are in the stream + return _from_generator(py.with_format(nothing).__iter__, stream, (;); features, split, kws...) +end + +function from_generator(source::Py; kws...) + pyisinstance(source, datasets.IterableDataset) && + return from_generator(IterableDataset(source); kws...) + return _generate(source; kws...) +end + +Dataset(stream::IterableDataset; kws...) = from_generator(stream; kws...) + +# `datasets` names the cache of `from_generator` by hashing the generator function, which pickles +# as `fingerprint` in a `JuliaFunction`, along with `features` and `split`. (Its own `fingerprint` +# keyword would name the cache by the fingerprint alone.) +function _from_generator(generator::Py, source, hashed_kws; fingerprint = nothing, + cache_dir = nothing, num_proc = nothing, kws...) + if num_proc isa Integer && num_proc > 1 + throw(ArgumentError( + "`from_generator` does not support `num_proc=$num_proc` for a Julia source or an " * + "`IterableDataset`: it generates their examples in this process, since `datasets`' " * + "worker processes cannot call back into Julia. Omit `num_proc`; to parallelize a " * + "Julia function, have it work on batches with several threads, e.g. with " * + "`Threads.@threads`.")) + end + fp = fingerprint === nothing ? _content_hash(source, hashed_kws) : fingerprint + if !(fp isa AbstractString) + @info "`from_generator` can't fingerprint its source ($(fp.msg)), so the dataset is \ + rebuilt on every call, in a temporary directory unless `cache_dir` is given. Pass \ + `fingerprint` to cache it." maxlog=1 + fp = _random_fingerprint() + cache_dir = something(cache_dir, datasets.fingerprint.get_temporary_cache_files_directory()) + end + return _generate(pyhfds.JuliaFunction(generator, fp); cache_dir, kws...) +end + +# `datasets` raises an error of the generator as the cause of a `DatasetGenerationError`, which +# doesn't say what it is, so throw the cause instead. +function _generate(generator::Py; kws...) + py = try + datasets.Dataset.from_generator(generator; kws...) + catch e + (e isa PyException && pyisinstance(e.v, datasets.exceptions.DatasetGenerationError) && + !pyis(e.v.__cause__, pybuiltins.None)) || rethrow() + throw(PyException(e.v.__cause__)) + end + return _wrap_toplevel(py) +end + +# A Python function returning the next element of `itr`, converted by `convert`, and +# `pyhfds.END` past the last one. +function _pynext(convert, itr) + started = false + state = nothing + return pyfunc() do + r = started ? iterate(itr, state) : iterate(itr) + started = true + r === nothing && return pyhfds.END + x, state = r + return convert(x) + end +end + +_example_to_py(x::Union{AbstractDict,NamedTuple,Py}) = jl2py(x) +_example_to_py(x) = throw(ArgumentError( + "an example must be an `AbstractDict` or a `NamedTuple` of column values, got $(typeof(x))")) + +# Each column at once: a stacked array becomes one numpy array, its observations along the first +# axis, which `pyhfds.batch_rows` splits. +function _batch_to_py(batch::Union{AbstractDict,NamedTuple}) + py = pydict() + for (k, v) in pairs(batch) + v isa AbstractArray || throw(ArgumentError( + "column \"$k\" of a batch must be an `AbstractArray` of observations, got $(typeof(v))")) + py[string(k)] = jl2py(v) + end + return py +end +_batch_to_py(x::Py) = x +_batch_to_py(x) = throw(ArgumentError( + "with `batched = true`, each element must be an `AbstractDict` or a `NamedTuple` of " * + "columns, got $(typeof(x))")) diff --git a/test/dataset.jl b/test/dataset.jl index 23e00f3..cd6e00c 100644 --- a/test/dataset.jl +++ b/test/dataset.jl @@ -215,6 +215,56 @@ end end end + @testset "from_generator (#109)" begin + mktempdir() do dir + rows = [(; x = i, y = Float32(i) / 2, s = "r$i") for i in 1:3] + ds = Dataset.from_generator(rows; cache_dir = dir) + @test ds isa Dataset + @test ds.format["type"] == "julia" + @test ds[:] == Dict("x" => [1, 2, 3], "y" => Float32[0.5, 1, 1.5], "s" => ["r1", "r2", "r3"]) + @test ds[:]["y"] isa Vector{Float32} + + # cached under a fingerprint of the examples + @test HuggingFaceDatasets.from_generator(copy(rows); cache_dir = dir).cache_files == + ds.cache_files + rows[3] = (; x = 3, y = 2.0f0, s = "r3") + ds2 = Dataset.from_generator(rows; cache_dir = dir) + @test ds2[3]["y"] == 2 + @test ds2.cache_files != ds.cache_files + @test Dataset.from_generator(rows; features = ds.features, cache_dir = dir)[3]["y"] == 2 + + # a given fingerprint loads the cached dataset without iterating the source + f = Dataset.from_generator(rows; fingerprint = "rows-v1", cache_dir = dir) + @test Dataset.from_generator(() -> error("not iterated"); fingerprint = "rows-v1", + cache_dir = dir)[:] == f[:] + + # a function returning batches of columns, whose arrays' last axis indexes the rows + b = Dataset.from_generator(; batched = true, cache_dir = dir) do + (Dict("ids" => reshape(UInt16.(1:8) .+ UInt16(8k), 2, 4), "n" => (1:4) .+ 4k) + for k in 0:1) + end + @test length(b) == 8 + @test b[1:8]["ids"] == reshape(1:16, 2, 8) + @test b[1:8]["ids"] isa Matrix{UInt16} + @test b[:]["n"] == 1:8 + + pygen = pyeval("lambda n: ({'a': i} for i in range(n))", @__MODULE__) + @test Dataset.from_generator(pygen; gen_kwargs = pydict(n = 3), cache_dir = dir)[:]["a"] == + [0, 1, 2] + + # an error of the generator is thrown instead of `datasets`' `DatasetGenerationError` + @test_throws ArgumentError Dataset.from_generator(rows; num_proc = 2) + err = try Dataset.from_generator([1, 2]; cache_dir = dir) catch e; e end + @test err isa PyException && occursin("an example must be", sprint(showerror, err)) + err = try + Dataset.from_generator([Dict("a" => [1, 2], "b" => [1])]; batched = true, cache_dir = dir) + catch e + e + end + @test err isa PyException && occursin("different lengths", sprint(showerror, err)) + end + end + @testset "type introspection unaffected by type-level getproperty" begin @test fieldnames(Dataset) == (:py, :jltransform) @test fieldnames(DatasetDict) == (:py, :jltransform, :splits) diff --git a/test/iterabledataset.jl b/test/iterabledataset.jl index 2b5ff7a..bb67b91 100644 --- a/test/iterabledataset.jl +++ b/test/iterabledataset.jl @@ -111,12 +111,6 @@ end @test first(itds)["x"] isa UInt16 # rows keep the column's type too @test [r["x"] for r in itds.filter((x, i) -> iseven(i); with_indices=true)] == [1, 3] - - # the `"julia"` formatter stacks a list of numpy scalars by dtype - fmt = datasets.formatting.get_formatter("julia") - a = fmt._consolidate(pylist([np.uint16(1), np.uint16(2)])) - @test pyconvert(String, a.dtype.name) == "uint16" - @test pyconvert(String, fmt._consolidate(pylist([np.uint16(1), np.int8(2)])).dtype.name) == "object" end @testset "julia format keeps each column's dtype" begin @@ -177,3 +171,37 @@ end tr = with_jltransform(itds, row -> py2jl(row)["x"] + 1) @test collect(tr) == [2, 3, 4] end + +@testset "materialize with from_generator (#109)" begin + mktempdir() do dir + src = Dataset((; label = [0, 1, 0], x = Float32[1, 2, 3], text = ["a", "bb", "c"])) + src = set_format!(Dataset(src.py.cast_column("label", datasets.ClassLabel(names = pylist(["neg", "pos"])))), + "julia") + stream = src.to_iterable_dataset() + ds = @test_logs (:info, r"can't fingerprint its source") Dataset(stream) + @test ds isa Dataset + @test ds.format["type"] == "julia" + @test ds[:] == src[:] + @test ds.features["label"] isa ClassLabel + # rebuilt on every call, in a temporary directory + @test Dataset(stream).cache_files != ds.cache_files + tmp = pyconvert(String, datasets.fingerprint.get_temporary_cache_files_directory()) + @test startswith(ds.cache_files[1]["filename"], tmp) + @test Dataset.from_generator(stream.py; cache_dir = dir)[:] == src[:] # a Python stream + + blocks = map(stream; batched = true, remove_columns = ["text"]) do b + Dict("ids" => reshape(UInt16.(1:2length(b["x"])), 2, :)) + end + ds = Dataset.from_generator(blocks; fingerprint = "blocks-v1", cache_dir = dir, split = "validation") + @test ds[1:3]["ids"] == [1 3 5; 2 4 6] + @test ds[1:3]["ids"] isa Matrix{UInt16} + @test string(ds.py.split) == "validation" + @test Dataset(blocks; fingerprint = "blocks-v1", cache_dir = dir, split = "validation").cache_files == + ds.cache_files + @test_throws ArgumentError Dataset(blocks; num_proc = 2) + + gen = pyeval("lambda: iter([{'x': 1}, {'x': 2}])", @__MODULE__) + named = IterableDataset(datasets.IterableDataset.from_generator(gen; split = "test")) + @test string(Dataset(named; cache_dir = dir).py.split) == "test" # the stream's split + end +end