It looks like there isn’t currently an official way to do this:
Just in case, @lhoestq
The closest direct answer I found is from datasets
maintainer lhoestq
on an almost identical question about resuming an interrupted Dataset.from_generator()
build: “It’s not currently possible”, with the suggestion to split the work into multiple Dataset
objects so that one failure does not invalidate everything.
So I would probably keep your overall streaming/preprocessing approach, but move the failure/restart boundary outside one giant Dataset.from_generator()
call.
The simplest default is:
source shard/chunk
↓
expensive processing
↓
independently finalized output shard
↓
next shard/chunk
Then a restart only needs to redo the last unfinished unit instead of rebuilding the entire dataset.
If your source already has natural shards/files, I would use those first. If it does not, I would create bounded output chunks myself.
The important distinction is:
“some Arrow data has been written to disk” is not the same thing as “I have a resumable checkpoint.”
Dataset.from_generator()
does write progressively while building, but the public API does not expose those intermediate builder files as a supported resume point.
A useful decision flow is therefore:
Does the source already have stable shards/files?
├── yes
│ └── process one source shard
│ → finalize one output shard
│ → skip completed shards on restart
│
└── no
└── is source order stable and roughly 1 input → 1 output?
├── yes
│ └── create bounded output chunks
│ → finalize each chunk independently
│
└── no
└── use stable IDs / a manifest / source state
instead of relying only on a row counter
That is also close to a later suggestion from lhoestq
for streaming preprocessing: stream the dataset, apply .map(...)
, then process it shard-by-shard and write each shard separately.
So, for your current case, my default order would be:
That preserves the basic idea of generating the dataset progressively; it just makes the unit of successful work smaller than the entire Dataset.from_generator()
build.