plaid.storage.writer¶

def save_to_disk(
    output_folder: Union[str, Path],
    sample_constructor: Callable[[Any], Sample],
    ids: Mapping[str, Any],
    infos: Optional[Infos] = None,
    backend: str = "hf_datasets",
    pb_defs: Optional[dict[str, ProblemDefinition]] = None,
    num_proc: int = 1,
    verbose: bool = False,
    overwrite: bool = False,
) -> None:
    """Save a PLAID dataset to local disk using the specified backend.

    This function preprocesses the dataset, extracts schemas, and saves
    the dataset to disk using the chosen backend.  It also saves metadata, infos,
    and problem definitions.

    The user provides a simple function ``sample_constructor`` that takes a single
    identifier and returns a :class:`~plaid.Sample`, together with a dictionary
    ``ids`` mapping split names to sliceable sequences of identifiers.
    PLAID handles iteration, generator creation, and parallel sharding
    internally.

    Example::

        from plaid import Sample
        from plaid.storage import save_to_disk

        def sample_constructor(file_path):
            sample = Sample()
            sample.add_tree(load_my_data(file_path))
            return sample

        save_to_disk(
            "output/",
            sample_constructor=sample_constructor,
            ids={
                "train": train_file_paths,
                "test":  test_file_paths,
            },
            infos=Infos(
                owner="owner",
                license="license",
            ),
            num_proc=6,
        )

    Args:
        output_folder: Path to the output directory where the dataset will be saved.
        sample_constructor: A callable that takes a single identifier (of any type)
            and returns a :class:`~plaid.Sample`.
        ids: Dictionary mapping split names (e.g. ``"train"``, ``"test"``) to
            sliceable sequences of sample identifiers.  Each sequence must
            support ``__getitem__`` and ``__len__`` (list, tuple, numpy array,
            …).  The identifiers can be of any type: integers, file paths,
            strings, tuples, etc.
        backend: Storage backend to use (``'cgns'``, ``'hf_datasets'``, or
            ``'zarr'``).
        infos: Dataset information to save with the dataset. If ``None``, a
            placeholder :class:`~plaid.Infos` is created with
            ``owner="unknown", license="unknown"``.
        pb_defs: Optional mapping from problem definition identifiers to definitions.
        num_proc: Number of processes to use for parallel writing.  When
            ``num_proc > 1`` PLAID automatically shards the identifier
            sequences and distributes work across workers.
        verbose: If True, enables verbose output during processing.
        overwrite: If True, overwrites existing output directory.
    """
    assert backend in available_backends(), (
        f"backend {backend} not among available ones: {available_backends()}"
    )
    # ---- validate ids: must be sliceable sequences ---------------------------
    for split_name, split_ids in ids.items():
        if not (hasattr(split_ids, "__getitem__") and hasattr(split_ids, "__len__")):
            raise TypeError(
                f"ids for split '{split_name}' must be a sliceable sequence "
                f"(with __getitem__ and __len__), got {type(split_ids).__name__}. "
                f"Use a list, tuple, or numpy array of sample identifiers."
            )

    # ---- build generators from sample_constructor -----------------------------------
    generators: dict[str, Callable[..., Generator[Sample, None, None]]] = {}
    for split_name in ids:
        generators[split_name] = _SampleFuncGenerator(sample_constructor)

    # ---- auto-shard when running in parallel ---------------------------------
    gen_kwargs: Optional[dict[str, dict[str, Any]]] = None
    if num_proc > 1:
        gen_kwargs = _build_gen_kwargs(ids, num_proc)
    else:
        # For sequential execution, wrap ids into a single shard so the
        # generator receives them via shards_ids
        gen_kwargs = {
            split_name: {"shards_ids": [list(split_ids)]}
            for split_name, split_ids in ids.items()
        }

    output_folder = Path(output_folder)
    _check_folder(output_folder, overwrite)

    # CGNS stores each sample as a complete CGNS tree. It does not need the
    # constant/variable split used by columnar backends, so avoid the expensive
    # preprocessing pass that flattens every generated tree and detects constant
    # leaves.  Sample counts can be derived directly from the declared ids.
    if backend == "cgns":
        variable_schema = None
        num_samples = {
            split_name: len(split_ids) for split_name, split_ids in ids.items()
        }
    else:
        flat_cst, variable_schema, constant_schema, num_samples, cgns_types = (
            preprocess(
                generators, gen_kwargs=gen_kwargs, num_proc=num_proc, verbose=verbose
            )
        )
        save_metadata_to_disk(
            output_folder, flat_cst, variable_schema, constant_schema, cgns_types
        )

    # Inject the actual on-disk storage backend and sample counts so the
    # written ``infos.yaml`` always reflects how the dataset was saved,
    # overriding any inherited values from the input ``infos``.
    if infos is None:
        infos = Infos(
            owner="unknown",
            license="unknown",
        )
    infos_data = infos.model_dump(exclude_none=True)
    infos_data["num_samples"] = num_samples
    infos_data["storage_backend"] = backend
    infos = Infos.validate_persisted(infos_data)

    save_infos_to_disk(output_folder, infos)

    if pb_defs is not None:
        save_problem_definitions_to_disk(output_folder, pb_defs)

    backend_spec = get_backend(backend)
    backend_spec.generate_to_disk(
        output_folder,
        generators,
        variable_schema,
        gen_kwargs=gen_kwargs,
        num_proc=num_proc,
        verbose=verbose,
    )

def push_to_hub(
    repo_id: str,
    local_dir: Union[str, Path],
    num_workers: int = 1,
    viewer: bool = False,
    pretty_name: Optional[str] = None,
    dataset_long_description: Optional[str] = None,
    illustration_urls: Optional[list[str]] = None,
    arxiv_paper_urls: Optional[list[str]] = None,
) -> None:  # pragma: no cover
    """Push a local PLAID dataset to Hugging Face Hub.

    This function uploads a previously saved dataset from local disk to Hugging Face Hub,
    including data, metadata, infos, and problem definitions. It automatically detects the
    backend used for saving and configures the dataset card.

    Args:
        repo_id: Hugging Face repository ID (e.g., 'username/dataset-name').
        local_dir: Local directory containing the saved dataset.
        num_workers: Number of workers for parallel upload.
        viewer: If True, enables dataset viewer on Hub.
        pretty_name: Optional pretty name for the dataset.
        dataset_long_description: Optional detailed description.
        illustration_urls: Optional list of illustration URLs.
        arxiv_paper_urls: Optional list of arXiv paper URLs.
    """
    infos = load_infos_from_disk(local_dir)

    backend = infos.storage_backend
    if backend is None:
        raise ValueError("Missing 'storage_backend' in persisted infos")

    backend_spec = get_backend(backend)
    backend_spec.push_local_to_hub(repo_id, local_dir, num_workers=num_workers)
    backend_spec.configure_dataset_card(
        repo_id,
        infos,
        local_dir,
        # variable_schema,
        viewer,
        pretty_name,
        dataset_long_description,
        illustration_urls,
        arxiv_paper_urls,
    )

    if backend != "cgns":
        push_local_metadata_to_hub(repo_id, local_dir)

    push_infos_to_hub(repo_id, infos)

    push_local_problem_definitions_to_hub(repo_id, local_dir)