Title: better correspondence between cached and saved datasets created using from_generator
State: open
Labels: description: New feature or request
name: enhancement

### Feature request

At the moment `.from_generator` can only create a dataset that lives in the cache. The cached dataset cannot be loaded with `load_from_disk` because the cache folder is missing `state.json`. So the only way to convert this cached dataset to a regular is to use `save_to_disk` which needs to create a copy of the cached dataset. For large datasets this can end up wasting a lot of space. In my case the saving operation failed so I am stuck with a large cached dataset and no clear way to convert to a `Dataset` that I can use. The requested feature is to provide a way to be able to load a cached dataset using `.load_from_disk`. Alternatively `.from_generator` can create the dataset at a specified location so that it can be loaded from there with `.load_from_disk`.

### Motivation

I have the following workflow which has exposed some awkwardness about the Datasets saving/caching.

1. I created a cached dataset using `.from_generator` which was cached in a folder. This dataset is rather large (~600GB) with many shards.
2. I tried to save this dataset using `.save_to_disk` to another location so that I can use later as a `Dataset`. This essentially creates another copy (for a total of 1.2TB!) of what is already in the cache... In my case the saving operation keeps dying for some reason and I am stuck with a cached dataset and no copy.
3. Now I am trying to "save" the existing cached dataset but it is not clear how to access the cached files after `.from_generator` has finished e.g. from a different process. I should not be even looking at the cache but I really do not want to waste another 2hr to generate the set so that if fails agains (I already did this couple of times). 
- I tried `.load_from_disk` but it does not work with cached files and complains that this is not a `Dataset` (!).
- I looked at `.from_file` which takes one file but the cached file has many (shards) so I am not sure how to make this work. 
- I tried `.load_dataset` but this seems to either try to "download" a copy (of a file which is already in the local file system!) which I will then need to save or I need to use `streaming=False` to create an `IterableDataset `which then I need to convert (using the cache) to `Dataset` so that I can save it. With both options I  will end up with 3 copies of the same dataset for a total of ~2TB! I am hoping here is another way to do this...

Maybe I am missing something here: I looked at docs and forums but no luck. I have a bunch of arrow files cached by `Dataset.from_generator` and no clean way to make them into a `Dataset` that I can use. 

This all could be so much easer if `load_from_disk` can recognize the cached files and produce a `Dataset`: after the cache is created I would not have to "save" it again and I can just load it when I need.  At the moment `load_from_disk` needs `state.json` which is lacking in the cache folder. So perhaps `.from_generator` could be made to "finalize" (e.g. create `state.json`) the dataset once it is done so that it can be loaded easily. Or provide `.from_generator` with a `save_to_dir` parameter in addition to `cache_dir` which can be used for the whole process including creating the `state.json` at the end. 

As a proof of concept I just created `state.json` by hand and `load_from_disk` worked using the cache! So it seems to be the missing piece here.

### Your contribution

Time permitting I can look into `.from_generator` to see if adding  `state.json` is feasible.

Comments:
Until this lands, in case anyone needs this, I made this helper (based on HF sources) that merely serializes the metadata from a dataset, thus allowing to bypass the completely redundant additional save.

```python
import json
import posixpath
import os
from typing import Optional
from fsspec.core import url_to_fs
from datasets.utils.typing import PathLike

from datasets import Dataset, config
from datasets.info import asdict

def save_dataset_metadata(
    dataset: Dataset,
    dataset_path: PathLike,
    storage_options: Optional[dict] = None
):
    """
    Saves only the dataset metadata (state.json and dataset_info.json) to a dataset directory.

    This does NOT rewrite or move any Arrow shard files.
    It reuses the existing data files referenced by `dataset.cache_files`.

    Args:
        dataset (`datasets.Dataset`):
            The in-memory Dataset instance whose current cache-backed data files
            should be exposed as a load_from_disk()-compatible dataset.
        dataset_path (`path-like`):
            Path (e.g. `dataset/train`) or remote URI (e.g. `s3://my-bucket/dataset/train`)
            of the dataset directory where the metadata will be saved to.
        storage_options (`dict`, *optional*):
            Key/value pairs to be passed on to the file-system backend, if any.
    """
    fs, _ = url_to_fs(dataset_path, **(storage_options or {}))
    fs.makedirs(dataset_path, exist_ok=True)

    # Build json-serializable state (copied from HF save_to_disk, but without shard writing)
    state = {
        key: dataset.__dict__[key]
        for key in [
            "_fingerprint",
            "_format_columns",
            "_format_kwargs",
            "_format_type",
            "_output_all_columns",
        ]
    }
    state["_split"] = str(dataset.split) if dataset.split is not None else dataset.split

    # Relies on `dataset.cache_files` to point at the real Arrow shards.
    state["_data_files"] = [
        {"filename": os.path.relpath(cache_filename["filename"], dataset_path)} for cache_filename in dataset.cache_files
    ]

    # Same json-serializability check as HF
    for k in state["_format_kwargs"].keys():
        json.dumps(state["_format_kwargs"][k])

    # Dataset info (same as HF)
    dataset_info = asdict(dataset._info)

    # Write state.json
    with fs.open(
        posixpath.join(dataset_path, config.DATASET_STATE_JSON_FILENAME), "w", encoding="utf-8"
    ) as state_file:
        json.dump(state, state_file, indent=2, sort_keys=True)

    # Write dataset_info.json
    with fs.open(
        posixpath.join(dataset_path, config.DATASET_INFO_FILENAME), "w", encoding="utf-8"
    ) as dataset_info_file:
        # Sort only the top-level keys, like HF does
        sorted_keys_dataset_info = {key: dataset_info[key] for key in sorted(dataset_info)}
        json.dump(sorted_keys_dataset_info, dataset_info_file, indent=2)
```

Usage:

```python
from .save_dataset_metadata import save_dataset_metadata

# ...

dataset = datasets.Dataset.from_generator(..., cache_dir="xxx")

# ...

save_dataset_metadata(dataset, Path(directory_path))

```

@lhoestq Do you think this issue could get some traction to be resolved? Currently if you have a large dataset created in cache and need to persist it, I don't think it's great that you have to use double disk space and hours on re-serializing the dataset, when in reality only the metadata is missing. Imo Datasets should be able to persist metadata only in this case.

Thanks for considering.

Have you tried using pickle ?

```python
ds = Dataset.from_genetator(...)  # generates and caches the dataset
with open("dataset.pickle", "wb") as f:
    pickle.dump(ds, f)  # saves the Dataset obj including the cache files locations (and not the data! unless some of them live in memory because of add_column or add_item for example)

# later
with open("dataset.pickle", "rb") as f:
    ds = pickle.load(ds, f)  # reload the Dataset obj from the cache files
```

I'm not a big fan of using pickle for this kind of use cases but it should do the job.

Anyways, moving forward it would be cool to change how save_to_disk / load_from_disk work to make them compatible with load_dataset. This implies replacing the state.json file with the usual README.md which can list the data files per config/split. save_to_disk would still materialize the full dataset on disk in the output directory though.

The difference would be load_from_disk would use the dataset directory as cache directory to keep it self-contained and avoid copies, contrary to load_dataset which copies the data to a directory per dataset inside `~/.cache/huggingface/datasets/`. Contributions on this are welcome on this subject btw !