Shard parquet in download_and_prepare

[lhoestq]

Following #4724 (needs to be merged first)

It's good practice to shard parquet files to enable parallelism with spark/dask/etc.

I added the max_shard_size parameter to download_and_prepare (default to 500MB for parquet, and None for arrow).

from datasets import *

output_dir = "./output_dir"  # also supports "s3://..."
builder = load_dataset_builder("squad")
builder.download_and_prepare(output_dir, file_format="parquet", max_shard_size="5MB")

Implementation details

The examples are written to a parquet file until ParquetWriter._num_bytes > max_shard_size. When this happens, a new writer is instantiated to start writing the next shard. At the end, all the shards are renamed to include the total number of shards in their names: {builder.name}-{split}-{shard_id:05d}-of-{num_shards:05d}.parquet

I also added the MAX_SHARD_SIZE config variable (default to 500MB)

TODO:

• docstrings
• docs
• tests

cc @severo

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[lhoestq]

This is ready for review cc @mariosasko :) please let me know what you think !

[mariosasko]

Thanks!

[mariosasko]

The final shard size can easily be off (significantly) using this logic, no? By default, writer._num_bytes is only updated every 10k examples (WRITER_BATCH_SIZE), so if the max_shard_size is small...

[lhoestq]

I think right now this is the builder's responsibility to specify if 10,000 is too much (DEFAULT_WRITER_BATCH_SIZE class attribute).

Though I agree this is not ideal. I think it's ok to have it this way in this PR (since it would be off if the max_shard_size is very very small), but it would be nice to have something smarter in general

[mariosasko]

One nit. Besides that, all looks good!