Process-parallel DataLoader (num_workers) support for Dataset

.gitignore

@@ -7,4 +7,5 @@ Manifest.toml
 .vscode
 .ipynb_checkpoints
 LocalPreferences.toml
-.claude/
+.claude/
+__pycache__/

CHANGELOG.md

@@ -7,6 +7,39 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- `MLUtils.DataLoader(ds::Dataset; num_workers=N)` works out of the box for process-parallel
+  loading, including over `mapobs`/`ObsView`-wrapped datasets. Building on the `Serialization`
+  support added in 0.4.0, the GIL-safe serialization it relies on — running the dataset's
+  `pickle` on the main task and loading this package on the workers — is handled upstream by
+  MLUtils, so no wrapper type is exported here. An `@info` notes when an in-memory dataset is
+  first materialized to a temporary Arrow file for this. Requires MLUtils ≥ 0.4.11.
+
+## [0.4.1] - 2026-07-06
+
+### Added
+- Julia views over a dataset's schema. `ds.features` now returns a `Features` view (an
+  `AbstractDict{String, Any}`) instead of a raw `Py` mapping: indexing a column yields a wrapped
+  `ClassLabel`/`Value` leaf (other feature types stay raw `Py`), each forwarding attribute and
+  method access to Python (`cl.names`, `cl.num_classes`, `cl.int2str(i)`, `cl.str2int(s)`,
+  `v.dtype`). The views can be built from Julia (`ClassLabel(names=[…])`, `Value("int64")`,
+  `Features(Dict(…))`) and handed back to Python via `jl2py` (e.g. a `features=` schema
+  argument). Public but unexported Julian conveniences look up a column's `ClassLabel` in one
+  call: `class_names(ds, col)`, `int2str(ds, col, i)`, `str2int(ds, col, s)` (label ids are
+  0-based data, passed through with no index offset), plus `features(ds)` as the function form
+  of `ds.features`. `Features`/`ClassLabel`/`Value` and these helpers are all public but not
+  exported — the Pythonic idioms (`ds.features`, method chaining, `datasets.ClassLabel(…)`) are
+  the primary interface.
+
+### Changed
+- Depend on `MLCore` (≥ 1.1) instead of `MLUtils` for the `getobs`/`numobs` observation
+  interface, and drop the package's own `getobs(::Py, ::Integer)` method (type piracy on
+  `PythonCall.Py`). `getobs` for Python observation containers is now provided by MLCore's
+  `PythonCall` extension. `MLUtils.DataLoader` still works, since MLUtils re-exports the
+  same `getobs`/`numobs` from MLCore.
+
+## [0.4.0] - 2026-07-02
+
 This is a breaking release (0.3.x → 0.4.0). The headline change is that datasets are
 now returned in the `"julia"` format by default, so observations come back as native
 Julia values instead of raw Python objects. See **Breaking** below before upgrading.
@@ -37,18 +70,6 @@ Julia values instead of raw Python objects. See **Breaking** below before upgrad
   `BoundsError`, instead of `AssertionError` — update any code catching `AssertionError`.
 
 ### Added
-- Julia views over a dataset's schema. `ds.features` now returns a `Features` view (an
-  `AbstractDict{String, Any}`) instead of a raw `Py` mapping: indexing a column yields a wrapped
-  `ClassLabel`/`Value` leaf (other feature types stay raw `Py`), each forwarding attribute and
-  method access to Python (`cl.names`, `cl.num_classes`, `cl.int2str(i)`, `cl.str2int(s)`,
-  `v.dtype`). The views can be built from Julia (`ClassLabel(names=[…])`, `Value("int64")`,
-  `Features(Dict(…))`) and handed back to Python via `jl2py` (e.g. a `features=` schema
-  argument). Public but unexported Julian conveniences look up a column's `ClassLabel` in one
-  call: `class_names(ds, col)`, `int2str(ds, col, i)`, `str2int(ds, col, s)` (label ids are
-  0-based data, passed through with no index offset), plus `features(ds)` as the function form
-  of `ds.features`. `Features`/`ClassLabel`/`Value` and these helpers are all public but not
-  exported — the Pythonic idioms (`ds.features`, method chaining, `datasets.ClassLabel(…)`) are
-  the primary interface.
 - `Serialization` support for `Dataset`: `Serialization.serialize`/`deserialize` now work,
   so a `Dataset` can cross a process boundary — the prerequisite for process-parallel data
   loaders (e.g. a `MLUtils.DataLoader(ds; num_workers=N)` that spreads `getobs` over worker
@@ -153,11 +174,6 @@ Julia values instead of raw Python objects. See **Breaking** below before upgrad
   `Dataset` (was `pyds`) and `DatasetDict` (was `pyd`), for consistency. This field is
   not part of the public API (it is shadowed by `getproperty`), so the change is
   non-breaking for documented usage.
-- Depend on `MLCore` (≥ 1.1) instead of `MLUtils` for the `getobs`/`numobs` observation
-  interface, and drop the package's own `getobs(::Py, ::Integer)` method (type piracy on
-  `PythonCall.Py`). `getobs` for Python observation containers is now provided by MLCore's
-  `PythonCall` extension. `MLUtils.DataLoader` still works, since MLUtils re-exports the
-  same `getobs`/`numobs` from MLCore.
 
 ### Fixed
 - Fixed a segfault triggered by REPL tab-completion on a `DatasetDict` (`d[<TAB>`), which
@@ -188,5 +204,7 @@ Baseline. Changes up to and including this release are recorded in the
 [git history](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/commits/main) and the
 [GitHub releases](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/releases).
 
-[Unreleased]: https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/compare/v0.3.4...HEAD
+[Unreleased]: https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/compare/v0.4.1...HEAD
+[0.4.1]: https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/compare/v0.4.0...v0.4.1
+[0.4.0]: https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/compare/v0.3.4...v0.4.0
 [0.3.4]: https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/releases/tag/v0.3.4

docs/src/guide.md

@@ -11,7 +11,8 @@ end
 
 This guide covers how the wrapper relates to the underlying Python `datasets` library,
 the `"julia"` format and the transform pipeline, array orientation and working with
-images, and how to feed a dataset into a Julia data loader.
+images, and how to feed a dataset into a Julia data loader — including process-parallel
+(`num_workers`) loading that scales past CPython's GIL.
 
 The examples below build small datasets in memory with the [`Dataset`](@ref) constructor
 (which accepts a `Dict` or `NamedTuple` of columns) so that they are self-contained and
@@ -292,7 +293,7 @@ Note that the format and the custom transform share the same slot: setting the `
 format installs `py2jl` as the transform, and `with_jltransform` then replaces it — which
 is why the transform above calls `py2jl` itself. For layering additional per-batch
 processing on top of the `"julia"` format, prefer `MLUtils.mapobs` (see below), as in the
-[`examples/flux_mnist.jl`](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/blob/main/examples/flux_mnist.jl)
+[`examples/flux_mnist/flux_mnist.jl`](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/blob/main/examples/flux_mnist/flux_mnist.jl)
 script.
 
 The order of operations when you index is:
@@ -422,7 +423,7 @@ numeric array**, not an `ImageCore` colorview. Combined with the axis reversal a
   while a ragged (variable-size) column falls back to a `Vector` of per-row arrays.
 
 That raw layout is what you want for feeding a model (see
-[Integration with MLUtils and data loaders](@ref) and the MNIST example), but to *look* at
+[Data Loaders](@ref) and the MNIST example), but to *look* at
 an image you turn it back into a colorview and undo the transpose with `permutedims`. Two
 small helpers cover the common cases:
 
@@ -467,11 +468,11 @@ The same `to_rgb` helper handles color datasets (e.g. `to_rgb(ds[1]["img"])` on
 and it round-trips exactly: the recovered array is pixel-for-pixel identical to the image
 `datasets` decodes on the Python side.
 
-## Integration with MLUtils and data loaders
+## Integration with MLUtils
 
 A [`Dataset`](@ref) implements the length/`getindex` interface expected by
-[MLUtils.jl](https://github.com/JuliaML/MLUtils.jl), so `numobs`, `getobs`, `mapobs`, and
-data loaders such as `Flux.DataLoader` work directly:
+[MLUtils.jl](https://github.com/JuliaML/MLUtils.jl), so `numobs`, `getobs`, and `mapobs`
+work directly:
 
 ```jldoctest guide
 julia> using MLUtils
@@ -495,23 +496,106 @@ julia> getobs(mapped, 1:4)
  40
 ```
 
-Putting it together for MNIST: the default `"julia"` format already stacks the image
-column into a `(W, H, N)` numeric array, so you `mapobs` a transform that rescales it to
-`Float32` (plus `Flux.onehotbatch` for the labels), then hand the result to a
-`Flux.DataLoader`:
+Because of this, a `Dataset` can be handed straight to a data loader — see below.
+
+## Data Loaders
+
+A [`Dataset`](@ref) is a valid MLUtils data container, so you can hand one straight to
+`MLUtils.DataLoader` (re-exported by Flux as `Flux.DataLoader`), including through a
+`mapobs`/`ObsView` transform. The loader's own options are covered in the
+[MLUtils Data Loaders guide](https://juliaml.github.io/MLUtils.jl/stable/guide/iteration/);
+the notes below are specific to a `Dataset`.
+
+### Iterating batches
+
+Each batch is the dataset's own `"julia"`-format batch: the columns for the batch's rows,
+stacked and keyed by name (numeric array columns stack into an N-D array with the
+observation axis last). Small, self-contained example:
+
+```jldoctest guide
+julia> loader = DataLoader(Dataset((; x=1:6, y=10:10:60)); batchsize=2);
+
+julia> batch = first(loader);
+
+julia> batch["x"], batch["y"]
+([1, 2], [10, 20])
+```
+
+Most training loops want `(input, target)` pairs rather than a `Dict`, so wrap the dataset
+in a `mapobs` that reshapes each observation; the loader then collates the tuples
+column-wise into an `(inputs, targets)` batch you can destructure with `for (x, y) in
+loader`:
+
+```jldoctest guide
+julia> reshape_obs(o) = (o["x"] * 100, o["y"]);   # per-observation: Dict -> (input, target)
+
+julia> loader = DataLoader(mapobs(reshape_obs, Dataset((; x=1:6, y=10:10:60))); batchsize=2);
+
+julia> first(loader)
+([100, 200], [10, 20])
+```
+
+Against the Hub, the same shapes hold — swap the in-memory `Dataset` for a `load_dataset`:
 
 ```julia-repl
-julia> train_data = load_dataset("ylecun/mnist", split="train");
+julia> using MLUtils
+
+julia> ds = load_dataset("ylecun/mnist", split="train");   # "julia" format by default
 
-julia> train_data = mapobs(mnist_transform, train_data)[:];   # lazily map, then materialize
+julia> loader = DataLoader(ds; batchsize=128, shuffle=true);
 
-julia> train_loader = Flux.DataLoader(train_data; batchsize=128, shuffle=true);
+julia> batch = first(loader);   # batch["image"] is (28, 28, 128), batch["label"] is length-128
 ```
 
-Materializing with `[:]` loads the whole (transformed) dataset into memory, which is
-fastest for small datasets like MNIST. Dropping the `[:]` keeps loading on-the-fly, which
-is slower per epoch but avoids holding everything in memory.
+### Materializing vs. reading on the fly
+
+For small datasets you can **materialize** into memory up front with `[:]` (faster per
+epoch, no decoding during iteration); otherwise the loader reads **on the fly**, keeping
+memory flat:
+
+```julia-repl
+julia> materialized = DataLoader(ds[:]; batchsize=128, shuffle=true);   # whole dataset in RAM
+
+julia> on_the_fly  = DataLoader(ds; batchsize=128, shuffle=true);       # decode per batch
+```
+
+Layer a per-batch transform with `mapobs`; with `[:]` on the *mapped* dataset the transform
+runs once up front, without it the transform runs per batch during iteration:
+
+```julia-repl
+julia> mnist_transform(o) = (Float32.(o["image"]) ./ 255, o["label"]);
+
+julia> train = mapobs(mnist_transform, ds);
+
+julia> loader = DataLoader(train[:]; batchsize=128, shuffle=true);   # or drop [:] for on-the-fly
+```
+
+### Parallel loading (`num_workers`)
+
+Reading on the fly calls into CPython, and CPython's global interpreter lock (**GIL**)
+serializes those reads, so thread-based `parallel=true` gives ~1× on a `Dataset`. Use
+`num_workers=N` instead (needs MLUtils ≥ 0.4.11): it spreads `getobs` over `N` worker
+**processes**, each with its own interpreter and GIL, so loading scales past the GIL:
+
+```julia-repl
+julia> loader = DataLoader(ds; batchsize=128, num_workers=4);   # 4 worker processes
+
+julia> for batch in loader
+           # ... training step ...
+       end
+```
 
-This snippet needs the Hub and the Flux/ImageCore stack, so it is not run as a doctest; a
-complete, runnable version — including `mnist_transform` and the training loop — lives in
-[`examples/flux_mnist.jl`](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/blob/main/examples/flux_mnist.jl).
+This works because a `Dataset` serializes by *reference* to its on-disk Arrow files
+(an in-memory dataset is first materialized to a temporary Arrow file, announced with an
+`@info`), so if you set a `jltransform` it must be serializable. It composes with a
+`mapobs`/`ObsView` wrapper at any nesting — `DataLoader(mapobs(f, ds); num_workers=4)` runs
+`f` on the workers. Batches may arrive **out of input order** across workers; combine with
+`shuffle=true` (or don't rely on order). Worker processes are launched on demand under the
+current `--project`; call `MLUtils.close_dataloader_pool()` to tear them down. Process
+parallelism only pays off when `getobs` is the bottleneck (large images, heavy decode) —
+for cheap observations the IPC overhead can make it slower than serial. See the MLUtils
+guide for the mechanics and tradeoffs.
+
+A complete example — `mapobs` transform, `num_workers`, and a training loop, benchmarked
+against PyTorch — lives in
+[`examples/flux_mnist/`](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/tree/main/examples/flux_mnist).

examples/Project.toml → examples/flux_mnist/Project.toml

@@ -1,6 +1,16 @@
 [deps]
+DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 HuggingFaceDatasets = "d94b9a45-fdf5-4270-b024-5cbb9ef7117d"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+
+[sources]
+HuggingFaceDatasets = {path = "../.."}
+
+[compat]
+DataStructures = "0.19"
+MLUtils = "0.4"