Build a Dataset One Example at a Time Without Loading All Data Into Memory
Opened by aravind-tonita (closed)

**Is your feature request related to a problem? Please describe.**
I have a very large dataset stored on disk in a custom format.  I have some custom code that reads one data example at a time and yields it in the form of a dictionary.  I want to construct a `Dataset` with all examples, and then save it to disk.  I later want to load the saved `Dataset` and use it like any other HuggingFace dataset, get splits, wrap it in a PyTorch `DataLoader`, etc.  **Crucially, I do not ever want to materialize all the data in memory while building the dataset.**

**Describe the solution you'd like**
I would like to be able to do something like the following.  Notice how each example is read and then immediately added to the dataset.  We do not store all the data in memory when constructing the `Dataset`.  If it helps, I will know the schema of my dataset before hand.
```

# Initialize an empty Dataset, possibly from a known schema.
dataset = Dataset()

# Read in examples one by one using a custom data streamer.
for example_dict in custom_example_dict_streamer("/path/to/raw/data"):

    # Add this example to the dict but do not store it in memory.
    dataset.add_item(example_dict)

# Save the final dataset to disk as an Arrow-backed dataset.
dataset.save_to_disk("/path/to/dataset")

...

# I'd like to be able to later `load_from_disk` and use the loaded Dataset
# just like any other memory-mapped pyarrow-backed HuggingFace dataset...
loaded_dataset = Dataset.load_from_disk("/path/to/dataset")
loaded_dataset.set_format(type="torch", columnns=["foo", "bar", "baz"])
dataloader = torch.utils.data.DataLoader(loaded_dataset, batch_size=16)
...

```

**Describe alternatives you've considered**
I initially tried to read all the data into memory, construct a Pandas DataFrame and then call `Dataset.from_pandas`.  This would not work as it requires storing all the data in memory.  It seems that there is an `add_item` method already -- I tried to implement something like the desired API written above, but I've not been able to initialize an empty `Dataset` (this seems to require several layers of constructing `datasets.table.Table` which requires constructing a `pyarrow.lib.Table`, etc).  I also considered writing my data to multiple sharded CSV files or JSON files and then using `from_csv` or `from_json`.  I'd prefer not to do this because (1) I'd prefer to avoid the intermediate step of creating these temp CSV/JSON files and (2) I'm not sure if `from_csv` and `from_json` use memory-mapping.

Do you have any suggestions on how I'd be able to achieve this use case?  Does something already exist to support this?  Thank you very much in advance!

Comments: Hi! Yes, the problem with `add_item` is that it keeps examples in memory, so you are left with these options:
* writing a dataset loading script in which you iterate over `custom_example_dict_streamer` and yield the examples (in `_generate examples`)
* storing the data in a JSON/CSV/Parquet/TXT file and using `Dataset.from_{format}`
* using `add_item` + `save_to_disk` on smaller chunks: 
    ```python
    from datasets import Dataset, concatenate_datasets
    MAX_SAMPLES_IN_MEMORY = 1000
    samples_in_dset = 0
    dset = Dataset.from_dict({"col1": [], "col2": []})  # empty dataset
    path_to_save_dir = "path/to/save/dir"
    num_chunks = 0
    for example_dict in custom_example_dict_streamer("/path/to/raw/data"):
        dset = dset.add_item(example_dict)
        samples_in_dset += 1
        if samples_in_dset == MAX_SAMPLES_IN_MEMORY:
            samples_in_dset = 0
            dset.save_to_disk(f"{path_to_save_dir}{num_chunks}")
            num_chunks =+ 1
            dset = Dataset.from_dict({"col1": [], "col2": []})  # empty dataset
    if samples_in_dset > 0:
        dset.save_to_disk(f"{path_to_save_dir}{num_chunks}")
        num_chunks =+ 1
    loaded_dsets = []  # memory-mapped
    for chunk_num in range(num_chunks):
        dset = Dataset.load_from_disk(f"{path_to_save_dir}{chunk_num}") 
        loaded_dsets.append(dset)
    final_dset = concatenate_datasets(dset)
    ```
    If you still have issues with this approach, you can try to delete unused datasets with `gc.collect()` to free some memory. 

This is really elegant, thank you @mariosasko!  I will try this.