Expose enable_sample_tables in prepare_dataset and as --no-sample-tables CLI flag

[pei-li-hedgehog]

Closes #231.

Summary

SqliteIndexWriter already supports enable_sample_tables=False to skip creating the samples and sample_parts tables, but that knob isn't plumbed through BaseWebdatasetFactory.prepare_dataset() or the energon prepare CLI. On very large datasets (100M+ samples) the SQLite inserts and the post-insert btree builds dominate preparation time. For datasets consumed purely by the integer-indexed loader (ShardInfosITarReader), the samples table isn't queried at training time — checkpoint/resume uses SliceState integer offsets resolved via .info.json + .tar.idx.

Concretely, on a 244M-row / 4000-shard text webdataset we hit a 13-hour K8s pod timeout in the indexing phase. We need a way to skip the SQLite cost for monolithic pretraining datasets while keeping everything the integer-indexed loader needs (.tar.idx, .info.json, split config).

Changes

• BaseWebdatasetFactory.prepare_dataset() accepts enable_sample_tables=True (default — backwards compatible), passes through to SqliteIndexWriterAggregator.
• energon prepare --no-sample-tables exposes the same option. Mutually exclusive with --tar-index-only (which operates on an already-prepared dataset, different semantics).
• append_samples / append_parts become no-ops when the writer was constructed with enable_sample_tables=False, so producers can keep emitting rows unmodified.
• SqliteIndexWriter always drops stale samples / sample_parts when reset_tables=True, regardless of whether the new run intends to repopulate them. Avoids stale tables when re-preparing with --no-sample-tables after a previous full prepare.

Out of scope

Sample-key lookups (WebdatasetFileStore / SqliteITarEntryReader), polylithic joins (which INNER JOIN ... ON samples.sample_key = ... at metadataset-prepare time), and aux-data access on CrudeSample datasets all rely on the SQLite sample tables and will not work for datasets prepared with --no-sample-tables. That's the intended trade-off — documented in the new docstring/help text. Failures are loud (clear SQLite no such table errors), not silent.

A larger follow-up would also skip emitting IndexSample / IndexSamplePart items from _preprocess_tar (saving IPC overhead in addition to SQLite cost). Out of scope here.

Testing

All four existing test_prepare_dataset* tests pass (backwards compatibility).

Two new tests:

• test_prepare_dataset_no_sample_tables — covers the prepare path:

  • --no-sample-tables prepare succeeds; .info.json and per-tar .tar.idx are produced.
  • index.sqlite exists but does not contain the samples / sample_parts tables.
  • --no-sample-tables + --tar-index-only is rejected with the documented UsageError.
  • Constructing a SqliteIndexReader against the prepared dataset and calling sample-key methods (get_sample_pointer_by_key, get_sample_count, get_sample_part) raises MissingSamplesTableError with a descriptive message — not the raw sqlite3.OperationalError: no such table: samples.

• test_prepare_dataset_no_sample_tables_save_restore — covers the training-time checkpoint/resume path on a --no-sample-tables dataset:

  • Reference: uninterrupted iteration of 20 samples.
  • Save state mid-stream after 10 samples via save_state_rank(); continue iterating to capture the next 10 (post_save).
  • Build a fresh loader, restore_state_rank(state), iterate 10 samples (post_restore).
  • Asserts first_half + post_save == reference and post_restore == post_save.
  • This empirically verifies that ShardInfosITarReader + SliceState (the integer-indexed loader/state path) work end-to-end on a dataset prepared without the SQLite samples tables.

ruff check and ruff format --check clean on all touched files.

[voegtlel]

Hi @pei-li-hedgehog , thank you for this feature! Basically looks good to me, I'd just reduce the docs a bit, to be more on the spot.

[voegtlel]

One more thing: We may need to additionally add a proper error message when loading such a dataset as aux fails, i.e. if the sqlite doesn't contain the samples table, but we try to load it, explaining that the dataset was created without an index, and that it needs to be prepared accordingly.

[pei-li-hedgehog]

Thanks for the review @voegtlel! All four points addressed.

Good that you mentioned about adding a proper error message when loading such a dataset as aux fails. I've made the following change:

• New MissingSamplesTableError in indexing.py. Message names the cause (--no-sample-tables / enable_sample_tables=False) and points to the fix (re-prepare without the flag).
• SqliteIndexReader now runs a one-time check on first use of any samples-touching method. get_media_metadata is unaffected — the media_metadata table is independent.
• Test extended to cover the error type + message for three of the methods.

[voegtlel]

Just this very small edit

[pei-li-hedgehog]

Suggestions applied. Thx!

[philipp-fischer]

I'd rather vote for improving the speed of creating those sqlite tables instead of introducing a flag to skip them. Skipping them means we will have datasets that are incompatible with some of energon's feature such as energon mount.

[psycherun-kaiko]

How about doing both? I agree improving the indexing speed would be very helpful in general, but in some cases the sample-level indices might really be unnecessary.

Moreover, this path is opt-in. We could additionally make the potential failures more clear to the user via docs, or failing fast.

What do you think?