plaid.storage.zarr.writer

Zarr dataset writer module.

This module provides functionality for writing and managing datasets in Zarr format for the PLAID library. It includes utilities for generating datasets from sample generators, saving them to disk with optimized chunking, uploading to Hugging Face Hub, and configuring dataset cards with metadata and usage examples.

Key features: - Parallel and sequential dataset generation from generators - Automatic chunking for efficient storage - Integration with Hugging Face Hub for dataset sharing - Dataset card generation with splits, features, and documentation

Functions

write_sample(split_root, sample, var_features_keys, ...)	Write a single PLAID sample to a Zarr group on disk.
generate_datasetdict_to_disk(→ None)	Generates and saves a dataset dictionary to disk in Zarr format.
push_local_datasetdict_to_hub(→ None)	Pushes a local dataset directory to Hugging Face Hub.
configure_dataset_card(→ None)	Configures and pushes a dataset card to Hugging Face Hub for a zarr backend dataset.

Module Contents

write_sample(split_root, sample, var_features_keys, sample_counter)

Write a single PLAID sample to a Zarr group on disk.

This function serializes one Sample instance into a dedicated Zarr group under the given split root. Each sample is written as:

sample_<zero-padded index>/

Only variable features listed in var_features_keys are written. Feature paths are flattened before being used as Zarr array names.

Behavior:

• A new Zarr group named sample_{sample_counter:09d} is created.

• Each selected feature is written as a Zarr array if its value is not None.

• NumPy arrays with Unicode dtype (dtype.kind == 'U') are converted to UTF-8 encoded byte arrays to ensure stable storage (notably for Zarr v3).

• Chunk sizes are automatically determined using _auto_chunks with a target chunk size of approximately 5 million elements.

Parameters:

• split_root – Open Zarr group corresponding to a dataset split (e.g. zarr.open_group(..., mode="a")).

• sample (Sample) – PLAID Sample object to serialize.

• var_features_keys (list[str]) – List of feature paths (as defined in the variable schema) to extract and write for this sample.

• sample_counter (int) – Global index of the sample within the split, used to generate the group name and ensure deterministic ordering.

Notes

• The function assumes split_root already exists and is writable.

• No schema validation is performed at write time.

• Missing features (None values) are silently skipped.

• The function is side-effect only and returns None.

Raises:

• zarr.errors.ContainsGroupError – If a sample group with the same name already exists.

• OSError / IOError – If an underlying filesystem or Zarr write error occurs.

generate_datasetdict_to_disk(output_folder: str | pathlib.Path, generators: dict[str, Callable[Ellipsis, Generator[plaid.Sample, None, None]]], variable_schema: dict[str, dict], gen_kwargs: dict[str, dict[str, list[plaid.types.IndexType]]] | None = None, num_proc: int = 1, verbose: bool = False) → None

Generates and saves a dataset dictionary to disk in Zarr format.

This function processes sample generators for different dataset splits, converts samples to dictionaries, and writes them to Zarr arrays on disk. It supports both sequential and parallel processing modes. In parallel mode, gen_kwargs must be provided with batch information for each split.

Parameters:

• output_folder (Union[str, Path]) – Base directory where the dataset will be saved. A ‘data’ subdirectory will be created inside this folder.

• generators (dict[str, Callable[..., Generator[Sample, None, None]]]) – Dictionary mapping split names (e.g., “train”, “test”) to generator functions that yield Sample objects.

• variable_schema (dict[str, dict]) – Schema describing the structure and types of variables/features in the samples.

• gen_kwargs (Optional[dict[str, dict[str, list[IndexType]]]]) – Optional generator arguments for parallel processing. Must include “shards_ids” for each split when num_proc > 1. Required for parallel execution.

• num_proc (int, optional) – Number of processes to use for parallel processing. Defaults to 1 (sequential). Must be > 1 only when gen_kwargs is provided.

• verbose (bool, optional) – Whether to display progress bars during processing. Defaults to False.

Returns:

This function does not return a value; it writes the dataset directly

to disk.

Return type:

None

push_local_datasetdict_to_hub(repo_id: str, local_dir: str | pathlib.Path, num_workers: int = 1) → None

Pushes a local dataset directory to Hugging Face Hub.

This function uploads the contents of a local directory to a specified Hugging Face repository as a dataset. It uses the HfApi to handle large folder uploads with configurable parallelism.

Parameters:

• repo_id (str) – The Hugging Face repository ID where the dataset will be uploaded (e.g., “username/dataset_name”).

• local_dir (str or Path) – Path to the local directory containing the dataset files to upload.

• num_workers (int, optional) – Number of worker threads to use for uploading. Defaults to 1.

Returns:

This function does not return a value; it uploads the dataset directly

to Hugging Face Hub.

Return type:

None

configure_dataset_card(repo_id: str, infos: dict[str, dict[str, str]], local_dir: str | pathlib.Path, viewer: bool | None = None, pretty_name: str | None = None, dataset_long_description: str | None = None, illustration_urls: list[str] | None = None, arxiv_paper_urls: list[str] | None = None) → None

Configures and pushes a dataset card to Hugging Face Hub for a zarr backend dataset.

This function generates a dataset card in YAML format with metadata, features, splits information, and usage examples. It automatically detects splits and sample counts from the local directory structure, then pushes the card to the specified Hugging Face repository.

Parameters:

• repo_id (str) – The Hugging Face repository ID where the dataset card will be pushed.

• infos (dict[str, dict[str, str]]) – Dictionary containing dataset metadata, including legal information like license.

• local_dir (Union[str, Path]) – Path to the local directory containing the dataset files, expected to have a ‘data’ subdirectory with split folders.

• variable_schema (Optional[dict]) – Schema describing the variables/features in the dataset, used to generate the features section in the card.

• viewer (Optional[bool]) – Unused parameter for viewer configuration.

• pretty_name (Optional[str]) – A human-readable name for the dataset to display in the card.

• dataset_long_description (Optional[str]) – A detailed description of the dataset to include in the card.

• illustration_urls (Optional[list[str]]) – List of URLs to images that illustrate the dataset, displayed in the card.

• arxiv_paper_urls (Optional[list[str]]) – List of arXiv URLs for papers related to the dataset, included as sources.

Returns:

This function does not return a value; it pushes the dataset card

directly to Hugging Face Hub.

Return type:

None