LeRobot v3 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 v3 dataset as a lazy DataFrame with one row per frame.

    Reads the per-episode metadata under ``meta/episodes`` 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.

    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 = [name for name, feat_info in info["features"].items() if feat_info["dtype"] == "video"]
        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}`")

        # An MP4 shard packs many episodes back to back, so the shard's internal
        # frame numbering is NOT the parquet's episode-local `frame_index` (which
        # resets to 0 each episode). Seeking by `frame_index` only happens to work
        # for the first episode in each shard. Instead, seek by absolute timestamp:
        # `from_timestamp` (where this episode begins in the shard) + the per-frame
        # episode-local `timestamp`. That keeps a single coordinate system end to end.
        fps = float(info["fps"])
        tolerance_s = 1.0 / fps / 2.0  # half a frame period: any closer frame is unambiguously "the" frame

        # To increase parallelism, reduce batch size
        df = df.into_batches(16)
        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"),
                    lit(tolerance_s),
                    lit(0),  # image_width: 0 disables resize (decode at native resolution)
                    lit(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 v3 episode metadata as a lazy DataFrame (one row per episode).

    This reads the `meta/episodes/**/*.parquet` path under the dataset root.

    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
            (the chunk/file indices locating each episode's own metadata shard).
            These are bookkeeping for random access into the sharded metadata
            and carry no analytical value once the rows are loaded. Defaults to
            False.
        include_stats: If True, keep the per-episode ``stats/*`` columns
            (per-feature min/max/mean/std/quantiles). Defaults to False.
        include_video_metadata: If True, keep the per-episode ``videos/{key}/*``
            columns (the chunk/file indices and from/to timestamps locating each
            episode's footage within its video shard). 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)
    df = daft.read_parquet(f"{root}/meta/episodes/**/*.parquet", io_config=io_config)
    if not include_meta:
        df = df.exclude(*(c for c in df.column_names if c.startswith("meta/")))
    if not include_stats:
        df = df.exclude(*(c for c in df.column_names if c.startswith("stats/")))

    # Get the video keys
    video_keys = set(name for name, feat_info in info["features"].items() if feat_info["dtype"] == "video")

    for key in video_keys:
        file_name_expr = (
            lit(f"{root}/videos/{key}/chunk-")
            + lpad(col(f"videos/{key}/chunk_index").cast(DataType.string), 3, "0")
            + lit("/file-")
            + lpad(col(f"videos/{key}/file_index").cast(DataType.string), 3, "0")
            + lit(".mp4")
        )

        df = df.with_column(f"videos/{key}/video", video_file(file_name_expr, verify=False, io_config=io_config))

    if not include_video_metadata:
        df = df.exclude(*(c for c in df.column_names if c.startswith("videos/") and not c.endswith("/video")))

    return df

@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/**`` 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)