create_zarr_arrays

create_zarr_arrays

Shared zarr replay buffer creation utilities.

Provides a unified pipeline for creating zarr stores from any dataset schema, with WebP compression for uint8 images and lz4 for numerical arrays.

create_zarr_arrays

create_zarr_arrays(data_group, schema, image_codec, numeric_compressor, image_frames_per_shard=DEFAULT_IMAGE_FRAMES_PER_SHARD)

Create zarr arrays with codec selection based on data type.

Uses WebP compression for uint8 image arrays and Blosc/lz4 for numerical arrays.

Parameters:

Name	Type	Description	Default
data_group	Group	Zarr group to create arrays in.	required
schema	DatasetSchema	Dataset schema providing array specifications.	required
image_codec	WebPCodec	WebP codec for uint8 image arrays.	required
numeric_compressor	BloscCodec	Blosc codec for numerical arrays.	required
image_frames_per_shard	int | None	Number of independently compressed image frames stored in each shard, or None to disable sharding.	DEFAULT_IMAGE_FRAMES_PER_SHARD

Raises:

Type	Description
TypeError	If image_frames_per_shard is not an integer.
ValueError	If image_frames_per_shard is not a positive multiple of the image chunk length.

Source code in src/versatil/data/preprocessing/create_zarr_arrays.py

def create_zarr_arrays(
    data_group: zarr.Group,
    schema: DatasetSchema,
    image_codec: WebPCodec,
    numeric_compressor: BloscCodec,
    image_frames_per_shard: int | None = DEFAULT_IMAGE_FRAMES_PER_SHARD,
) -> None:
    """Create zarr arrays with codec selection based on data type.

    Uses WebP compression for uint8 image arrays and Blosc/lz4 for
    numerical arrays.

    Args:
        data_group: Zarr group to create arrays in.
        schema: Dataset schema providing array specifications.
        image_codec: WebP codec for uint8 image arrays.
        numeric_compressor: Blosc codec for numerical arrays.
        image_frames_per_shard: Number of independently compressed image frames
            stored in each shard, or None to disable sharding.

    Raises:
        TypeError: If image_frames_per_shard is not an integer.
        ValueError: If image_frames_per_shard is not a positive multiple of the
            image chunk length.
    """
    specs = schema.get_zarr_array_specs()
    for key, spec in specs.items():
        dtype = str if spec["dtype"] == "str" else getattr(np, spec["dtype"])
        if is_uint8_image_spec(spec):
            chunks = (1, *spec["shape"][1:])
            shards = get_image_shard_shape(
                image_chunks=chunks,
                image_frames_per_shard=image_frames_per_shard,
            )
            data_group.create_array(
                name=key,
                shape=spec["shape"],
                chunks=chunks,
                shards=shards,
                dtype=dtype,
                serializer=image_codec,
                compressors=None,
            )
        else:
            data_group.create_array(
                name=key,
                shape=spec["shape"],
                chunks=spec["chunks"],
                dtype=dtype,
                compressors=[numeric_compressor] if spec["needs_compressor"] else None,
            )

create_zarr_replay_buffer

create_zarr_replay_buffer(schema, episodes, total_episodes=None, image_frames_per_shard=DEFAULT_IMAGE_FRAMES_PER_SHARD)

Create a zarr replay buffer from an iterable of episodes.

Handles store creation, array setup with proper codecs, episode insertion, and metadata. Uses WebP for uint8 images and lz4 for numerical data.

Parameters:

Name	Type	Description	Default
schema	DatasetSchema	Dataset schema with zarr_path and array specs.	required
episodes	Iterable[dict[str, ndarray]]	Iterable yielding episode dicts mapping keys to arrays.	required
total_episodes	int | None	Total number of episodes for progress reporting.	None
image_frames_per_shard	int | None	Number of independently compressed image frames stored in each shard, or None to disable sharding.	DEFAULT_IMAGE_FRAMES_PER_SHARD

Raises:

Type	Description
TypeError	If image_frames_per_shard is not an integer.
ValueError	If image_frames_per_shard is not a positive multiple of the image chunk length.

Source code in src/versatil/data/preprocessing/create_zarr_arrays.py

def create_zarr_replay_buffer(
    schema: DatasetSchema,
    episodes: Iterable[dict[str, np.ndarray]],
    total_episodes: int | None = None,
    image_frames_per_shard: int | None = DEFAULT_IMAGE_FRAMES_PER_SHARD,
) -> None:
    """Create a zarr replay buffer from an iterable of episodes.

    Handles store creation, array setup with proper codecs, episode insertion,
    and metadata. Uses WebP for uint8 images and lz4 for numerical data.

    Args:
        schema: Dataset schema with zarr_path and array specs.
        episodes: Iterable yielding episode dicts mapping keys to arrays.
        total_episodes: Total number of episodes for progress reporting.
        image_frames_per_shard: Number of independently compressed image frames
            stored in each shard, or None to disable sharding.

    Raises:
        TypeError: If image_frames_per_shard is not an integer.
        ValueError: If image_frames_per_shard is not a positive multiple of the
            image chunk length.
    """
    logging.info(
        msg=f"Creating Zarr dataset at {schema.zarr_path} "
        f"from schema {schema.__class__.__name__}"
    )

    store = zarr.storage.LocalStore(schema.zarr_path)
    root = zarr.open_group(store=store, mode="w")
    data_group = root.create_group("data")
    meta_group = root.create_group("meta")

    image_codec = WebPCodec(level=WEBP_QUALITY)
    numeric_compressor = BloscCodec(
        cname="lz4", clevel=5, shuffle=BloscShuffle.noshuffle
    )
    create_zarr_arrays(
        data_group=data_group,
        schema=schema,
        image_codec=image_codec,
        numeric_compressor=numeric_compressor,
        image_frames_per_shard=image_frames_per_shard,
    )

    episode_ends: list[int] = []
    cumulative_len = 0

    with threadpool_limits(limits=1):
        for i, episode_data in enumerate(episodes):
            if total_episodes is not None and i % 50 == 0:
                logging.info(msg=f"  Processing episode {i + 1}/{total_episodes}...")
            for key, array in episode_data.items():
                data_group[key].append(array)
            cumulative_len += len(next(iter(episode_data.values())))
            episode_ends.append(cumulative_len)

    meta_group.create_array(
        "episode_ends",
        data=np.array(episode_ends),
        chunks=(len(episode_ends),),
        compressors=None,
    )

    logging.info(
        msg=f"Created Zarr dataset with {len(episode_ends)} episodes, "
        f"{cumulative_len} total steps."
    )