Add the option of saving in parquet instead of arrow

[arita37]

Feature request

In dataset.save_to_disk('/path/to/save/dataset'),

add the option to save in parquet format

dataset.save_to_disk('/path/to/save/dataset', format="parquet"),

because arrow is not used for Production Big data.... (only parquet)

Motivation

because arrow is not used for Production Big data.... (only parquet)

Your contribution

I can do the testing !

[Dref360]

I think Dataset.to_parquet is what you're looking for.

Let me know if I'm wrong

[arita37]

No, it does not save the metadata json. We have to recode all meta json load/save with another custome functions. save_to_disk and load should have option with “Parquet” instead of “arrow” since “arrow” is never user for production (only parquet). Thanks !

…

On May 17, 2024, at 5:38, Frédéric Branchaud-Charron ***@***.***> wrote:  I think Dataset.to_parquet is what you're looking for. Let me know if I'm wrong — Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because you authored the thread.

[lhoestq]

You can use to_parquet and ds.info.write_to_directory() to save the dataset info

[arita37]

Ok, What about loading ? Should we do in 2 steps ?

…

On Jun 14, 2024, at 1:09, Quentin Lhoest ***@***.***> wrote:  You can use to_parquet and ds.info.write_to_directory() to save the dataset info — Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because you authored the thread.

[lhoestq]

Yes, and there is DatasetInfo.from_directory(). to reload the info

[arita37]

Isn’t easier to combine both into load_dataset and save_dataset with parquet options. 2) another question, How can we download large dataset into disk directly without loading all in memory (!)

…

On Jun 14, 2024, at 19:54, Quentin Lhoest ***@***.***> wrote:  Yes, and there is DatasetInfo.from_directory(). to reload the info — Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because you authored the thread.

[lhoestq]

load_dataset doesn't load the dataset in memory, it progressively writes to disk in Arrow format and then memory maps the Arrow files. This allows to load datasets bigger than memory and without filling your RAM

[arita37]

Sure. How memory map is managed ? Managed by the OS ? Why the need of save_dataset() ?

…

On Jun 15, 2024, at 0:06, Quentin Lhoest ***@***.***> wrote:  load_dataset doesn't load the dataset in memory, it progressively writes to disk in Arrow format and then memory maps the Arrow files. This allows to load datasets bigger than memory and without filling your RAM — Reply to this email directly, view it on GitHub, or unsubscribe. You are receiving this because you authored the thread.

[Adenialzz]

can to_parquet do auto sharding ?

[lhoestq]

Not at the moment, but you can shard yourself using

num_shards = 32
for index in range(num_shards):
    shard = ds.shard(index=index, num_shards=num_shards)
    shard.to_parquet("data-{index:05d}.parquet")

[Adenialzz]

Thanks! That's good!

[jamesharrisivi]

@lhoestq

Related to this, I'm doing processing on a very large dataset and I save it in shards, when I load the parquet shards it takes time but doesn't load it all in memory and reindexes it as one dataset which is what I want, so I can store in the cloud and its ready for lazy loading in inference.
dataset = datasets.load_dataset(
"parquet",
data_files=f"dataset_parquet/*.parquet",
keep_in_memory=False,
storage_options=storage_options,
)

But if I call save_to_disk it rematerializes the whole dataset in RAM (blowing up memory usage and causing OOM)? Is there a way to save a very large dataset so it doesn't need to be reindexed like above. I feel this should be possible as once it's loaded it's cached and I can see cached arrow (but it cant be transported).

[lhoestq]

save_to_disk iterates on chunks from the data to write them to disk shard by shard, it doesn't load the dataset in memory.

Maybe try to lower the max_shard_size ?

[jamesharrisivi]

Ok, is it a bug with my particular dataset.
I do max_shard_size="500Mb" but still the RAM usage starts increasing to the full dataset even though the .arrow chunks are 500mb.
This also happens at the end of my map function (which maps to a much larger number of rows), during tqdm the % until 100 the RAM is constant then rematerialize.
I know it must be possible as if i do
dataset = datasets.load_dataset(
"parquet",
data_files=f"dataset_parquet/*.parquet",
keep_in_memory=False,
storage_options=storage_options,
)
it takes a long time but the second time it is fast, and no RAM increase, but I know I cant transfer the cached folder for a production inference server.