LeRobot datasets with Daft#

@PublicAPI
def read(
    dataset_uri: str,
    io_config: IOConfig | None = None,
    include_stats: bool = False,
    load_video_frames: str | list[str] | bool = False,
) -> DataFrame:
    """Read a LeRobot v2 or v3 dataset as a lazy DataFrame with one row per frame.

    Reads per-episode metadata and the per-frame sensor data under ``data``,
    joins them on ``episode_index``, and broadcasts each episode's metadata
    across its frames. Optionally decodes the matching video frame for one or
    more camera keys into an image column.

    v2.0/v2.1 datasets store one Parquet/MP4 file per episode. v3 packs many
    episodes into shared shards.

    Args:
        dataset_uri: Huggingface repo id (``org/name``), or a local / remote
            directory (``s3://...``, ``hf://datasets/...``).
        io_config: Optional IO configuration for remote reads.
        include_stats: If True, keep the per-episode ``stats/*`` columns
            (per-feature min/max/mean/std/quantiles). Defaults to False.
        load_video_frames: Which camera keys to decode into image columns,
            aligned to each frame's timestamp. Defaults to False (decode
            nothing). Pass True to decode every video feature, a single key
            (``"observation.image"``), or a list of keys. Decoding requires the
            optional ``av`` (PyAV) and ``Pillow`` dependencies.

    Returns:
        Lazy DataFrame with one row per frame: the frame's sensor columns, the
        broadcast episode metadata, and one image column per decoded video key.
    """
    root = _normalize_dataset_root(dataset_uri)
    info = _read_info(root, io_config=io_config)

    # Keep the per-episode video metadata (notably `videos/{key}/from_timestamp`,
    # the time within the shard where each episode's footage begins). We need it
    # to translate episode-local frame timestamps into absolute shard timestamps
    # when decoding, and drop these internal columns again before returning.
    episode_df = read_episodes(
        dataset_uri, io_config=io_config, include_stats=include_stats, include_video_metadata=True
    )
    df = load_episode_frames(episode_df, dataset_uri, io_config=io_config)

    # Load video frames into memory
    if load_video_frames is not False:
        if load_video_frames is True:
            video_keys = _video_keys(info)
        elif isinstance(load_video_frames, str):
            video_keys = [load_video_frames]
        elif isinstance(load_video_frames, list) and all(isinstance(k, str) for k in load_video_frames):
            video_keys = load_video_frames
        else:
            raise ValueError(f"Invalid value provided for argument load_video_frames=`{load_video_frames}`")

        # Seek by absolute timestamp inside the MP4: `from_timestamp` (where this
        # episode begins in the file) + the per-frame episode-local `timestamp`.
        # v3 shards pack many episodes back to back, so `from_timestamp` is the
        # offset within the shard. v2 stores one MP4 per episode, so
        # `from_timestamp` is 0 and the episode-local timestamp is already
        # absolute in that file. Seeking by `frame_index` is wrong for v3 (it
        # only happens to work for the first episode in each shard).
        fps = float(info["fps"])
        tolerance_s = 1.0 / fps / 2.0  # half a frame period: any closer frame is unambiguously "the" frame

        df = df.into_batches(_DECODE_BATCH_SIZE)
        for k in video_keys:
            df = df.with_column(
                k,
                _decode_lerobot_video_timestamp(
                    col(f"videos/{k}/video"),
                    col(f"videos/{k}/from_timestamp"),
                    col("timestamp"),
                    tolerance_s,
                    0,  # image_width: 0 disables resize (decode at native resolution)
                    0,  # image_height: 0 disables resize
                ),
            )
            df = df.exclude(f"videos/{k}/video")

    # Drop the internal per-episode video metadata we kept above (chunk/file index,
    # from/to timestamp). This restores read_episodes' default of hiding these.
    df = df.exclude(*(c for c in df.column_names if c.startswith("videos/") and not c.endswith("/video")))

    return df

@PublicAPI
def read_episodes(
    dataset_uri: str,
    io_config: IOConfig | None = None,
    include_meta: bool = False,
    include_stats: bool = False,
    include_video_metadata: bool = False,
) -> DataFrame:
    """Read LeRobot episode metadata as a lazy DataFrame (one row per episode).

    v3 reads ``meta/episodes/**/*.parquet``. v2.0/v2.1 reads
    ``meta/episodes.jsonl``. Extra per-episode fields present in that metadata
    are kept as columns.

    Args:
        dataset_uri: Huggingface repo id (`org/name`),
            or a local / remote directory (`s3://...`, `hf://datasets/...`)
        io_config: Optional IO configuration for remote reads.
        include_meta: If True, keep the internal ``meta/episodes/*`` columns
            (v3 chunk/file indices locating each episode's own metadata shard).
            No effect on v2, which has no such bookkeeping columns. Defaults to
            False.
        include_stats: If True, keep per-episode statistics. On v3 these are
            ``stats/*`` columns in the episode parquet; on v2.1 they are joined
            from ``meta/episodes_stats.jsonl``. Defaults to False.
        include_video_metadata: If True, keep the per-episode ``videos/{key}/*``
            columns (chunk/file indices and from/to timestamps locating each
            episode's footage). On v2, ``from_timestamp`` is 0 because each
            episode has its own MP4. Defaults to False.

    Returns:
        Lazy DataFrame of episode metadata, one row per episode. Always includes
        a ``videos/{key}/video`` file-handle column per video feature; the
        ``include_*`` flags control which additional column families are kept.
    """
    root = _normalize_dataset_root(dataset_uri)
    info = _read_info(root, io_config=io_config)
    if _is_v3(info):
        return _read_episodes_v3(
            root,
            info,
            io_config=io_config,
            include_meta=include_meta,
            include_stats=include_stats,
            include_video_metadata=include_video_metadata,
        )
    return _read_episodes_v2(
        root,
        info,
        io_config=io_config,
        include_stats=include_stats,
        include_video_metadata=include_video_metadata,
    )

@PublicAPI
def load_episode_frames(
    episodes: DataFrame,
    dataset_uri: str,
    io_config: IOConfig | None = None,
) -> DataFrame:
    """Expand an episode-level DataFrame into a frame-level DataFrame.

    Reads the per-frame parquet under ``data/**`` (v3 shared shards or v2
    per-episode files) and joins it to the provided episode metadata on
    ``episode_index``, producing one row per frame. Episode metadata is
    broadcast across each episode's frames.

    Filter ``episodes`` before calling this to expand only the episodes you need;
    only the surviving episodes contribute to the join.

    Args:
        episodes: Episode-level DataFrame, typically from :func:`read_episodes`
            (optionally filtered). Must contain an ``episode_index`` column.
        dataset_uri: The same dataset identifier passed to :func:`read_episodes`
            (Huggingface repo id ``org/name``, or a local / remote directory such
            as ``s3://...`` or ``hf://datasets/...``).
        io_config: Optional IO configuration for remote reads.

    Returns:
        Lazy DataFrame with one row per frame.
    """
    root = _normalize_dataset_root(dataset_uri)

    frame_df = daft.read_parquet(f"{root}/data/**", io_config=io_config)
    df = episodes.join(frame_df, on=["episode_index"])
    df = df.exclude("data/chunk_index", "data/file_index")
    return df

@PublicAPI
def read_tasks(dataset_uri: str, io_config: IOConfig | None = None) -> DataFrame:
    """Load task metadata as a DataFrame.

    Prefers ``meta/tasks.parquet`` (current LeRobot default). Falls back to legacy
    ``meta/tasks.jsonl`` when the Parquet file is missing.
    """
    root = _normalize_dataset_root(dataset_uri)

    pq_url = f"{root}/meta/tasks.parquet"
    try:
        return daft.read_parquet(pq_url, io_config=io_config)
    except (OSError, DaftCoreException, FileNotFoundError):
        return daft.read_json(f"{root}/meta/tasks.jsonl", io_config=io_config)