I am creating a dataset by streaming another one and processing each sample one by one in a generator function used by Dataset.from_generator(). It is very time consuming and I can’t wait until the end to save it, I always end up losing everything. I need to save it while it’s being created, and I can’t find a way to do this.
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.
If the source already has natural shards
This would be my first choice because it adds very little machinery.
Conceptually:
from pathlib import Path
import os
out_dir = Path("processed")
out_dir.mkdir(exist_ok=True)
num_shards = ds.num_shards
for shard_idx in range(num_shards):
final_path = out_dir / f"{shard_idx:05d}-of-{num_shards:05d}.parquet"
# Already committed by a previous run.
if final_path.exists():
continue
tmp_path = final_path.with_suffix(".parquet.tmp")
shard = ds.shard(
num_shards=num_shards,
index=shard_idx,
)
processed = shard.map(expensive_processing)
processed.to_parquet(tmp_path)
# Local-filesystem example: publish only after the file is complete.
os.replace(tmp_path, final_path)
The useful property here is not specifically Parquet. It is that final_path means:
this unit completed successfully
while an absent final file means:
redo this unit
That gives you a very simple restart protocol without needing to recover internals from Dataset.from_generator().
There are two caveats I would attach to this example.
First, os.replace() is a local-filesystem example. I would not assume identical atomic-rename semantics for arbitrary object stores or remote filesystems.
Second, shard size still matters. The current IterableDataset.to_parquet() path should not be assumed to be a universal constant-memory writer for an arbitrarily huge logical shard. So I would make each restart unit reasonably bounded rather than replacing “one enormous dataset” with “one enormous shard.”
If there are no useful source shards
Then I would probably make the checkpoint boundary explicit myself.
For a simple pipeline where:
- source ordering is stable,
- processing is deterministic enough for replay,
- one input produces one output,
something like this is sufficient:
from pathlib import Path
import os
import pyarrow as pa
import pyarrow.parquet as pq
ROWS_PER_PART = 1000
out_dir = Path("processed")
out_dir.mkdir(exist_ok=True)
def committed_rows():
total = 0
for path in sorted(out_dir.glob("part-*.parquet")):
total += pq.ParquetFile(path).metadata.num_rows
return total
# Important: this should ideally be the cheap/replayable source,
# before the expensive operation.
source_iter = iter(raw_source)
# Replay only the already-committed source prefix.
already_done = committed_rows()
for _ in range(already_done):
next(source_iter)
buffer = []
part_idx = len(list(out_dir.glob("part-*.parquet")))
for example in source_iter:
result = expensive_processing(example)
buffer.append(result)
if len(buffer) >= ROWS_PER_PART:
final_path = out_dir / f"part-{part_idx:06d}.parquet"
tmp_path = out_dir / f"part-{part_idx:06d}.parquet.tmp"
pq.write_table(
pa.Table.from_pylist(buffer),
tmp_path,
)
# Publish only the completed file.
os.replace(tmp_path, final_path)
buffer.clear()
part_idx += 1
if buffer:
final_path = out_dir / f"part-{part_idx:06d}.parquet"
tmp_path = out_dir / f"part-{part_idx:06d}.parquet.tmp"
pq.write_table(
pa.Table.from_pylist(buffer),
tmp_path,
)
os.replace(tmp_path, final_path)
The exact value 1000 is not important.
I would choose the chunk size from the amount of processing you are willing to repeat after a crash, balanced against creating too many tiny files.
For example, if each example costs several seconds to compute, even 100 examples may be a useful commit boundary. If each example is cheap, much larger chunks may make more sense.
The simple row-count resume above stops being sufficient if the transform:
- drops rows,
- expands one input into multiple outputs,
- changes ordering,
- shuffles or interleaves sources,
- depends on mutable external data,
- or the upstream dataset itself can change between runs.
In those cases I would record something more explicit, such as stable source IDs and/or source revision plus a small manifest describing completed chunks.
`IterableDataset.state_dict()` is useful, but solves a different half of the problem
There is now an official resume mechanism for IterableDataset:
state = ds.state_dict()
# later
ds.load_state_dict(state)
The documentation explains that the state tracks the current shard and the position inside that shard. On resume, already-consumed shards are skipped, then the current shard is read again from its beginning until the saved position is reached.
That means I would separate two questions:
Where should I resume reading input?
↓
IterableDataset state / stable IDs / source shard
Which outputs are definitely complete?
↓
finalized Parquet files / manifest / other durable outputs
state_dict() can help with the first one.
It does not by itself make the output from Dataset.from_generator() durable.
There is another subtlety here if your expensive work is inside the generator itself.
For example:
def generator():
for item in source:
result = expensive_processing(item)
yield result
When the current shard is replayed to reach the saved position, that expensive prefix may also be executed again.
I tested a small version of this behavior on a recent datasets revision:
consume outputs 0..4
save state
resume
The resumed outputs correctly began at 5, but the generator-side work for 0..4 was executed again while fast-forwarding.
In a second simple canary I separated the pipeline:
cheap resumable source
↓
.map(expensive_processing)
and after resume the cheap source prefix was reread, but the expensive .map() function only ran for newly emitted examples.
So if your expensive operation currently lives directly inside the generator, separating the cheap/replayable source from the heavy transform may substantially reduce restart cost.
I would treat that as a design option rather than a universal guarantee, though. My test was deliberately simple: non-batched .map(), no shuffle/interleave, no multiprocessing.
The official docs also document resume caveats for some transformation combinations, so a small kill/resume test with the exact pipeline is worthwhile before committing to a very long run.
There have also been recent fixes around repeated IterableDataset resumes; for example, a 2026 bug caused a second resume to restart from the beginning because state tracking stopped advancing after the first resume. That issue is closed and the corresponding fix is listed in the later release notes.
So for long-running jobs I would test at least:
run
→ checkpoint
→ resume
→ checkpoint again
→ resume again
rather than only testing the first resume.
One crash-safety detail that is easy to miss
If you maintain both:
- an output file, and
- a separate “source position” checkpoint,
the ordering matters.
I would prefer:
write output
↓
close/finalize output
↓
commit source progress
rather than:
commit source progress
↓
write output
because this failure:
source checkpoint says "10000 done"
↓
process crashes before output is finalized
can create a silent hole: the next run may skip source rows whose corresponding output was never safely committed.
If the output is finalized first, a crash before updating the source checkpoint tends to move the failure mode toward reprocessing/duplication, which is generally easier to detect and make idempotent with:
- deterministic output names,
- stable source IDs,
- or a small manifest.
For a simple local workflow, the existence of the finalized part itself can often act as the commit record, which avoids maintaining two independent pieces of state.
Things that look like checkpoints but I would not rely on as checkpoints
A few datasets controls are adjacent to this problem but do not provide the same guarantee.
cache_dir
This controls where builder/cache artifacts are stored. It does not turn an incomplete from_generator() build into a documented resumable Dataset.
writer_batch_size
This controls the number of examples buffered/written per Arrow batch and trades memory usage against write granularity/performance.
It is useful operationally, but I would not interpret:
writer_batch_size=1000
as:
durable checkpoint every 1000 examples
Those are different guarantees.
max_shard_size
Similarly, physical Arrow shard creation inside one builder run is not necessarily the same thing as having independently committed restart units.
.incomplete
The builder implementation uses incomplete build locations internally. I would treat these as implementation artifacts, not as an application-level recovery API.
In a small hard-kill canary I ran, an incomplete directory could remain after the process was forcibly killed, but invoking the same Dataset.from_generator() again regenerated from the beginning rather than continuing from the partial build.
A normal Python exception cleaned the incomplete build up.
After a successful build, by contrast, calling the same build again with the same cache identity reused the completed cache without re-running the generator.
That distinction seems useful:
completed cache reuse âś“
mid-build resume no supported path found
There is even a separate open issue about making a fully completed from_generator() cache easier to reopen like a normal save_to_disk() Dataset, which is another reason I would not build application-level checkpointing around private cache contents.
Why I prefer finalized files over keeping one huge Parquet writer open
If this processing can run for hours or days, I would prefer:
part-000000.parquet complete
part-000001.parquet complete
part-000002.parquet complete
part-000003.tmp current work
over one Parquet file that remains open for the entire job.
The useful recovery invariant becomes very simple:
only finalized files count.
Then a crash can normally discard the unfinished temporary part and restart from the last completed boundary.
PyArrow also exposes lower-level dataset-writing primitives if you eventually need more control over batches and file sizing, so you do not necessarily have to construct a complete Hugging Face Dataset object for every intermediate step.
I would only drop to that level if the simple per-shard/per-chunk approach becomes limiting; it probably is not necessary as the first change.
So, for your current case, my default order would be:
- If the streaming source already has useful shards, process and finalize one shard at a time.
- Otherwise, create bounded finalized output chunks yourself.
- Keep input resume state separate from output commit state.
- If replay is expensive, move the heavy operation after the cheapest resumable boundary you can get.
- Run a tiny forced-stop → resume → forced-stop → resume test before starting the full job.
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.