# Save a dataset while it's being created by from\_generator()
**URL:**
**Author:**  [@maliced](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fmaliced)%5C
**Post date:** [August 26, 2026, 7:52am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsave-a-dataset-while-its-being-created-by-from-generator%2F179254%2F1 "2026-08-26T07:52:23Z")
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.
---
**Author:**  [@John6666](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FJohn6666)%5C
**Post date:** [August 26, 2026, 12:23pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsave-a-dataset-while-its-being-created-by-from-generator%2F179254%2F2 "2026-08-26T12:23:03Z")
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:
```plaintext
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:
```plaintext
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:
>
> ```python
> 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:
>
> ```python
> 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`:
>
> ```python
> 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:
>
> ```plaintext
> 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:
>
> ```python
> 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:
>
> ```plaintext
> 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:
>
> ```plaintext
> 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:
>
> ```plaintext
> 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:
>
> 1. an output file, and
> 2. a separate “source position” checkpoint,
>
> the ordering matters.
>
> I would prefer:
>
> ```plaintext
> write output
> ↓
> close/finalize output
> ↓
> commit source progress
>
> ```
>
> rather than:
>
> ```plaintext
> commit source progress
> ↓
> write output
>
> ```
>
> because this failure:
>
> ```plaintext
> 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:
>
> ```python
> writer_batch_size=1000
>
> ```
>
> as:
>
> ```plaintext
> 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:
>
> ```plaintext
> 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:
>
> ```plaintext
> 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:
1. **If the streaming source already has useful shards, process and finalize one shard at a time.**
2. Otherwise, **create bounded finalized output chunks yourself.**
3. Keep **input resume state** separate from **output commit state**.
4. If replay is expensive, move the heavy operation _after_ the cheapest resumable boundary you can get.
5. 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.
---
**Author:**  [@thordata](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fthordata)%5C
**Post date:** [September 4, 2026, 9:12am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fsave-a-dataset-while-its-being-created-by-from-generator%2F179254%2F3 "2026-09-04T09:12:00Z")
I’d save each batch as a small shard instead of waiting for the final Dataset object. Keep a simple manifest of completed batches, so you can resume from the next one if the job stops. It’s a bit less tidy, but you won’t lose the whole run.