plaid.storage.zarr.reader¶

def sample_generator(
    repo_id: str, split: str, ids: Iterable[int], selected_features: list[str]
) -> Iterator[dict[str, Any]]:  # pragma: no cover
    """Generates samples from a Zarr dataset on Hugging Face Hub.

    Args:
        repo_id (str): The Hugging Face repository ID.
        split (str): The dataset split name.
        ids (Iterable[int]): Iterable of sample IDs to generate.
        selected_features (list[str]): List of features to include.

    Yields:
        dict: Dictionary mapping feature names to Zarr arrays.
    """
    base_url = (
        f"https://huggingface.co/datasets/{repo_id}/resolve/main/data/{split}/sample_"
    )
    for idx in ids:
        sample = {}
        for feat in selected_features:
            flatten_feat = flatten_path(feat)
            mapper = fsspec.get_mapper(base_url + f"{idx:09d}/{flatten_feat}")
            sample[feat] = zarr.open(mapper, mode="r")

        yield sample

def create_zarr_iterable_dataset(
    repo_id: str, split: str, ids: Iterable[int], selected_features: list[str]
) -> IterableDataset:  # pragma: no cover
    """Creates an IterableDataset from Zarr data on Hugging Face Hub.

    Args:
        repo_id (str): The Hugging Face repository ID.
        split (str): The dataset split name.
        ids (Iterable[int]): Iterable of sample IDs.
        selected_features (list[str]): List of features to include.

    Returns:
        IterableDataset: An iterable dataset for streaming access.
    """

    def wrapped_gen():
        yield from sample_generator(repo_id, split, ids, selected_features)

    return IterableDataset.from_generator(
        wrapped_gen,
        split=NamedSplit(split),
        features=None,
    )

def init_datasetdict_from_disk(
    path: Union[str, Path],
) -> ZarrDatasetDict:
    """Initializes dataset dictionaries from local Zarr files.

    Args:
        path (Union[str, Path]): Path to the local directory containing the dataset.

    Returns:
        ZarrDatasetDict: Dictionary mapping split names to ZarrDataset objects.
    """
    local_path = Path(path) / "data"
    split_names = [p.name for p in local_path.iterdir() if p.is_dir()]
    return {
        sn: ZarrDataset(
            zarr.open(zarr.storage.LocalStore(local_path / sn), mode="r"), path
        )
        for sn in split_names
    }

def download_datasetdict_from_hub(
    repo_id: str,
    local_dir: Union[str, Path],
    split_ids: Optional[dict[str, Iterable[int]]] = None,
    features: Optional[list[str]] = None,
    overwrite: bool = False,
) -> str:  # pragma: no cover
    """Downloads dataset from Hugging Face Hub to local directory.

    Args:
        repo_id (str): The Hugging Face repository ID.
        local_dir (Union[str, Path]): Local directory to download to.
        split_ids (Optional[dict[str, Iterable[int]]]): Optional split IDs for selective download.
        features (Optional[list[str]]): Optional features for selective download.
        overwrite (bool): Whether to overwrite existing directory.

    Returns:
        str: Path to the local directory where the dataset has been downloaded.
    """
    output_folder = Path(local_dir)

    if output_folder.is_dir():
        if overwrite:
            shutil.rmtree(local_dir)
            logger.warning(f"Existing {local_dir} directory has been reset.")
        elif any(local_dir.iterdir()):
            raise ValueError(
                f"directory {local_dir} already exists and is not empty. Set `overwrite` to True if needed."
            )

    allow_patterns, ignore_patterns = _zarr_patterns(repo_id, split_ids, features)

    return snapshot_download(
        repo_id=repo_id,
        repo_type="dataset",
        allow_patterns=allow_patterns,
        ignore_patterns=ignore_patterns,
        local_dir=local_dir,
    )

def init_datasetdict_streaming_from_hub(
    repo_id: str,
    split_ids: Optional[dict[str, Iterable[int]]] = None,
    features: Optional[list[str]] = None,
) -> dict[str, IterableDataset]:  # pragma: no cover
    """Initializes streaming dataset dictionaries from Hugging Face Hub.

    This function creates iterable datasets that stream Zarr data directly from
    the Hugging Face Hub without downloading files locally. It supports selective
    loading of specific splits and features for memory-efficient data access.
    Note that streaming mode is not compatible with private Hugging Face mirrors.

    Args:
        repo_id (str): The Hugging Face repository ID (e.g., "username/dataset_name").
        split_ids (Optional[dict[str, Iterable[int]]]): Optional dictionary mapping split names
            to lists of sample IDs to include. If None, all samples from all splits
            are included.
        features (Optional[list[str]]): Optional list of feature names to include.
            If None, all features from the variable schema are included.

    Returns:
        dict[str, IterableDataset]: Dictionary mapping split names to IterableDataset
            objects for streaming data access.
    """
    hf_endpoint = os.getenv("HF_ENDPOINT", "").strip()
    if hf_endpoint:
        raise RuntimeError("Streaming mode not compatible with private mirror.")

    if features is not None:
        selected_features = features
    else:
        yaml_path = hf_hub_download(
            repo_id=repo_id,
            filename="variable_schema.yaml",
            repo_type="dataset",
        )
        with open(yaml_path, "r", encoding="utf-8") as f:
            variable_schema = yaml.safe_load(f)
        selected_features = list(variable_schema.keys())

    if split_ids is not None:
        selected_ids = split_ids
    else:
        infos = load_infos_from_hub(repo_id=repo_id)
        selected_ids = {
            split: range(n_samples) for split, n_samples in infos.num_samples.items()
        }

    return {
        split: create_zarr_iterable_dataset(repo_id, split, ids, selected_features)
        for split, ids in selected_ids.items()
    }