Title: Add the option of saving in parquet instead of arrow 
Author: arita37 (NONE)
State: open
Created: 2024-05-16T13:35:51Z
Labels: enhancement

### 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 !

Comments:
I think [`Dataset.to_parquet`](https://huggingface.co/docs/datasets/v1.10.2/package_reference/main_classes.html#datasets.Dataset.to_parquet) is what you're looking for.

Let me know if I'm wrong 

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.


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

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.


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

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.


`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

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.


can `to_parquet` do auto sharding ?

Not at the moment, but you can shard yourself using

```python
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")
```

Thanks! That's good!

@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).

`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` ?

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.

For `map()` it works a similar way: it writes batch by batch. You can lower the `writer_batch_size=` in `map()` (default is 1000)
Feel free to open a new issue to discuss RAM usage in `save_to_disk()` and `map()`, ideally if you can provide a minimal code that reproduces your issue

OK, could it be because I'm loading the initial dataset from GCS?

I understand how it should work  i.e. not balloon from 800mb to 15gb when call .save_to_disk or at end of map. Its loading the full dataset/rematerializing, and not respecting the shard_size or writer_batch_size. When i use shard_size 100mb, the shards are indeed 100mg, just peak memory is the full dataset.

Ok I'd have to create a very large public dataset to recreate it. 

I ended up manually sharding chunks to disk and using 
```
datasets.load_dataset(
                "parquet",
                data_files=f"{dataset_path}/*.parquet")
```
This is less than ideal as it means it is very slow loading/indexing the first time. But the RSS remains as expected, only a bit larger than the chunks. My use case is I have one k8s service preparing the large dataset, and another to run inference.

I think it could be related to this: https://stasosphere.com/entrepreneur-being/301-mmap-memory-leak-investigation/
The RSS usage still looks like its growing and growing even though its evictable.. And k8s monitors it so crashes the pod?
@lhoestq Have you seen this problem before of using large datasets in production.


It's also possible to encounter that issue if you're using a special kind of disk that has a caching mechanism that loads memory mapped files in RAM.