Skip to content

feat(py + cli): download datasets as files by default, add --zip flag#20

Open
ravivooda wants to merge 1 commit into
mainfrom
ravivooda/dataset-download-no-zip
Open

feat(py + cli): download datasets as files by default, add --zip flag#20
ravivooda wants to merge 1 commit into
mainfrom
ravivooda/dataset-download-no-zip

Conversation

@ravivooda

Copy link
Copy Markdown
Collaborator

What

lightning dataset download now downloads a dataset version's files into a directory by default, and only produces a .zip when the new --zip flag is passed.

Why

The command always packaged the version into a .zip via shutil.make_archive and then deleted the raw files. Two problems:

  • The deflate step dominates wall time on real datasets and provides little benefit on already-incompressible data.
  • Users get a .zip they must unzip again before the data is usable — there was no way to get the plain files.

Change

  • download_dataset (CLI): output is now <target-path>/<dataset>_<version>/ by default; --zip writes <target-path>/<dataset>_<version>.zip (previous behaviour).
  • _download_dataset_version: new as_zip: bool = False kwarg. When False, files are written straight into target_path; when True, the previous temp-dir + make_archive path is used. Backward compatible (new kwarg defaults off).

Test results

Measured against a single-file 819 MB dataset (MLP_input_output.npy), same payload, cold each run:

Mode Output Wall time
default (--no-zip) directory of files 22.9 s
--zip single .zip (760 MB) 47.7 s

Breakdown of the old (--zip) path: ~14.5 s single-stream download + ~34.3 s make_archive deflate. Dropping the forced zip removes that ~34 s and yields directly-usable files. Downloaded bytes verified byte-identical in both modes (MD5 match).

Note: the remaining time is a single urllib.urlretrieve per file (sequential, one stream). Parallel / chunked downloads are a natural follow-up, not in this PR.

Unit tests

  • test_download_dataset_version / test_download_dataset_version_no_token_no_cluster — updated to pass as_zip=True, keeping coverage of the archive path.
  • test_download_dataset_version_default_downloads_files_no_zipnew: asserts the default does not call make_archive and writes files into the target directory.
  • test_dataset_download_help / test_datasets_download_help — updated for the new help text; assert --zip is documented.

All 7 pass.

🤖 Generated with Claude Code

`lightning dataset download` always packaged the version into a .zip via
shutil.make_archive, deleting the raw files afterwards. The deflate step
dominates wall time on large datasets (~34s of a 48s / 819MB download) and
forces users to unzip again before use.

Default to downloading the files straight into a directory and gate the
archive behind an opt-in `--zip` flag. The `as_zip` option is threaded
through the CLI, `datasets.download_dataset`, and `_download_dataset_version`
(default False); when set it keeps the previous temp-dir + make_archive
behaviour.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@ravivooda
ravivooda force-pushed the ravivooda/dataset-download-no-zip branch from 937da6a to 136a797 Compare July 15, 2026 17:56
@ravivooda
ravivooda marked this pull request as ready for review July 16, 2026 23:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant