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,
    sample_callback: Optional[SampleCallback] = None,
    worker_initializer: Optional[Callable[[], None]] = None,
) -> 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

        # When ``num_proc > 1`` protect the call with the standard main guard
        # (see the ``num_proc`` note below).
        if __name__ == "__main__":
            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.  Under the ``spawn``
            start method (default on macOS/Windows), the call must be protected
            by the standard main guard (``if __name__ == "__main__":``),
            otherwise it re-executes on worker import and recursively spawns
            processes.
        verbose: If True, enables verbose output during processing.
        overwrite: If True, overwrites existing output directory.
        sample_callback: Optional callback, available for the ``'cgns'`` backend,
            invoked with the written sample and its context. In parallel mode,
            it must be picklable and process-safe.
        worker_initializer: Optional callable executed once when each worker
            process starts. It must be picklable when using the ``spawn`` start
            method.
    """
    assert backend in available_backends(), (
        f"backend {backend} not among available ones: {available_backends()}"
    )
    # ---- validate the (optional) per-sample output hook ----------------------
    if sample_callback is not None and backend != "cgns":
        raise NotImplementedError(
            f"sample_callback is currently only supported for the 'cgns' "
            f"backend, got '{backend}'."
        )
    # ---- 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,
                worker_initializer=worker_initializer,
            )
        )
        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_kwargs: dict[str, Any] = {}
    if sample_callback is not None:
        backend_kwargs["sample_callback"] = sample_callback

    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,
        worker_initializer=worker_initializer,
        **backend_kwargs,
    )

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)