From 96c206325db0f093fa281c50b5ce3db4fae4ad77 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Tue, 16 Jun 2026 17:31:14 +0200 Subject: [PATCH 1/4] Extract Rerun reader --- docs/nav.md | 1 + docs/reading-data/index.md | 1 + docs/reading-data/rerun.md | 119 +++ docs/reference/optional-dependencies.md | 2 + pyproject.toml | 5 + src/refiner/__init__.py | 3 + src/refiner/execution/engine.py | 91 +- src/refiner/pipeline/__init__.py | 3 + src/refiner/pipeline/_rerun_io.py | 82 ++ src/refiner/pipeline/pipeline.py | 81 ++ src/refiner/pipeline/sources/__init__.py | 2 + src/refiner/pipeline/sources/base.py | 7 +- .../pipeline/sources/readers/__init__.py | 2 + src/refiner/pipeline/sources/readers/rerun.py | 969 ++++++++++++++++++ tests/readers/test_rerun_reader.py | 588 +++++++++++ uv.lock | 55 +- 16 files changed, 1995 insertions(+), 16 deletions(-) create mode 100644 docs/reading-data/rerun.md create mode 100644 src/refiner/pipeline/_rerun_io.py create mode 100644 src/refiner/pipeline/sources/readers/rerun.py create mode 100644 tests/readers/test_rerun_reader.py diff --git a/docs/nav.md b/docs/nav.md index 4b2ba88f..065282a9 100644 --- a/docs/nav.md +++ b/docs/nav.md @@ -20,6 +20,7 @@ - [HDF5](reading-data/hdf5.md) - [Zarr](reading-data/zarr.md) - [MCAP](reading-data/mcap.md) +- [Rerun](reading-data/rerun.md) - [Tabular files](reading-data/tabular-files.md) - [Files and videos](reading-data/files-and-videos.md) - [Hugging Face](reading-data/hugging-face.md) diff --git a/docs/reading-data/index.md b/docs/reading-data/index.md index 62639d59..3a1f92ca 100644 --- a/docs/reading-data/index.md +++ b/docs/reading-data/index.md @@ -22,6 +22,7 @@ pipeline = mdr.read_lerobot("hf://datasets/lerobot/aloha_sim_transfer_cube_human | One HDF5 file per episode, or grouped HDF5 demos | `read_hdf5` | [HDF5](hdf5.md) | | Zarr replay buffer with episode boundaries | `read_zarr` | [Zarr](zarr.md) | | MCAP robotics or autonomy logs | `read_mcap` | [MCAP](mcap.md) | +| Rerun RRD recordings | `read_rerun` | [Rerun](rerun.md) | | Parquet, JSON, JSONL, CSV tables | `read_parquet`, `read_json`, `read_jsonl`, `read_csv` | [Tabular Files](tabular-files.md) | | Raw files or media files | `read_files`, `read_videos` | [Files and Videos](files-and-videos.md) | | Hugging Face datasets table | `read_hf_dataset` | [Hugging Face](hugging-face.md) | diff --git a/docs/reading-data/rerun.md b/docs/reading-data/rerun.md new file mode 100644 index 00000000..6ce1a86a --- /dev/null +++ b/docs/reading-data/rerun.md @@ -0,0 +1,119 @@ +--- +title: "Rerun reader" +description: "Read Rerun RRD recordings as columnar or robotics episode rows" +--- + +# Rerun reader + +Use `read_rerun` for `.rrd` files written by Rerun. + +```python +import refiner as mdr + +pipeline = mdr.read_rerun( + "s3://bucket/run/**/*.rrd", + output="robotics", + fps=30, +) +``` + +Install `macrodata-refiner[rerun]` to use this reader. Add storage extras such +as `s3` when reading remote paths. + +Directory inputs are filtered to paths ending in `.rrd`. RRD files are planned +as atomic files, so workers parallelize across recordings instead of splitting +one recording by byte range. + +## Recording rows + +With `output="recording"`, `read_rerun` emits one row per Rerun recording +segment: + +```python +rows = mdr.read_rerun( + "/data/run/*.rrd", + output="recording", + contents=("/action/**", "/observation/**"), + timelines=("frame",), +) +``` + +Each row includes: + +| Column | Meaning | +| --- | --- | +| `episode_id` | Rerun recording id or segment id. | +| `rerun` | `RerunRecording` value with Arrow-backed `Tabular` tables by timeline. | +| `file_path` | Source RRD path, unless `file_path_column=None`. | + +`contents` is passed to Rerun's content filter. `timelines` limits the timeline +tables returned. If `timelines` is omitted, the reader materializes all timeline +indexes reported by the Rerun schema. + +Set `materialize_tables=False` with `output="recording"` when you only need the +recording metadata and do not want to materialize Arrow timeline/static tables. +That mode keeps the row lightweight while preserving the source recording +identity and chunk metadata. + +## Robotics rows + +With `output="robotics"`, the reader creates rows that can be passed to +`to_robot_rows(...)` and robotics writers: + +```python +robot_rows = ( + mdr.read_rerun( + "/data/episodes/*.rrd", + output="robotics", + fps=30, + robot_type="unknown", + ) + .to_robot_rows( + episode_id_key="episode_id", + nested_frames_key="frames", + fps_key="fps", + robot_type_key="robot_type", + video_keys={ + "observation.images.top": "cam.top", + "observation.images.left_wrist": "cam.left_wrist", + }, + ) +) +``` + +The default robotics mapping reads scalar components under `/action/**` into +the frame `action` vector, scalar components under `/observation/state/**` into +`observation.state`, and encoded images under `/cam/**` into top-level video +sources such as `cam.top`. + +Use explicit selections when vector order or camera names matter: + +```python +mdr.read_rerun( + "episode.rrd", + output="robotics", + actions=("/robot/actions/gripper", "/robot/actions/arm"), + states=("/robot/state/qpos", "/robot/state/gripper"), + videos={"observation.images.top": "/robot/cameras/top"}, + fps=30, +) +``` + +`actions` and `states` define vector order. `videos` maps output video keys to +Rerun encoded-image entity paths. If `contents` is omitted, explicit selections +also define the minimal Rerun content filter for those categories. + +## Decoding + +Scalar action and state columns are read from Arrow list arrays and converted +to frame vectors. Encoded images remain lazy `VideoFrameSequence` values; JPEG +or PNG bytes are decoded frame-by-frame only when a downstream video writer or +consumer iterates the sequence. + +## Sharding + +Rerun SDK queries require a complete local RRD file. For that reason, +`read_rerun` sets file-atomic sharding, like other container readers such as +HDF5 and MCAP. `target_shard_bytes` groups whole RRD files into shard buckets, +and `num_shards` can request a target number of file buckets when there are +enough files. diff --git a/docs/reference/optional-dependencies.md b/docs/reference/optional-dependencies.md index 399d632b..51356f46 100644 --- a/docs/reference/optional-dependencies.md +++ b/docs/reference/optional-dependencies.md @@ -16,6 +16,7 @@ Install extras based on the data and operations you use. | `hdf5` | HDF5 reader support. | | `zarr` | Zarr reader and writer support. | | `mcap` | MCAP robotics log reader support, including ROS2, protobuf, and H.264 video decoding. | +| `rerun` | Rerun RRD reader support. | | `video` | Video decode/write support. | | `text` | Common Crawl text readers. | | `s3` | S3 filesystem support. | @@ -30,5 +31,6 @@ pip install macrodata-refiner[hf,video] pip install macrodata-refiner[datasets] pip install macrodata-refiner[hdf5,zarr] pip install macrodata-refiner[mcap] +pip install macrodata-refiner[rerun] pip install macrodata-refiner[hand_tracking] ``` diff --git a/pyproject.toml b/pyproject.toml index 4de39281..12211ef1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,10 @@ mcap = [ "mcap-ros2-support", "pillow", ] +rerun = [ + "pillow", + "rerun-sdk[datafusion]>=0.33,<0.34", +] s3 = [ "s3fs", ] @@ -87,6 +91,7 @@ all = [ "macrodata-refiner[hdf5]", "macrodata-refiner[hf]", "macrodata-refiner[mcap]", + "macrodata-refiner[rerun]", "macrodata-refiner[video]", "macrodata-refiner[zarr]", "macrodata-refiner[text]", diff --git a/src/refiner/__init__.py b/src/refiner/__init__.py index 15dba63f..b6ea9ede 100644 --- a/src/refiner/__init__.py +++ b/src/refiner/__init__.py @@ -31,6 +31,7 @@ "read_lerobot": "refiner.pipeline", "read_mcap": "refiner.pipeline", "read_parquet": "refiner.pipeline", + "read_rerun": "refiner.pipeline", "read_tfds": "refiner.pipeline", "read_tfrecords": "refiner.pipeline", "read_videos": "refiner.pipeline", @@ -66,6 +67,7 @@ "read_lerobot", "read_mcap", "read_parquet", + "read_rerun", "read_tfds", "read_tfrecords", "read_videos", @@ -140,6 +142,7 @@ def __dir__() -> list[str]: read_lerobot, read_mcap, read_parquet, + read_rerun, read_tfds, read_tfrecords, read_videos, diff --git a/src/refiner/execution/engine.py b/src/refiner/execution/engine.py index 0006fa15..b9bf4f80 100644 --- a/src/refiner/execution/engine.py +++ b/src/refiner/execution/engine.py @@ -9,6 +9,7 @@ from refiner.pipeline.data.block import Block, StreamItem from refiner.pipeline.data.datatype import schema_with_dtypes +from refiner.pipeline.data.shard import SHARD_ID_COLUMN from refiner.pipeline.data.tabular import Tabular from refiner.pipeline.steps import ( CastStep, @@ -31,7 +32,7 @@ from refiner.execution.operators.vectorized import ( apply_vectorized_ops, ) -from refiner.pipeline.data.row import Row +from refiner.pipeline.data.row import DictRow, Row _DEFAULT_VECTORIZED_CHUNK_ROWS = 2048 @@ -301,27 +302,50 @@ def _execute_vector_segment( pending_rows = RowBuffer() current_chunk_rows = max(1, int(vectorized_chunk_rows)) estimated_row_bytes: float | None = None - segment_changes_rows = any( - isinstance(op, (FilterExprStep, FnTableStep)) for op in ops + row_projection_ops, row_remaining_ops = _split_row_projection_ops(ops) + row_projected_schema = ( + _vector_segment_schema(input_schema, row_projection_ops) + if row_projection_ops + else input_schema ) - def _run_block(block: Tabular) -> Tabular: - return_row_indices = block.needs_row_indices and segment_changes_rows + def _run_block(block: Tabular, block_ops: Sequence[VectorizedOp]) -> Tabular: + block_changes_rows = any( + isinstance(op, (FilterExprStep, FnTableStep)) for op in block_ops + ) + return_row_indices = block.needs_row_indices and block_changes_rows if not return_row_indices: table = apply_vectorized_ops( block.table, - ops, + block_ops, on_shard_delta=on_shard_delta, ) return block.with_table(table) table, row_indices = apply_vectorized_ops( block.table, - ops, + block_ops, on_shard_delta=on_shard_delta, return_row_indices=True, ) return block.with_table(table, row_indices=row_indices) + def _rows_to_block(batch: list[Row]) -> tuple[Tabular, Sequence[VectorizedOp]]: + try: + return _tabular_from_rows(batch, schema=input_schema), ops + except (pa.ArrowInvalid, pa.ArrowTypeError, TypeError) as err: + if not row_projection_ops: + raise + projected = [ + _apply_row_projection_ops(row, row_projection_ops) for row in batch + ] + try: + return ( + _tabular_from_rows(projected, schema=row_projected_schema), + row_remaining_ops, + ) + except (pa.ArrowInvalid, pa.ArrowTypeError, TypeError) as fallback_err: + raise err from fallback_err + def _chunk_rows_for_budget() -> int: if ( max_vectorized_block_bytes is None @@ -338,11 +362,7 @@ def _run_pending_chunk(target_rows: int) -> Iterator[Tabular]: while True: batch = pending_rows.peek(rows_for_try) try: - block = ( - Tabular.from_rows(batch, schema=input_schema) - if not batch - else batch[0].tabular_type.from_rows(batch, schema=input_schema) - ) + block, block_ops = _rows_to_block(batch) except pa.ArrowMemoryError: if rows_for_try <= 1: raise @@ -368,7 +388,7 @@ def _run_pending_chunk(target_rows: int) -> Iterator[Tabular]: continue try: - out = _run_block(block) + out = _run_block(block, block_ops) except pa.ArrowMemoryError: if rows_for_try <= 1: raise @@ -420,7 +440,7 @@ def _yield_tabular_chunks(block: Tabular) -> Iterator[Tabular]: continue try: - out = _run_block(chunk) + out = _run_block(chunk, ops) except pa.ArrowMemoryError: if chunk_rows <= 1: raise @@ -457,6 +477,49 @@ def _yield_tabular_chunks(block: Tabular) -> Iterator[Tabular]: yield from _drain_rows(force=True) +def _tabular_from_rows( + rows: list[Row], + *, + schema: pa.Schema | None, +) -> Tabular: + return ( + Tabular.from_rows(rows, schema=schema) + if not rows + else rows[0].tabular_type.from_rows(rows, schema=schema) + ) + + +def _split_row_projection_ops( + ops: Sequence[VectorizedOp], +) -> tuple[tuple[SelectStep | DropStep, ...], Sequence[VectorizedOp]]: + projection: list[SelectStep | DropStep] = [] + for index, op in enumerate(ops): + if not isinstance(op, (SelectStep, DropStep)): + return tuple(projection), ops[index:] + projection.append(op) + return tuple(projection), () + + +def _apply_row_projection_ops( + row: Row, + ops: Sequence[SelectStep | DropStep], +) -> Row: + out = row + for op in ops: + if isinstance(op, SelectStep): + out = DictRow( + { + column: out[column] + for column in op.columns + if column != SHARD_ID_COLUMN + }, + shard_id=out.shard_id, + ) + continue + out = out.drop(*(column for column in op.columns if column != SHARD_ID_COLUMN)) + return out + + def _chunk_output_rows(rows: Iterable[Row], block_rows: int) -> Iterator[list[Row]]: pending: list[Row] = [] for row in rows: diff --git a/src/refiner/pipeline/__init__.py b/src/refiner/pipeline/__init__.py index 856ec288..cf7389bd 100644 --- a/src/refiner/pipeline/__init__.py +++ b/src/refiner/pipeline/__init__.py @@ -18,6 +18,7 @@ "read_lerobot": "refiner.pipeline.pipeline", "read_mcap": "refiner.pipeline.pipeline", "read_parquet": "refiner.pipeline.pipeline", + "read_rerun": "refiner.pipeline.pipeline", "read_tfds": "refiner.pipeline.pipeline", "read_tfrecords": "refiner.pipeline.pipeline", "read_videos": "refiner.pipeline.pipeline", @@ -48,6 +49,7 @@ "read_lerobot", "read_mcap", "read_parquet", + "read_rerun", "read_tfds", "read_tfrecords", "read_videos", @@ -88,6 +90,7 @@ def __dir__() -> list[str]: read_lerobot, read_mcap, read_parquet, + read_rerun, read_tfds, read_tfrecords, read_videos, diff --git a/src/refiner/pipeline/_rerun_io.py b/src/refiner/pipeline/_rerun_io.py new file mode 100644 index 00000000..5388c364 --- /dev/null +++ b/src/refiner/pipeline/_rerun_io.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +import os +import tempfile +from pathlib import Path +from typing import cast + +from refiner.io import DataFile +from refiner.pipeline.data.tabular import Tabular + + +class LocalRrd: + def __init__(self, source: DataFile) -> None: + self.source = source + self.tmpdir: tempfile.TemporaryDirectory[str] | None = None + self.path: Path | None = None + + def open(self) -> Path: + if self.path is not None: + return self.path + if self.source.is_local: + self.path = Path(self.source.abs_path()) + return self.path + self.tmpdir = tempfile.TemporaryDirectory(prefix="refiner-rerun-") + name = os.path.basename(self.source.path) or "recording.rrd" + self.path = Path(self.tmpdir.name) / name + self.source.copy(str(self.path)) + return self.path + + def close(self) -> None: + if self.tmpdir is not None: + self.tmpdir.cleanup() + self.tmpdir = None + self.path = None + + def __enter__(self) -> Path: + return self.open() + + def __exit__(self, *args: object) -> None: + self.close() + + def __del__(self) -> None: + try: + self.close() + except Exception: + pass + + def __getstate__(self) -> dict[str, object]: + return { + "source": self.source, + "path": str(self.path) if self.path is not None else None, + } + + def __setstate__(self, state: dict[str, object]) -> None: + self.source = cast(DataFile, state["source"]) + self.tmpdir = None + path = state.get("path") + self.path = Path(path) if isinstance(path, str) else None + + +@dataclass(frozen=True, slots=True) +class RerunRecording: + """Columnar Rerun recording data loaded from one RRD segment.""" + + segment_id: str + source_path: str + tables: Mapping[str, Tabular] + static: Tabular | None = None + source_file: DataFile | None = None + local_source: LocalRrd | None = None + application_id: str | None = None + recording_id: str | None = None + contents: tuple[str, ...] | None = None + timelines: tuple[str, ...] | None = None + include_static: bool = True + use_source_chunks: bool = True + source_recording_count: int | None = None + + +__all__ = ["LocalRrd", "RerunRecording"] diff --git a/src/refiner/pipeline/pipeline.py b/src/refiner/pipeline/pipeline.py index 4ee33797..0822f1b0 100644 --- a/src/refiner/pipeline/pipeline.py +++ b/src/refiner/pipeline/pipeline.py @@ -41,6 +41,7 @@ JsonReader, McapReader, ParquetReader, + RerunReader, TfdsReader, TfrecordReader, ZarrReader, @@ -48,6 +49,12 @@ from refiner.pipeline.sources.readers.hdf5 import MissingPolicy from refiner.pipeline.sources.readers.lerobot import LeRobotEpisodeReader from refiner.pipeline.sources.readers.mcap import SyncMethod +from refiner.pipeline.sources.readers.rerun import ( + DEFAULT_RERUN_ACTION_PREFIX, + DEFAULT_RERUN_CAMERA_PREFIX, + DEFAULT_RERUN_STATE_PREFIX, + RerunOutputMode, +) from refiner.pipeline.sources.items import ItemsSource from refiner.pipeline.sources.task import TaskSource, TaskStep from refiner.pipeline.data import datatype @@ -710,6 +717,7 @@ def launch_cloud( gpu: GPU | None = None, sync_local_dependencies: bool = False, dependencies: Sequence[str] | None = None, + extra_dependencies: Sequence[str] | None = None, refiner_extras: Sequence[str] | None = None, secrets: SecretInput | None = None, env: Mapping[str, object | None] | None = None, @@ -729,6 +737,7 @@ def launch_cloud( dependencies: Additional packages to install in the cloud runtime. Entries are requirement strings such as `"torch"` or `"ego-vision[models]==0.1.2"`. + extra_dependencies: Compatibility alias for ``dependencies``. refiner_extras: Additional macrodata-refiner extras to install in the cloud runtime. Built-in blocks automatically declare the extras they require; pass this for extras used outside those @@ -746,6 +755,11 @@ def launch_cloud( """ from refiner.launchers.cloud import CloudLauncher + if dependencies is not None and extra_dependencies is not None: + raise ValueError("Pass only one of dependencies or extra_dependencies") + if dependencies is None: + dependencies = extra_dependencies + launcher = CloudLauncher( pipeline=self, name=name, @@ -1217,6 +1231,73 @@ def read_mcap( ) +def read_rerun( + inputs: DataFileSetLike, + *, + fs: AbstractFileSystem | None = None, + storage_options: Mapping[str, Any] | None = None, + recursive: bool = False, + target_shard_bytes: int = DEFAULT_TARGET_SHARD_BYTES, + num_shards: int | None = None, + file_path_column: str | None = "file_path", + output: RerunOutputMode = "recording", + contents: str | Sequence[str] | None = None, + timelines: Sequence[str] | None = None, + primary_timeline: str | None = None, + include_static: bool = True, + materialize_tables: bool = True, + include_recording: bool | None = None, + fill_latest_at: bool = False, + action_prefix: str = DEFAULT_RERUN_ACTION_PREFIX, + state_prefix: str = DEFAULT_RERUN_STATE_PREFIX, + camera_prefix: str = DEFAULT_RERUN_CAMERA_PREFIX, + actions: PathSelection | None = None, + states: PathSelection | None = None, + videos: PathSelection | None = None, + fps: float | None = None, + robot_type: str | None = None, +) -> RefinerPipeline: + """Create a pipeline with a Rerun RRD reader source. + + RRD files are planned as atomic input shards. With ``output="recording"``, + each emitted row preserves the selected Rerun data as Arrow-backed + ``Tabular`` tables grouped by timeline under the ``rerun`` field. Set + ``materialize_tables=False`` when you only need source chunk metadata. + With ``output="robotics"``, the reader additionally derives common + robotics episode fields from configurable Rerun entity prefixes so the + rows can be passed through ``to_robot_rows(...)``. Pass ``actions``, + ``states``, or ``videos`` to pin exact entity paths and output order + instead of using prefix-derived defaults. + """ + return RefinerPipeline( + source=RerunReader( + inputs, + fs=fs, + storage_options=storage_options, + recursive=recursive, + target_shard_bytes=target_shard_bytes, + num_shards=num_shards, + file_path_column=file_path_column, + output=output, + contents=contents, + timelines=timelines, + primary_timeline=primary_timeline, + include_static=include_static, + materialize_tables=materialize_tables, + include_recording=include_recording, + fill_latest_at=fill_latest_at, + action_prefix=action_prefix, + state_prefix=state_prefix, + camera_prefix=camera_prefix, + actions=actions, + states=states, + videos=videos, + fps=fps, + robot_type=robot_type, + ) + ) + + def read_parquet( inputs: DataFileSetLike, *, diff --git a/src/refiner/pipeline/sources/__init__.py b/src/refiner/pipeline/sources/__init__.py index 3ef0c3e7..90857a84 100644 --- a/src/refiner/pipeline/sources/__init__.py +++ b/src/refiner/pipeline/sources/__init__.py @@ -9,6 +9,7 @@ LeRobotEpisodeReader, McapReader, ParquetReader, + RerunReader, TfdsReader, TfrecordReader, ZarrReader, @@ -25,6 +26,7 @@ "LeRobotEpisodeReader", "McapReader", "ParquetReader", + "RerunReader", "TfdsReader", "TfrecordReader", "ZarrReader", diff --git a/src/refiner/pipeline/sources/base.py b/src/refiner/pipeline/sources/base.py index 75e85b52..cff34ee3 100644 --- a/src/refiner/pipeline/sources/base.py +++ b/src/refiner/pipeline/sources/base.py @@ -12,7 +12,7 @@ from refiner.worker.metrics.api import log_throughput _INTERNAL_SHARD_ID_KEY = "__shard_id" -SourceUnit: TypeAlias = Row | Tabular +SourceUnit: TypeAlias = Row | list[Row] | Tabular class BaseSource(ABC): @@ -73,6 +73,8 @@ def _io_refiner_extras(self) -> tuple[str, ...]: def _unit_num_rows(unit: SourceUnit) -> int: if isinstance(unit, Row): return 1 + if isinstance(unit, list): + return len(unit) if isinstance(unit, Tabular): return int(unit.num_rows) raise TypeError(f"Unsupported source unit type: {type(unit)!r}") @@ -82,6 +84,9 @@ def _with_shard_id(unit: SourceUnit, shard_id: str) -> SourceUnit: if isinstance(unit, Row): return unit.update(**{_INTERNAL_SHARD_ID_KEY: shard_id}) + if isinstance(unit, list): + return [row.update(**{_INTERNAL_SHARD_ID_KEY: shard_id}) for row in unit] + if isinstance(unit, Tabular): table = unit.table if table.num_rows == 0: diff --git a/src/refiner/pipeline/sources/readers/__init__.py b/src/refiner/pipeline/sources/readers/__init__.py index 09a1184b..8a3eda84 100644 --- a/src/refiner/pipeline/sources/readers/__init__.py +++ b/src/refiner/pipeline/sources/readers/__init__.py @@ -7,6 +7,7 @@ from refiner.pipeline.sources.readers.lerobot import LeRobotEpisodeReader from refiner.pipeline.sources.readers.mcap import McapReader from refiner.pipeline.sources.readers.parquet import ParquetReader +from refiner.pipeline.sources.readers.rerun import RerunReader from refiner.pipeline.sources.readers.tfds import TfdsReader from refiner.pipeline.sources.readers.tfrecord import TfrecordReader from refiner.pipeline.sources.readers.zarr import ZarrReader @@ -23,6 +24,7 @@ "LeRobotRow", "McapReader", "ParquetReader", + "RerunReader", "TfdsReader", "TfrecordReader", "ZarrReader", diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py new file mode 100644 index 00000000..673de053 --- /dev/null +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -0,0 +1,969 @@ +from __future__ import annotations + +import concurrent.futures +from collections.abc import Iterable, Iterator, Mapping, Sequence +from pathlib import Path +from typing import Any, Literal, cast +import warnings + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc +from fsspec import AbstractFileSystem + +from refiner.io import DataFile +from refiner.io.fileset import DataFileSetLike +from refiner.pipeline._rerun_io import LocalRrd, RerunRecording +from refiner.pipeline.data.row import DictRow, Row +from refiner.pipeline.data.shard import FilePartsDescriptor, Shard +from refiner.pipeline.data.tabular import Tabular +from refiner.pipeline.sources.base import SourceUnit +from refiner.pipeline.sources.readers.base import BaseReader +from refiner.pipeline.sources.readers.utils import ( + DEFAULT_TARGET_SHARD_BYTES, + PathSelection, + path_selection_map, +) +from refiner.utils import check_required_dependencies +from refiner.video import VideoFrameSequence +from refiner.worker.context import logger + +RerunOutputMode = Literal["recording", "robotics"] + +_RERUN_SEGMENT_ID = "rerun_segment_id" +_RERUN_COMPONENT_METADATA = b"rerun:component" +_RERUN_ENTITY_PATH_METADATA = b"rerun:entity_path" +_ROBOTICS_ROW_COLUMNS = frozenset( + {"episode_id", "rerun", "frames", "fps", "robot_type"} +) +_RECORDING_ROW_COLUMNS = frozenset({"episode_id", "rerun"}) +DEFAULT_RERUN_ACTION_PREFIX = "/action" +DEFAULT_RERUN_STATE_PREFIX = "/observation/state" +DEFAULT_RERUN_CAMERA_PREFIX = "/cam" +# Amortize Rerun server startup for small files without staging an unbounded shard. +_MAX_STAGED_RRD_BATCH_BYTES = 512 * 1024 * 1024 +_MAX_STAGED_RRD_BATCH_FILES = 16 + + +def _reject_recording_robotics_options( + *, + primary_timeline: str | None, + action_prefix: str, + state_prefix: str, + camera_prefix: str, + actions: PathSelection | None, + states: PathSelection | None, + videos: PathSelection | None, + fps: float | None, + robot_type: str | None, +) -> None: + invalid = [] + if primary_timeline is not None: + invalid.append("primary_timeline") + if action_prefix != DEFAULT_RERUN_ACTION_PREFIX: + invalid.append("action_prefix") + if state_prefix != DEFAULT_RERUN_STATE_PREFIX: + invalid.append("state_prefix") + if camera_prefix != DEFAULT_RERUN_CAMERA_PREFIX: + invalid.append("camera_prefix") + if actions is not None: + invalid.append("actions") + if states is not None: + invalid.append("states") + if videos is not None: + invalid.append("videos") + if fps is not None: + invalid.append("fps") + if robot_type is not None: + invalid.append("robot_type") + if invalid: + raise ValueError( + "Rerun recording output does not use robotics options: " + + ", ".join(invalid) + ) + + +class RerunReader(BaseReader): + """Read Rerun RRD files as columnar recording rows or robotics episode rows.""" + + name = "read_rerun" + + def __init__( + self, + inputs: DataFileSetLike, + *, + fs: AbstractFileSystem | None = None, + storage_options: Mapping[str, Any] | None = None, + recursive: bool = False, + target_shard_bytes: int = DEFAULT_TARGET_SHARD_BYTES, + num_shards: int | None = None, + file_path_column: str | None = "file_path", + output: RerunOutputMode = "recording", + contents: str | Sequence[str] | None = None, + timelines: Sequence[str] | None = None, + primary_timeline: str | None = None, + include_static: bool = True, + materialize_tables: bool = True, + include_recording: bool | None = None, + fill_latest_at: bool = False, + action_prefix: str = DEFAULT_RERUN_ACTION_PREFIX, + state_prefix: str = DEFAULT_RERUN_STATE_PREFIX, + camera_prefix: str = DEFAULT_RERUN_CAMERA_PREFIX, + actions: PathSelection | None = None, + states: PathSelection | None = None, + videos: PathSelection | None = None, + fps: float | None = None, + robot_type: str | None = None, + ) -> None: + if output not in ("recording", "robotics"): + raise ValueError("output must be 'recording' or 'robotics'") + if output == "robotics" and not materialize_tables: + raise ValueError( + "materialize_tables=False is only supported for recording output" + ) + if output == "recording" and include_recording is False: + raise ValueError( + "include_recording=False is only supported for robotics output" + ) + if output == "recording": + _reject_recording_robotics_options( + primary_timeline=primary_timeline, + action_prefix=action_prefix, + state_prefix=state_prefix, + camera_prefix=camera_prefix, + actions=actions, + states=states, + videos=videos, + fps=fps, + robot_type=robot_type, + ) + if fps is not None: + fps = float(fps) + if not np.isfinite(fps) or fps <= 0: + raise ValueError("fps must be > 0") + super().__init__( + inputs, + fs=fs, + storage_options=storage_options, + recursive=recursive, + extensions=(".rrd",), + target_shard_bytes=target_shard_bytes, + num_shards=num_shards, + file_path_column=file_path_column, + split_by_bytes=False, + ) + self.output = output + self.contents = _contents(contents) + self.timelines = tuple(timelines) if timelines is not None else None + self.primary_timeline = primary_timeline + self.include_static = include_static + self.materialize_tables = materialize_tables + self.use_source_chunks = self.timelines is None + self.include_recording = ( + output == "recording" if include_recording is None else include_recording + ) + self.fill_latest_at = fill_latest_at + self.action_prefix = _normalize_entity_prefix(action_prefix) + self.state_prefix = _normalize_entity_prefix(state_prefix) + self.camera_prefix = _normalize_entity_prefix(camera_prefix) + self.actions_explicit = actions is not None + self.states_explicit = states is not None + self.videos_explicit = videos is not None + self.actions = _selection_map( + actions, + format_name="Rerun actions", + derive_names_from_paths=False, + ) + self.states = _selection_map( + states, + format_name="Rerun states", + derive_names_from_paths=False, + ) + self.videos = _selection_map(videos, format_name="Rerun videos") + reserved_row_columns = ( + _ROBOTICS_ROW_COLUMNS if output == "robotics" else _RECORDING_ROW_COLUMNS + ) + if file_path_column in reserved_row_columns: + raise ValueError( + f"file_path_column cannot use reserved Rerun {output} row " + f"column {file_path_column!r}" + ) + if output == "robotics": + reserved_video_names = set(reserved_row_columns) + if file_path_column is not None: + reserved_video_names.add(file_path_column) + video_collisions = set(self.videos).intersection(reserved_video_names) + if video_collisions: + raise ValueError( + "Rerun video output names cannot use reserved robotics row " + "columns: " + ", ".join(sorted(video_collisions)) + ) + self.fps = fps + self.robot_type = robot_type + + def _declared_refiner_extras(self) -> tuple[str, ...]: + return ("rerun",) + + def describe(self) -> dict[str, Any]: + description = super().describe() + description.update( + { + "output": self.output, + "contents": self.contents, + "timelines": self.timelines, + "include_static": self.include_static, + "materialize_tables": self.materialize_tables, + "include_recording": self.include_recording, + "fill_latest_at": self.fill_latest_at, + } + ) + if self.output == "robotics": + description.update( + { + "primary_timeline": self.primary_timeline, + "action_prefix": self.action_prefix, + "state_prefix": self.state_prefix, + "camera_prefix": self.camera_prefix, + "actions": dict(self.actions) if self.actions_explicit else None, + "states": dict(self.states) if self.states_explicit else None, + "videos": dict(self.videos) if self.videos_explicit else None, + "fps": self.fps, + "robot_type": self.robot_type, + } + ) + return description + + def read_shard(self, shard: Shard) -> Iterator[SourceUnit]: + descriptor = shard.descriptor + assert isinstance(descriptor, FilePartsDescriptor) + batch: list[tuple[DataFile, LocalRrd]] = [] + batch_bytes = 0 + for part in descriptor.parts: + source = self.fileset.resolve_file(part.source_index, part.path) + part_size = max(0, self.fileset.size(part.source_index, part.path)) + if batch and ( + len(batch) >= _MAX_STAGED_RRD_BATCH_FILES + or batch_bytes + part_size > _MAX_STAGED_RRD_BATCH_BYTES + ): + yield from self._read_staged_batch(batch) + batch = [] + batch_bytes = 0 + local_source = LocalRrd(source) + batch.append((source, local_source)) + batch_bytes += part_size + if batch: + yield from self._read_staged_batch(batch) + + def _read_staged_batch( + self, + local_files: Sequence[tuple[DataFile, LocalRrd]], + ) -> Iterator[SourceUnit]: + opened_files = _open_local_sources(local_files) + if self._retain_batch_local_sources(): + try: + units = list(self._read_files(opened_files)) + except BaseException: + _close_local_sources(opened_files) + raise + if not units: + _close_local_sources(opened_files) + return + try: + yield cast(list[Row], units) + finally: + _close_local_sources(opened_files) + return + + try: + yield from self._read_files(opened_files) + finally: + _close_local_sources(opened_files) + + def _read_files( + self, + local_files: Sequence[tuple[DataFile, Path, LocalRrd]], + ) -> Iterator[SourceUnit]: + if self.output == "recording" and not self.materialize_tables: + check_required_dependencies( + "read_rerun", + [("rerun", "rerun-sdk")], + dist="rerun", + ) + server_fallback: list[tuple[DataFile, Path, LocalRrd]] = [] + for ( + source, + local_path, + local_source, + store_entries, + ) in _scan_recording_entries(local_files): + rows = self._read_metadata_only_recording_rows( + source, + local_source, + store_entries, + ) + if rows: + yield from rows + else: + server_fallback.append((source, local_path, local_source)) + if server_fallback: + yield from self._read_files_with_server(server_fallback) + return + + yield from self._read_files_with_server(local_files) + + def _read_files_with_server( + self, + local_files: Sequence[tuple[DataFile, Path, LocalRrd]], + ) -> Iterator[SourceUnit]: + check_required_dependencies( + "read_rerun", + [("rerun", "rerun-sdk"), "datafusion"], + dist="rerun", + ) + import rerun as rr + + datasets = { + f"recording_{index}": (str(local_path),) + for index, (_source, local_path, _local_rrd) in enumerate(local_files) + } + with rr.server.Server(datasets=cast(Any, datasets)) as server: + client = server.client() + for dataset_name, (source, local_path, local_source) in zip( + datasets, local_files, strict=True + ): + dataset = client.get_dataset(dataset_name) + yield from self._read_dataset( + source, + local_path, + local_source, + dataset, + ) + + def _read_dataset( + self, + source: DataFile, + local_path: Path, + local_source: LocalRrd, + dataset: Any, + store_entries: Sequence[Any] | None = None, + ) -> Iterator[SourceUnit]: + if store_entries is None: + store_entries = ( + _recording_entries(local_path) + if self.output == "recording" or self.include_recording + else [] + ) + source_recording_count = len(store_entries) if store_entries else None + entries_by_recording_id = {entry.recording_id: entry for entry in store_entries} + timelines = self.timelines + if timelines is None: + timelines = self._timelines(dataset.schema()) + for segment_id in dataset.segment_ids(): + store = entries_by_recording_id.get(segment_id) + application_id = store.application_id if store is not None else None + recording_id = store.recording_id if store is not None else segment_id + view = dataset.filter_segments([segment_id]) + if self.output == "robotics": + yield self._robotics_row( + view, + segment_id=segment_id, + source_path=source.abs_path(), + source_file=source, + local_source=local_source, + application_id=application_id, + recording_id=recording_id, + timelines=timelines, + ) + else: + tables: dict[str, Tabular] = {} + static = None + if self.materialize_tables: + content_view = self._view_for_contents(view) + static = ( + _collect_table(content_view.reader(index=None)) + if self.include_static + else None + ) + tables = { + timeline: Tabular( + _collect_table( + content_view.reader( + index=timeline, + fill_latest_at=self.fill_latest_at, + ) + ) + ) + for timeline in timelines + } + yield self._recording_row( + segment_id=segment_id, + source=source, + local_source=local_source, + tables=tables, + static=Tabular(static) if static is not None else None, + application_id=application_id, + recording_id=recording_id, + source_recording_count=source_recording_count, + ) + + def _read_metadata_only_recording_rows( + self, + source: DataFile, + local_source: LocalRrd, + store_entries: Sequence[Any], + ) -> list[DictRow]: + rows = [] + source_recording_count = len(store_entries) + for store in store_entries: + recording_id = str(store.recording_id) + rows.append( + self._recording_row( + segment_id=recording_id, + source=source, + local_source=local_source, + tables={}, + static=None, + application_id=store.application_id, + recording_id=recording_id, + source_recording_count=source_recording_count, + ) + ) + return rows + + def _recording_row( + self, + *, + segment_id: str, + source: DataFile, + local_source: LocalRrd, + tables: Mapping[str, Tabular], + static: Tabular | None, + application_id: str | None, + recording_id: str | None, + source_recording_count: int | None, + ) -> DictRow: + data: dict[str, Any] = { + "episode_id": segment_id, + "rerun": RerunRecording( + segment_id=segment_id, + source_path=source.abs_path(), + tables=tables, + static=static, + source_file=source, + local_source=( + local_source if self._retain_batch_local_sources() else None + ), + application_id=application_id, + recording_id=recording_id, + source_recording_count=source_recording_count, + contents=self.contents, + timelines=self.timelines, + include_static=self.include_static, + use_source_chunks=self.use_source_chunks, + ), + } + self._with_file_path(data, source) + return DictRow(data) + + def _timelines(self, schema: Any) -> tuple[str, ...]: + return tuple(str(index.name) for index in schema.index_columns()) + + def _primary_timeline(self, timelines: Sequence[str]) -> str: + if self.primary_timeline is not None: + return self.primary_timeline + for timeline in timelines: + if timeline not in {"log_tick", "log_time", "real_time"}: + return timeline + if not timelines: + raise ValueError("Rerun recording has no timelines") + return timelines[0] + + def _view_for_contents(self, dataset_or_view: Any) -> Any: + if self.contents is None: + return dataset_or_view + return dataset_or_view.filter_contents(self.contents) + + def _robotics_contents(self) -> tuple[str, ...]: + if self.contents is not None: + return self.contents + contents: list[str] = [] + if self.actions_explicit: + contents.extend(self.actions.values()) + else: + contents.append(_prefix_contents(self.action_prefix)) + if self.states_explicit: + contents.extend(self.states.values()) + else: + contents.append(_prefix_contents(self.state_prefix)) + if self.videos_explicit: + contents.extend(self.videos.values()) + else: + contents.append(_prefix_contents(self.camera_prefix)) + return tuple(dict.fromkeys(contents)) + + def _robotics_row( + self, + view: Any, + *, + segment_id: str, + source_path: str, + source_file: DataFile, + local_source: LocalRrd, + application_id: str | None, + recording_id: str, + timelines: Sequence[str], + ) -> DictRow: + timeline = self._primary_timeline(timelines) + contents = self._robotics_contents() + content_view = view.filter_contents(contents) + table = _collect_table( + content_view.reader( + index=timeline, + fill_latest_at=self.fill_latest_at, + ) + ) + frames = _robotics_frame_table(table, timeline=timeline) + row: dict[str, Any] = { + "episode_id": segment_id, + } + if self.include_recording: + static = ( + _collect_table(content_view.reader(index=None)) + if self.include_static + else None + ) + row["rerun"] = RerunRecording( + segment_id=segment_id, + source_path=source_path, + tables={timeline: Tabular(table)}, + static=Tabular(static) if static is not None else None, + source_file=source_file, + local_source=( + local_source if self._retain_batch_local_sources() else None + ), + application_id=application_id, + recording_id=recording_id, + contents=tuple(contents), + timelines=(timeline,), + include_static=self.include_static, + use_source_chunks=False, + ) + if self.fps is not None: + row["fps"] = self.fps + if self.robot_type is not None: + row["robot_type"] = self.robot_type + + component_columns = _component_column_maps(table) + scalar_columns = component_columns.get("Scalars:scalars", {}) + action_columns = ( + _selected_columns( + scalar_columns, + self.actions, + format_name="Rerun action", + ) + if self.actions_explicit + else _prefixed_columns(scalar_columns, self.action_prefix) + ) + state_columns = ( + _selected_columns( + scalar_columns, + self.states, + format_name="Rerun state", + ) + if self.states_explicit + else _prefixed_columns(scalar_columns, self.state_prefix) + ) + if action_columns: + frames = frames.append_column( + "action", + _list_column(_singleton_scalar_matrix(table, action_columns)), + ) + if state_columns: + frames = frames.append_column( + "observation.state", + _list_column(_singleton_scalar_matrix(table, state_columns)), + ) + row["frames"] = Tabular(frames) + + image_columns = component_columns.get("EncodedImage:blob", {}) + camera_columns = ( + _selected_camera_columns(image_columns, self.videos) + if self.videos_explicit + else _camera_columns(image_columns, self.camera_prefix) + ) + reserved_video_names = set(_ROBOTICS_ROW_COLUMNS) + if self.file_path_column is not None: + reserved_video_names.add(self.file_path_column) + _validate_video_output_names(camera_columns, reserved=reserved_video_names) + for name, column in camera_columns.items(): + values = table.column(column).combine_chunks() + _require_dense_encoded_images(values, video_name=name) + row[name] = VideoFrameSequence( + lambda values=values: _iter_encoded_images(values), + fps=self.fps or 30.0, + frame_count=len(values), + ) + + if self.file_path_column is not None: + row[self.file_path_column] = source_path + return DictRow(row) + + def _retain_batch_local_sources(self) -> bool: + return ( + self.output == "recording" + and not self.materialize_tables + and self.use_source_chunks + ) + + +def _contents(contents: str | Sequence[str] | None) -> tuple[str, ...] | None: + if contents is None: + return None + if isinstance(contents, str): + return (contents,) + return tuple(contents) + + +def _normalize_entity_prefix(value: str) -> str: + value = value.strip() + if not value: + raise ValueError("Rerun entity prefixes must be non-empty") + stripped = value.strip("/") + return "/" if not stripped else "/" + stripped + + +def _prefix_contents(prefix: str) -> str: + return "/**" if prefix == "/" else f"{prefix}/**" + + +def _selection_map( + value: PathSelection | None, + *, + format_name: str, + derive_names_from_paths: bool = True, +) -> dict[str, str]: + return { + name: _normalize_entity_prefix(path) + for name, path in path_selection_map( + value, + format_name=format_name, + derive_names_from_paths=derive_names_from_paths, + ).items() + } + + +def _collect_table(df: Any) -> pa.Table: + table = df.to_arrow_table() + if _RERUN_SEGMENT_ID in table.column_names and table.num_rows > 0: + segment_ids = table.column(_RERUN_SEGMENT_ID) + if segment_ids.null_count: + table = table.filter(_is_valid(segment_ids)) + return table + + +def _close_local_sources( + local_files: Iterable[tuple[DataFile, Path, LocalRrd]], +) -> None: + for _source, _local_path, local_source in local_files: + local_source.close() + + +def _open_local_sources( + local_files: Sequence[tuple[DataFile, LocalRrd]], +) -> list[tuple[DataFile, Path, LocalRrd]]: + if len(local_files) <= 1: + return [ + (source, local_source.open(), local_source) + for source, local_source in local_files + ] + + local_sources = [local_source for _source, local_source in local_files] + max_workers = min(8, len(local_files)) + try: + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool: + local_paths = list(pool.map(_open_local_source, local_sources)) + except BaseException: + for local_source in local_sources: + local_source.close() + raise + return [ + (source, local_path, local_source) + for (source, local_source), local_path in zip( + local_files, local_paths, strict=True + ) + ] + + +def _open_local_source(local_source: LocalRrd) -> Path: + return local_source.open() + + +def _scan_recording_entries( + local_files: Sequence[tuple[DataFile, Path, LocalRrd]], +) -> list[tuple[DataFile, Path, LocalRrd, list[Any]]]: + if len(local_files) <= 1: + return [ + (source, local_path, local_source, _recording_entries(local_path)) + for source, local_path, local_source in local_files + ] + + local_paths = [local_path for _source, local_path, _local_source in local_files] + max_workers = min(8, len(local_files)) + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool: + store_entries = list(pool.map(_recording_entries, local_paths)) + return [ + (source, local_path, local_source, store_entry) + for (source, local_path, local_source), store_entry in zip( + local_files, store_entries, strict=True + ) + ] + + +def _recording_entries(local_path: Path) -> list[Any]: + import rerun as rr + + try: + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message="RRD file has no footer/manifest:.*", + ) + reader = rr.experimental.RrdReader(local_path) + internal = getattr(reader, "_internal", None) + store_entries = ( + internal.store_entries() + if internal is not None + and callable(getattr(internal, "store_entries", None)) + else reader.recordings() + ) + return [entry for entry in store_entries if entry.kind == "recording"] + except Exception as err: + logger.warning( + "Rerun recording metadata unavailable; falling back to server scan: {}", + type(err).__name__, + ) + return [] + + +def _metadata_text(metadata: Mapping[bytes, bytes], key: bytes) -> str | None: + value = metadata.get(key) + return value.decode("utf-8") if value is not None else None + + +def _component_column_maps(table: pa.Table) -> dict[str, dict[str, str]]: + by_component: dict[str, dict[str, str]] = {} + for field in table.schema: + metadata = field.metadata or {} + field_component = _metadata_text(metadata, _RERUN_COMPONENT_METADATA) + if field_component is None: + continue + entity_path = _metadata_text(metadata, _RERUN_ENTITY_PATH_METADATA) + if entity_path is not None: + by_component.setdefault(field_component, {})[entity_path] = field.name + return by_component + + +def _prefixed_columns(columns: Mapping[str, str], prefix: str) -> list[tuple[str, str]]: + return sorted( + ( + (path, column) + for path, column in columns.items() + if _matches_entity_prefix(path, prefix) + ), + key=lambda item: item[0], + ) + + +def _selected_columns( + columns: Mapping[str, str], + selected: Mapping[str, str], + *, + format_name: str, +) -> list[tuple[str, str]]: + out: list[tuple[str, str]] = [] + for _name, path in selected.items(): + column = columns.get(path) + if column is None: + raise KeyError(f"{format_name} entity path not found: {path}") + out.append((path, column)) + return out + + +def _matches_entity_prefix(entity_path: str, prefix: str) -> bool: + if prefix == "/": + return entity_path.startswith("/") + return entity_path == prefix or entity_path.startswith(f"{prefix}/") + + +def _camera_columns(columns: Mapping[str, str], prefix: str) -> dict[str, str]: + out: dict[str, str] = {} + output_paths: dict[str, str] = {} + for entity_path, column in columns.items(): + if not _matches_entity_prefix(entity_path, prefix): + continue + name = entity_path.strip("/").replace("/", ".") + existing = output_paths.get(name) + if existing is not None: + raise ValueError( + "Rerun camera paths derive the same output video name " + f"{name!r}: {existing!r} and {entity_path!r}; pass explicit " + "videos={...} to choose unique names" + ) + output_paths[name] = entity_path + out[name] = column + return out + + +def _selected_camera_columns( + by_entity_path: Mapping[str, str], + selected: Mapping[str, str], +) -> dict[str, str]: + out: dict[str, str] = {} + for name, path in selected.items(): + column = by_entity_path.get(path) + if column is None: + raise KeyError(f"Rerun video entity path not found: {path}") + out[name] = column + return out + + +def _validate_video_output_names( + selected: Mapping[str, str], + *, + reserved: set[str], +) -> None: + collisions = set(selected).intersection(reserved) + if collisions: + raise ValueError( + "Rerun video output names cannot use reserved robotics row columns: " + + ", ".join(sorted(collisions)) + ) + + +def _robotics_frame_table(table: pa.Table, *, timeline: str) -> pa.Table: + columns: dict[str, pa.ChunkedArray] = {} + if timeline in table.column_names: + columns["frame_index"] = table.column(timeline) + return pa.table(columns) + + +def _singleton_scalar_matrix( + table: pa.Table, + columns: Sequence[tuple[str, str]], +) -> np.ndarray: + values = np.full((table.num_rows, len(columns)), np.nan, dtype=np.float64) + for index, (_, column) in enumerate(columns): + _fill_singleton_list_array( + table.column(column).combine_chunks(), values[:, index] + ) + return values + + +def _list_column(values: np.ndarray) -> pa.Array: + if values.ndim != 2: + raise ValueError("Rerun vector columns must be 2D") + width = int(values.shape[1]) + if width <= 0: + offsets = pa.array(np.zeros(values.shape[0] + 1, dtype=np.int32)) + return pa.ListArray.from_arrays(offsets, pa.array([], type=pa.float64())) + flat_values = pa.array( + np.ascontiguousarray(values).reshape(-1), + type=pa.float64(), + ) + offsets = pa.array( + np.arange(0, len(flat_values) + width, width, dtype=np.int32), + ) + return pa.ListArray.from_arrays(offsets, flat_values) + + +def _fill_singleton_list_array(array: pa.Array, out: np.ndarray) -> None: + if len(array) != len(out): + raise ValueError("Rerun vector column length mismatch") + if len(array) == 0: + return + if not pa.types.is_list(array.type) and not pa.types.is_large_list(array.type): + raise TypeError(f"Expected a Rerun list component column, got {array.type}") + offsets = np.asarray(array.offsets) + starts = offsets[:-1] + ends = offsets[1:] + valid = np.asarray(_is_valid(array), dtype=bool) & (ends > starts) + if not valid.any(): + return + values = np.asarray(array.values) + out[valid] = values[starts[valid]] + + +def _require_dense_encoded_images(values: pa.Array, *, video_name: str) -> None: + if not pa.types.is_list(values.type) and not pa.types.is_large_list(values.type): + raise TypeError( + f"Expected a Rerun encoded image list column, got {values.type}" + ) + offsets = np.asarray(values.offsets) + missing = np.asarray(_is_valid(values), dtype=bool) == 0 + if len(offsets) > 1: + missing |= offsets[1:] <= offsets[:-1] + if missing.any(): + raise ValueError( + f"Rerun video {video_name!r} has missing frames on the primary timeline; " + "use fill_latest_at=True or select a denser timeline" + ) + + +def _iter_encoded_images(values: pa.Array) -> Iterator[np.ndarray]: + from io import BytesIO + + from PIL import Image + + if not pa.types.is_list(values.type) and not pa.types.is_large_list(values.type): + raise TypeError( + f"Expected a Rerun encoded image list column, got {values.type}" + ) + outer_offsets = np.asarray(values.offsets) + valid = np.asarray(_is_valid(values), dtype=bool) + inner = values.values + if pa.types.is_list(inner.type) or pa.types.is_large_list(inner.type): + inner_offsets = np.asarray(inner.offsets) + inner_values = inner.values + for index in range(len(values)): + if not valid[index]: + continue + outer_start = int(outer_offsets[index]) + outer_end = int(outer_offsets[index + 1]) + if outer_end <= outer_start: + continue + byte_start = int(inner_offsets[outer_start]) + byte_end = int(inner_offsets[outer_start + 1]) + data = np.asarray( + inner_values.slice(byte_start, byte_end - byte_start) + ).tobytes() + with Image.open(BytesIO(data)) as image: + yield np.asarray(image.convert("RGB"), dtype=np.uint8) + return + + for index in range(len(values)): + if not valid[index]: + continue + outer_start = int(outer_offsets[index]) + outer_end = int(outer_offsets[index + 1]) + if outer_end <= outer_start: + continue + value = values[index].as_py() + if not value: + continue + data = bytes(cast(bytes | bytearray | list[int], value[0])) + with Image.open(BytesIO(data)) as image: + yield np.asarray(image.convert("RGB"), dtype=np.uint8) + + +def _is_valid(values: pa.Array | pa.ChunkedArray) -> pa.Array | pa.ChunkedArray: + return pc.call_function("is_valid", [values]) + + +__all__ = [ + "DEFAULT_RERUN_ACTION_PREFIX", + "DEFAULT_RERUN_CAMERA_PREFIX", + "DEFAULT_RERUN_STATE_PREFIX", + "RerunReader", + "RerunRecording", + "RerunOutputMode", +] diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py new file mode 100644 index 00000000..2bc0e615 --- /dev/null +++ b/tests/readers/test_rerun_reader.py @@ -0,0 +1,588 @@ +from __future__ import annotations + +from io import BytesIO +from pathlib import Path +from typing import Any, cast + +import numpy as np +import pytest + +import refiner as mdr +from refiner.pipeline import Row +from refiner.pipeline.data.row import DictRow +from refiner.pipeline.sources.readers.rerun import RerunReader + +pytest.importorskip("rerun") + + +def _tiny_rrd(path: Path) -> None: + import rerun as rr + + rr.init("refiner_rerun_test", recording_id="episode-a") + rr.save(path) + frames = np.arange(3) + rr.send_columns( + "/action/x", + indexes=[rr.TimeColumn("frame", sequence=frames)], + columns=rr.Scalars.columns(scalars=np.asarray([1.0, 2.0, 3.0])), + ) + rr.send_columns( + "/action_extra/y", + indexes=[rr.TimeColumn("frame", sequence=frames)], + columns=rr.Scalars.columns(scalars=np.asarray([10.0, 20.0, 30.0])), + ) + rr.send_columns( + "/observation/state/y", + indexes=[rr.TimeColumn("frame", sequence=frames)], + columns=rr.Scalars.columns(scalars=np.asarray([4.0, 5.0, 6.0])), + ) + + +def _sparse_rrd(path: Path) -> None: + import rerun as rr + + rr.init("refiner_rerun_sparse_test", recording_id="episode-sparse") + rr.save(path) + rr.send_columns( + "/action/x", + indexes=[], + columns=rr.SeriesLines.columns(names=["x"]), + ) + rr.send_columns( + "/action/x", + indexes=[rr.TimeColumn("frame", sequence=np.asarray([0, 1, 2]))], + columns=rr.Scalars.columns(scalars=np.asarray([1.0, 2.0, 3.0])), + ) + rr.send_columns( + "/action/y", + indexes=[rr.TimeColumn("frame", sequence=np.asarray([1]))], + columns=rr.Scalars.columns(scalars=np.asarray([9.0])), + ) + + +def _custom_robotics_rrd(path: Path) -> None: + import rerun as rr + from PIL import Image + + rr.init("refiner_rerun_custom_robotics_test", recording_id="episode-custom") + rr.save(path) + frames = np.arange(2) + rr.send_columns( + "/robot/actions/z", + indexes=[rr.TimeColumn("frame", sequence=frames)], + columns=rr.Scalars.columns(scalars=np.asarray([30.0, 40.0])), + ) + rr.send_columns( + "/robot/actions/a", + indexes=[rr.TimeColumn("frame", sequence=frames)], + columns=rr.Scalars.columns(scalars=np.asarray([10.0, 20.0])), + ) + rr.send_columns( + "/robot/state/b", + indexes=[rr.TimeColumn("frame", sequence=frames)], + columns=rr.Scalars.columns(scalars=np.asarray([1.0, 2.0])), + ) + rr.send_columns( + "/robot/state/a", + indexes=[rr.TimeColumn("frame", sequence=frames)], + columns=rr.Scalars.columns(scalars=np.asarray([3.0, 4.0])), + ) + blobs: list[bytes] = [] + for color in ((1, 2, 3), (4, 5, 6)): + image = Image.new("RGB", (1, 1), color=color) + out = BytesIO() + image.save(out, format="PNG") + blobs.append(out.getvalue()) + rr.send_columns( + "/robot/cameras/top", + indexes=[rr.TimeColumn("frame", sequence=frames)], + columns=rr.EncodedImage.columns( + blob=blobs, + media_type=["image/png", "image/png"], + ), + ) + + +def _sparse_camera_rrd(path: Path) -> None: + import rerun as rr + from PIL import Image + + rr.init("refiner_rerun_sparse_camera_test", recording_id="episode-sparse-camera") + rr.save(path) + rr.send_columns( + "/robot/actions/x", + indexes=[rr.TimeColumn("frame", sequence=np.asarray([0, 1]))], + columns=rr.Scalars.columns(scalars=np.asarray([1.0, 2.0])), + ) + out = BytesIO() + Image.new("RGB", (1, 1), color=(1, 2, 3)).save(out, format="PNG") + rr.send_columns( + "/robot/cameras/top", + indexes=[rr.TimeColumn("frame", sequence=np.asarray([0]))], + columns=rr.EncodedImage.columns( + blob=[out.getvalue()], + media_type=["image/png"], + ), + ) + + +def _reserved_video_name_rrd(path: Path) -> None: + import rerun as rr + from PIL import Image + + rr.init("refiner_rerun_reserved_video_name_test", recording_id="episode-reserved") + rr.save(path) + out = BytesIO() + Image.new("RGB", (1, 1), color=(1, 2, 3)).save(out, format="PNG") + rr.send_columns( + "/frames", + indexes=[rr.TimeColumn("frame", sequence=np.asarray([0]))], + columns=rr.EncodedImage.columns( + blob=[out.getvalue()], + media_type=["image/png"], + ), + ) + + +def _colliding_camera_names_rrd(path: Path) -> None: + import rerun as rr + from PIL import Image + + rr.init("refiner_rerun_colliding_camera_names_test", recording_id="episode-cameras") + rr.save(path) + frames = np.arange(1) + blobs: list[bytes] = [] + for color in ((1, 2, 3), (4, 5, 6)): + out = BytesIO() + Image.new("RGB", (1, 1), color=color).save(out, format="PNG") + blobs.append(out.getvalue()) + for entity_path, blob in zip(("/cam/a.b", "/cam/a/b"), blobs, strict=True): + rr.send_columns( + entity_path, + indexes=[rr.TimeColumn("frame", sequence=frames)], + columns=rr.EncodedImage.columns( + blob=[blob], + media_type=["image/png"], + ), + ) + + +def test_read_rerun_recording_preserves_timeline_table(tmp_path: Path) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + unit = next(mdr.read_rerun(str(rrd), timelines=("frame",)).source.read()) + assert isinstance(unit, Row) + row = cast(Any, unit) + recording = row["rerun"] + + assert row["episode_id"] == "episode-a" + assert list(recording.tables) == ["frame"] + assert recording.tables["frame"].num_rows == 3 + assert recording.source_path == str(rrd) + assert row["file_path"] == str(rrd) + + +def test_read_rerun_recording_can_project_before_arrow_conversion( + tmp_path: Path, +) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + selected = ( + mdr.read_rerun(str(rrd), timelines=("frame",)).select("episode_id").take(1)[0] + ) + dropped = mdr.read_rerun(str(rrd), timelines=("frame",)).drop("rerun").take(1)[0] + + assert selected.to_dict() == {"episode_id": "episode-a"} + assert dropped["episode_id"] == "episode-a" + assert "rerun" not in dropped + + +def test_read_rerun_recording_rejects_reserved_file_path_column( + tmp_path: Path, +) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + with pytest.raises(ValueError, match="reserved Rerun recording row column 'rerun'"): + mdr.read_rerun(str(rrd), file_path_column="rerun") + + +def test_read_rerun_recording_preserves_sparse_rows(tmp_path: Path) -> None: + rrd = tmp_path / "sparse.rrd" + _sparse_rrd(rrd) + + row = cast(Any, next(mdr.read_rerun(str(rrd), timelines=("frame",)).source.read())) + table = row["rerun"].tables["frame"].table + + assert table.num_rows == 3 + assert table.column("frame").to_pylist() == [0, 1, 2] + + +def test_read_rerun_recording_can_skip_table_materialization(tmp_path: Path) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + row = cast( + Any, + next( + mdr.read_rerun( + str(rrd), + timelines=("frame",), + materialize_tables=False, + ).source.read() + ), + ) + recording = row["rerun"] + + assert recording.recording_id == "episode-a" + assert recording.tables == {} + assert recording.static is None + assert recording.source_file is not None + assert recording.timelines == ("frame",) + assert recording.use_source_chunks is False + + +def test_read_rerun_recording_without_materialized_tables_scans_metadata_once( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + calls = 0 + + def fake_recording_entries(*args: Any, **kwargs: Any) -> list[Any]: + nonlocal calls + del args, kwargs + calls += 1 + return [ + type( + "Store", + (), + {"recording_id": "episode-a", "application_id": "refiner"}, + )() + ] + + monkeypatch.setattr( + "refiner.pipeline.sources.readers.rerun._recording_entries", + fake_recording_entries, + ) + + row = cast( + Any, + next( + mdr.read_rerun( + str(rrd), + materialize_tables=False, + ).source.read() + ), + ) + + assert isinstance(row, list) + assert row[0]["episode_id"] == "episode-a" + assert calls == 1 + + +def test_read_rerun_rejects_ignored_mode_options(tmp_path: Path) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + with pytest.raises( + ValueError, + match="materialize_tables=False is only supported for recording output", + ): + mdr.read_rerun(str(rrd), output="robotics", materialize_tables=False) + + with pytest.raises( + ValueError, + match="include_recording=False is only supported for robotics output", + ): + mdr.read_rerun(str(rrd), output="recording", include_recording=False) + with pytest.raises( + ValueError, + match="Rerun recording output does not use robotics options: primary_timeline", + ): + mdr.read_rerun(str(rrd), output="recording", primary_timeline="frame") + with pytest.raises( + ValueError, + match="Rerun recording output does not use robotics options: actions, fps", + ): + mdr.read_rerun( + str(rrd), + output="recording", + actions=("/action/x",), + fps=30.0, + ) + + +def test_read_rerun_recording_without_materialized_tables_skips_server( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + monkeypatch.setattr( + "refiner.pipeline.sources.readers.rerun.RerunReader._read_files_with_server", + lambda *args, **kwargs: pytest.fail( + "metadata-only recording rows do not need the Rerun server" + ), + ) + + row = cast( + Any, + next( + mdr.read_rerun( + str(rrd), + timelines=("frame",), + materialize_tables=False, + ).source.read() + ), + ) + + assert row["episode_id"] == "episode-a" + assert row["rerun"].tables == {} + + +def test_read_rerun_batches_small_files_in_one_staged_read( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + first = tmp_path / "first.rrd" + second = tmp_path / "second.rrd" + first.write_bytes(b"first") + second.write_bytes(b"second") + reader = RerunReader( + [str(first), str(second)], + target_shard_bytes=1024 * 1024, + ) + batch_sizes: list[int] = [] + + def fake_read_files(self: RerunReader, local_files: Any) -> Any: + del self + local_files = tuple(local_files) + batch_sizes.append(len(local_files)) + for source, local_path, _local_rrd in local_files: + assert local_path.exists() + yield DictRow({"source": source.abs_path()}) + + monkeypatch.setattr(RerunReader, "_read_files", fake_read_files) + + rows = cast(list[Row], list(reader.read_shard(reader.list_shards()[0]))) + + assert batch_sizes == [2] + assert [row["source"] for row in rows] == [str(first), str(second)] + + +def test_read_rerun_robotics_mode_converts_to_robot_row(tmp_path: Path) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + row = cast( + Any, + mdr.read_rerun(str(rrd), output="robotics", fps=30.0) + .to_robot_rows( + episode_id_key="episode_id", + nested_frames_key="frames", + fps=30.0, + ) + .take(1)[0], + ) + + assert "rerun" not in row + assert row.episode_id == "episode-a" + assert row.num_frames == 3 + assert row.actions.to_pylist() == [[1.0], [2.0], [3.0]] + assert row.states.to_pylist() == [[4.0], [5.0], [6.0]] + + +def test_read_rerun_robotics_mode_without_recording_skips_recording_entries( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + monkeypatch.setattr( + "refiner.pipeline.sources.readers.rerun._recording_entries", + lambda *args, **kwargs: pytest.fail( + "robotics rows without recording payload do not need store metadata" + ), + ) + + row = cast( + Any, + mdr.read_rerun(str(rrd), output="robotics", fps=30.0).take(1)[0], + ) + + assert "rerun" not in row + assert row["frames"].num_rows == 3 + + +def test_read_rerun_robotics_mode_with_recording_includes_recording_payload( + tmp_path: Path, +) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + row = cast( + Any, + mdr.read_rerun( + str(rrd), + output="robotics", + include_recording=True, + timelines=("frame",), + fps=30.0, + ).take(1)[0], + ) + + recording = row["rerun"] + assert recording.recording_id == "episode-a" + assert list(recording.tables) == ["frame"] + assert recording.tables["frame"].num_rows == 3 + assert row["frames"].num_rows == 3 + + +def test_read_rerun_describe_includes_robotics_metadata(tmp_path: Path) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + description = mdr.read_rerun( + str(rrd), + output="robotics", + fps=12.5, + robot_type="testbot", + ).source.describe() + + assert description["fps"] == 12.5 + assert description["robot_type"] == "testbot" + + +def test_read_rerun_robotics_rejects_reserved_implicit_video_name( + tmp_path: Path, +) -> None: + rrd = tmp_path / "reserved-video.rrd" + _reserved_video_name_rrd(rrd) + + with pytest.raises( + ValueError, + match="Rerun video output names cannot use reserved robotics row columns: frames", + ): + mdr.read_rerun( + str(rrd), + output="robotics", + camera_prefix="/", + fps=30.0, + ).take(1) + + +def test_read_rerun_robotics_rejects_derived_video_name_collision( + tmp_path: Path, +) -> None: + rrd = tmp_path / "colliding-cameras.rrd" + _colliding_camera_names_rrd(rrd) + + with pytest.raises( + ValueError, + match="derive the same output video name", + ): + mdr.read_rerun( + str(rrd), + output="robotics", + camera_prefix="/cam", + fps=30.0, + ).take(1) + + +def test_read_rerun_robotics_mode_respects_explicit_selections( + tmp_path: Path, +) -> None: + rrd = tmp_path / "custom.rrd" + _custom_robotics_rrd(rrd) + + row = cast( + Any, + mdr.read_rerun( + str(rrd), + output="robotics", + actions=("/robot/actions/z", "/robot/actions/a"), + states={ + "first": "/robot/state/a", + "second": "/robot/state/b", + }, + videos={"observation.images.top": "/robot/cameras/top"}, + fps=5.0, + ).take(1)[0], + ) + + assert row["frames"].column("action").to_pylist() == [[30.0, 10.0], [40.0, 20.0]] + assert row["frames"].column("observation.state").to_pylist() == [ + [3.0, 1.0], + [4.0, 2.0], + ] + video = row["observation.images.top"] + assert video.frame_count == 2 + assert [frame[0, 0].tolist() for frame in video.iter_frame_arrays()] == [ + [1, 2, 3], + [4, 5, 6], + ] + + +def test_read_rerun_robotics_mode_rejects_sparse_video_stream( + tmp_path: Path, +) -> None: + rrd = tmp_path / "sparse-camera.rrd" + _sparse_camera_rrd(rrd) + + with pytest.raises( + ValueError, + match="missing frames on the primary timeline", + ): + mdr.read_rerun( + str(rrd), + output="robotics", + timelines=("frame",), + actions=("/robot/actions/x",), + videos={"observation.images.top": "/robot/cameras/top"}, + fps=5.0, + ).take(1) + + +def test_read_rerun_robotics_mode_with_explicit_timeline_uses_table_metadata( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + rrd = tmp_path / "custom.rrd" + _custom_robotics_rrd(rrd) + + monkeypatch.setattr( + "refiner.pipeline.sources.readers.rerun.RerunReader._timelines", + lambda *args, **kwargs: pytest.fail( + "explicit timelines should not require schema timeline discovery" + ), + ) + + row = cast( + Any, + mdr.read_rerun( + str(rrd), + output="robotics", + timelines=("frame",), + actions=("/robot/actions/z", "/robot/actions/a"), + states={ + "first": "/robot/state/a", + "second": "/robot/state/b", + }, + videos={"observation.images.top": "/robot/cameras/top"}, + fps=5.0, + ).take(1)[0], + ) + + assert row["frames"].column("action").to_pylist() == [[30.0, 10.0], [40.0, 20.0]] + assert row["frames"].column("observation.state").to_pylist() == [ + [3.0, 1.0], + [4.0, 2.0], + ] + video = row["observation.images.top"] + assert video.frame_count == 2 diff --git a/uv.lock b/uv.lock index 9ade185d..06335d65 100644 --- a/uv.lock +++ b/uv.lock @@ -867,6 +867,23 @@ nvtx = [ { name = "nvidia-nvtx", marker = "sys_platform == 'linux'" }, ] +[[package]] +name = "datafusion" +version = "52.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyarrow" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/db/d4/a5ad7b665a80008901892fde61dc667318db0652a955d706ddca3a224b5a/datafusion-52.3.0.tar.gz", hash = "sha256:2e8b02ad142b1a0d673f035d96a0944a640ac78275003d7e453cee4afe4a20a4", size = 205026, upload-time = "2026-03-16T10:54:07.739Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/55/63/1bb0737988cefa77274b459d64fa4b57ba4cf755639a39733e9581b5d599/datafusion-52.3.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:a73f02406b2985b9145dd97f8221a929c9ef3289a8ba64c6b52043e240938528", size = 31503230, upload-time = "2026-03-16T10:53:50.312Z" }, + { url = "https://files.pythonhosted.org/packages/d6/e3/ea3b79239953c3044d19d8e9581015da025b6640796db03799e435b17910/datafusion-52.3.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:118a1f0add6a3f91fcbc90c71819fe08750e2981637d5e7b346e099e94a20b8b", size = 28159497, upload-time = "2026-03-16T10:53:54.032Z" }, + { url = "https://files.pythonhosted.org/packages/24/c8/7d325feb4b7509ae03857fd7e164e95ec72e8c9f3dfd3178ec7f80d53977/datafusion-52.3.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:253ce7aee5fe84bd6ee290c20608114114bdb5115852617f97d3855d36ad9341", size = 30769154, upload-time = "2026-03-16T10:53:57.835Z" }, + { url = "https://files.pythonhosted.org/packages/37/ee/478689c69b3cb1ccabb2d52feac0c181f6cdf20b51a81df35344b1dab9a6/datafusion-52.3.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:2af3469d2f06959bec88579ab107a72f965de18b32e607069bbdd0b859ed8dbb", size = 33060335, upload-time = "2026-03-16T10:54:01.715Z" }, + { url = "https://files.pythonhosted.org/packages/1c/48/01906ab5c1a70373c6874ac5192d03646fa7b94d9ff06e3f676cb6b0f43f/datafusion-52.3.0-cp310-abi3-win_amd64.whl", hash = "sha256:9fb35738cf4dbff672dbcfffc7332813024cb0ad2ab8cda1fb90b9054277ab0c", size = 33765807, upload-time = "2026-03-16T10:54:05.728Z" }, +] + [[package]] name = "datasets" version = "4.8.5" @@ -2085,6 +2102,7 @@ all = [ { name = "mcap-ros2-support" }, { name = "numcodecs" }, { name = "pillow" }, + { name = "rerun-sdk", extra = ["datafusion"] }, { name = "s3fs" }, { name = "tensorflow" }, { name = "tensorflow-datasets" }, @@ -2120,6 +2138,10 @@ mcap = [ { name = "mcap-ros2-support" }, { name = "pillow" }, ] +rerun = [ + { name = "pillow" }, + { name = "rerun-sdk", extra = ["datafusion"] }, +] s3 = [ { name = "s3fs" }, ] @@ -2140,6 +2162,7 @@ testing = [ { name = "pillow" }, { name = "pytest" }, { name = "pytest-cov" }, + { name = "rerun-sdk", extra = ["datafusion"] }, { name = "s3fs" }, { name = "tensorflow" }, { name = "tensorflow-datasets" }, @@ -2193,6 +2216,7 @@ requires-dist = [ { name = "macrodata-refiner", extras = ["hf"], marker = "extra == 'datasets'" }, { name = "macrodata-refiner", extras = ["hf"], marker = "extra == 'hand-tracking'" }, { name = "macrodata-refiner", extras = ["mcap"], marker = "extra == 'all'" }, + { name = "macrodata-refiner", extras = ["rerun"], marker = "extra == 'all'" }, { name = "macrodata-refiner", extras = ["s3"], marker = "extra == 'all'" }, { name = "macrodata-refiner", extras = ["tensorflow"], marker = "extra == 'tfds'" }, { name = "macrodata-refiner", extras = ["text"], marker = "extra == 'all'" }, @@ -2209,18 +2233,20 @@ requires-dist = [ { name = "orjson" }, { name = "packaging" }, { name = "pillow", marker = "extra == 'mcap'" }, + { name = "pillow", marker = "extra == 'rerun'" }, { name = "pillow", marker = "extra == 'video'" }, { name = "pyarrow" }, { name = "pydantic", specifier = ">=2.0.0" }, { name = "pytest", marker = "extra == 'testing'", specifier = ">=8.0.0" }, { name = "pytest-cov", marker = "extra == 'testing'", specifier = ">=5.0.0" }, + { name = "rerun-sdk", extras = ["datafusion"], marker = "extra == 'rerun'", specifier = ">=0.33,<0.34" }, { name = "s3fs", marker = "extra == 's3'" }, { name = "tensorflow", marker = "extra == 'tensorflow'" }, { name = "tensorflow-datasets", marker = "extra == 'tfds'" }, { name = "warcio", marker = "extra == 'text'" }, { name = "zarr", marker = "extra == 'zarr'", specifier = ">=2.18,<3" }, ] -provides-extras = ["video", "hf", "datasets", "hand-tracking", "text", "hdf5", "zarr", "mcap", "s3", "gcs", "tensorflow", "tfds", "testing", "all"] +provides-extras = ["video", "hf", "datasets", "hand-tracking", "text", "hdf5", "zarr", "mcap", "rerun", "s3", "gcs", "tensorflow", "tfds", "testing", "all"] [package.metadata.requires-dev] dev = [ @@ -4150,6 +4176,33 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3b/5d/63d4ae3b9daea098d5d6f5da83984853c1bbacd5dc826764b249fe119d24/requests_oauthlib-2.0.0-py2.py3-none-any.whl", hash = "sha256:7dd8a5c40426b779b0868c404bdef9768deccf22749cde15852df527e6269b36", size = 24179, upload-time = "2024-03-22T20:32:28.055Z" }, ] +[[package]] +name = "rerun-sdk" +version = "0.33.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "attrs" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "pillow" }, + { name = "psutil" }, + { name = "pyarrow" }, + { name = "typing-extensions" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/31/17/5a521e86ac0064bd0f452e3e98e2422433511b54110423c0217d2cc1234f/rerun_sdk-0.33.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:97f123e3ef6aa69b60194bc566e5435c7d4040757ed4f58297ea46c8ef320c5c", size = 125707606, upload-time = "2026-05-29T09:42:53.584Z" }, + { url = "https://files.pythonhosted.org/packages/34/2f/2ca2599aca03b69fbcac7c8391ef50376968edd7c58b96de53a4b7f20624/rerun_sdk-0.33.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:8f734cf59419dcfbc46915bea6cec030224f16e96c3a597f0ccf7cb7b058dd43", size = 135271020, upload-time = "2026-05-29T09:43:00.106Z" }, + { url = "https://files.pythonhosted.org/packages/2e/ba/d70997b43e6db4f58c4326c29c6a6a384ddc6c2fe125f231c885ad9b3b1f/rerun_sdk-0.33.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:53d95609f8b330026bcd041bf6d11b46ee1c18b6fbde155135f291fe86328eeb", size = 139552018, upload-time = "2026-05-29T09:43:06.275Z" }, + { url = "https://files.pythonhosted.org/packages/14/a5/0cac294d16aff6c9a2f183f838428a0380b4d2fd9e053bb37b3041999ad5/rerun_sdk-0.33.0-cp310-abi3-win_amd64.whl", hash = "sha256:b152992a72ec240062c8c285bd30ab681b464a25efbe1464c66fdac82320de1f", size = 120418186, upload-time = "2026-05-29T09:43:13.733Z" }, +] + +[package.optional-dependencies] +datafusion = [ + { name = "datafusion" }, + { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "pandas", version = "3.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, +] + [[package]] name = "rich" version = "14.3.3" From 80251ade52882516b61cd420c403eb9fc552865d Mon Sep 17 00:00:00 2001 From: guipenedo Date: Wed, 17 Jun 2026 23:33:10 +0200 Subject: [PATCH 2/4] Emit Rerun robotics rows directly --- docs/reading-data/rerun.md | 43 ++++----- src/refiner/execution/engine.py | 26 +++--- src/refiner/pipeline/pipeline.py | 11 +-- src/refiner/pipeline/sources/readers/rerun.py | 43 +++++---- src/refiner/robotics/tabular.py | 4 +- tests/readers/test_rerun_reader.py | 90 +++++++++++++++---- 6 files changed, 139 insertions(+), 78 deletions(-) diff --git a/docs/reading-data/rerun.md b/docs/reading-data/rerun.md index 6ce1a86a..9c1b8981 100644 --- a/docs/reading-data/rerun.md +++ b/docs/reading-data/rerun.md @@ -57,34 +57,25 @@ identity and chunk metadata. ## Robotics rows -With `output="robotics"`, the reader creates rows that can be passed to -`to_robot_rows(...)` and robotics writers: +With `output="robotics"`, the reader creates robotics episode rows directly: ```python -robot_rows = ( - mdr.read_rerun( - "/data/episodes/*.rrd", - output="robotics", - fps=30, - robot_type="unknown", - ) - .to_robot_rows( - episode_id_key="episode_id", - nested_frames_key="frames", - fps_key="fps", - robot_type_key="robot_type", - video_keys={ - "observation.images.top": "cam.top", - "observation.images.left_wrist": "cam.left_wrist", - }, - ) +robot_rows = mdr.read_rerun( + "/data/episodes/*.rrd", + output="robotics", + fps=30, + robot_type="unknown", ) ``` The default robotics mapping reads scalar components under `/action/**` into -the frame `action` vector, scalar components under `/observation/state/**` into -`observation.state`, and encoded images under `/cam/**` into top-level video -sources such as `cam.top`. +the top-level `action` vector, scalar components under +`/observation/state/**` into top-level `observation.state`, and encoded images +under `/cam/**` into top-level video sources such as `cam.top`. + +Rows also include the original `rerun` recording sidecar by default. Set +`include_recording=False` only when you want a lightweight robotics projection +and do not need the source Rerun structure later. Use explicit selections when vector order or camera names matter: @@ -105,10 +96,10 @@ also define the minimal Rerun content filter for those categories. ## Decoding -Scalar action and state columns are read from Arrow list arrays and converted -to frame vectors. Encoded images remain lazy `VideoFrameSequence` values; JPEG -or PNG bytes are decoded frame-by-frame only when a downstream video writer or -consumer iterates the sequence. +Scalar action and state columns are read from Arrow list arrays and exposed as +top-level Arrow list arrays. Encoded images remain lazy `VideoFrameSequence` +values; JPEG or PNG bytes are decoded frame-by-frame only when a downstream +video writer or consumer iterates the sequence. ## Sharding diff --git a/src/refiner/execution/engine.py b/src/refiner/execution/engine.py index b9bf4f80..7b7f2550 100644 --- a/src/refiner/execution/engine.py +++ b/src/refiner/execution/engine.py @@ -298,7 +298,7 @@ def _execute_vector_segment( max_vectorized_block_bytes: int | None, on_shard_delta: ShardDeltaFn | None, input_schema: pa.Schema | None, -) -> Iterator[Tabular]: +) -> Iterator[Block]: pending_rows = RowBuffer() current_chunk_rows = max(1, int(vectorized_chunk_rows)) estimated_row_bytes: float | None = None @@ -329,12 +329,8 @@ def _run_block(block: Tabular, block_ops: Sequence[VectorizedOp]) -> Tabular: ) return block.with_table(table, row_indices=row_indices) - def _rows_to_block(batch: list[Row]) -> tuple[Tabular, Sequence[VectorizedOp]]: - try: - return _tabular_from_rows(batch, schema=input_schema), ops - except (pa.ArrowInvalid, pa.ArrowTypeError, TypeError) as err: - if not row_projection_ops: - raise + def _rows_to_block(batch: list[Row]) -> tuple[Block, Sequence[VectorizedOp]]: + if row_projection_ops: projected = [ _apply_row_projection_ops(row, row_projection_ops) for row in batch ] @@ -343,8 +339,11 @@ def _rows_to_block(batch: list[Row]) -> tuple[Tabular, Sequence[VectorizedOp]]: _tabular_from_rows(projected, schema=row_projected_schema), row_remaining_ops, ) - except (pa.ArrowInvalid, pa.ArrowTypeError, TypeError) as fallback_err: - raise err from fallback_err + except (pa.ArrowInvalid, pa.ArrowTypeError, TypeError): + if not row_remaining_ops: + return projected, () + raise + return _tabular_from_rows(batch, schema=input_schema), ops def _chunk_rows_for_budget() -> int: if ( @@ -356,7 +355,7 @@ def _chunk_rows_for_budget() -> int: budget_rows = int(max_vectorized_block_bytes / estimated_row_bytes) return max(1, min(current_chunk_rows, budget_rows)) - def _run_pending_chunk(target_rows: int) -> Iterator[Tabular]: + def _run_pending_chunk(target_rows: int) -> Iterator[Block]: nonlocal current_chunk_rows, estimated_row_bytes rows_for_try = max(1, target_rows) while True: @@ -369,6 +368,11 @@ def _run_pending_chunk(target_rows: int) -> Iterator[Tabular]: rows_for_try = max(1, rows_for_try // 2) current_chunk_rows = min(current_chunk_rows, rows_for_try) continue + if not isinstance(block, Tabular): + pending_rows.discard(rows_for_try) + if block: + yield block + return table = block.table if table.num_rows > 0: @@ -401,7 +405,7 @@ def _run_pending_chunk(target_rows: int) -> Iterator[Tabular]: yield out return - def _drain_rows(*, force: bool) -> Iterator[Tabular]: + def _drain_rows(*, force: bool) -> Iterator[Block]: while len(pending_rows) > 0: desired_rows = _chunk_rows_for_budget() if not force and len(pending_rows) < desired_rows: diff --git a/src/refiner/pipeline/pipeline.py b/src/refiner/pipeline/pipeline.py index 0822f1b0..87b93435 100644 --- a/src/refiner/pipeline/pipeline.py +++ b/src/refiner/pipeline/pipeline.py @@ -1263,11 +1263,12 @@ def read_rerun( each emitted row preserves the selected Rerun data as Arrow-backed ``Tabular`` tables grouped by timeline under the ``rerun`` field. Set ``materialize_tables=False`` when you only need source chunk metadata. - With ``output="robotics"``, the reader additionally derives common - robotics episode fields from configurable Rerun entity prefixes so the - rows can be passed through ``to_robot_rows(...)``. Pass ``actions``, - ``states``, or ``videos`` to pin exact entity paths and output order - instead of using prefix-derived defaults. + With ``output="robotics"``, the reader emits robotics episode rows directly, + with top-level ``action``, ``observation.state``, and video fields derived + from configurable Rerun entity prefixes. The original ``rerun`` sidecar is + included by default; pass ``include_recording=False`` for a lightweight + projection. Pass ``actions``, ``states``, or ``videos`` to pin exact entity + paths and output order instead of using prefix-derived defaults. """ return RefinerPipeline( source=RerunReader( diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index 673de053..85bf7f08 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -24,6 +24,7 @@ PathSelection, path_selection_map, ) +from refiner.robotics.row import _robot_row_converter from refiner.utils import check_required_dependencies from refiner.video import VideoFrameSequence from refiner.worker.context import logger @@ -34,7 +35,15 @@ _RERUN_COMPONENT_METADATA = b"rerun:component" _RERUN_ENTITY_PATH_METADATA = b"rerun:entity_path" _ROBOTICS_ROW_COLUMNS = frozenset( - {"episode_id", "rerun", "frames", "fps", "robot_type"} + { + "episode_id", + "rerun", + "frames", + "action", + "observation.state", + "fps", + "robot_type", + } ) _RECORDING_ROW_COLUMNS = frozenset({"episode_id", "rerun"}) DEFAULT_RERUN_ACTION_PREFIX = "/action" @@ -160,7 +169,7 @@ def __init__( self.materialize_tables = materialize_tables self.use_source_chunks = self.timelines is None self.include_recording = ( - output == "recording" if include_recording is None else include_recording + True if include_recording is None else include_recording ) self.fill_latest_at = fill_latest_at self.action_prefix = _normalize_entity_prefix(action_prefix) @@ -512,7 +521,7 @@ def _robotics_row( application_id: str | None, recording_id: str, timelines: Sequence[str], - ) -> DictRow: + ) -> Row: timeline = self._primary_timeline(timelines) contents = self._robotics_contents() content_view = view.filter_contents(contents) @@ -522,7 +531,6 @@ def _robotics_row( fill_latest_at=self.fill_latest_at, ) ) - frames = _robotics_frame_table(table, timeline=timeline) row: dict[str, Any] = { "episode_id": segment_id, } @@ -574,16 +582,13 @@ def _robotics_row( else _prefixed_columns(scalar_columns, self.state_prefix) ) if action_columns: - frames = frames.append_column( - "action", - _list_column(_singleton_scalar_matrix(table, action_columns)), + row["action"] = _list_column( + _singleton_scalar_matrix(table, action_columns) ) if state_columns: - frames = frames.append_column( - "observation.state", - _list_column(_singleton_scalar_matrix(table, state_columns)), + row["observation.state"] = _list_column( + _singleton_scalar_matrix(table, state_columns) ) - row["frames"] = Tabular(frames) image_columns = component_columns.get("EncodedImage:blob", {}) camera_columns = ( @@ -606,7 +611,14 @@ def _robotics_row( if self.file_path_column is not None: row[self.file_path_column] = source_path - return DictRow(row) + return _robot_row_converter( + episode_id_key="episode_id", + fps_key="fps", + robot_type_key="robot_type", + action_key="action", + state_key="observation.state", + video_keys={name: name for name in camera_columns}, + )(DictRow(row)) def _retain_batch_local_sources(self) -> bool: return ( @@ -840,13 +852,6 @@ def _validate_video_output_names( ) -def _robotics_frame_table(table: pa.Table, *, timeline: str) -> pa.Table: - columns: dict[str, pa.ChunkedArray] = {} - if timeline in table.column_names: - columns["frame_index"] = table.column(timeline) - return pa.table(columns) - - def _singleton_scalar_matrix( table: pa.Table, columns: Sequence[tuple[str, str]], diff --git a/src/refiner/robotics/tabular.py b/src/refiner/robotics/tabular.py index d967b5e1..d4ee1b75 100644 --- a/src/refiner/robotics/tabular.py +++ b/src/refiner/robotics/tabular.py @@ -2,6 +2,7 @@ from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass +from dataclasses import is_dataclass from typing import Any, cast import pyarrow as pa @@ -165,7 +166,8 @@ def _row_index_key(rows: Sequence[Row]) -> str: def _is_side_data_value(value: Any) -> bool: return value is not _MISSING and ( - isinstance(value, Tabular) or isinstance(value, VideoSource) + isinstance(value, Tabular | VideoSource | pa.Array | pa.ChunkedArray) + or (not isinstance(value, type) and is_dataclass(value)) ) diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index 2bc0e615..c9aa6092 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -10,7 +10,9 @@ import refiner as mdr from refiner.pipeline import Row from refiner.pipeline.data.row import DictRow +from refiner.pipeline.expressions import col from refiner.pipeline.sources.readers.rerun import RerunReader +from refiner.robotics.row import RoboticsRow pytest.importorskip("rerun") @@ -381,20 +383,21 @@ def test_read_rerun_robotics_mode_converts_to_robot_row(tmp_path: Path) -> None: row = cast( Any, - mdr.read_rerun(str(rrd), output="robotics", fps=30.0) - .to_robot_rows( - episode_id_key="episode_id", - nested_frames_key="frames", - fps=30.0, - ) - .take(1)[0], + mdr.read_rerun(str(rrd), output="robotics", fps=30.0).take(1)[0], ) - assert "rerun" not in row + assert isinstance(row, RoboticsRow) + assert "rerun" in row + assert "frames" not in row assert row.episode_id == "episode-a" assert row.num_frames == 3 assert row.actions.to_pylist() == [[1.0], [2.0], [3.0]] assert row.states.to_pylist() == [[4.0], [5.0], [6.0]] + assert row.to_frame_table().column("action").to_pylist() == [ + [1.0], + [2.0], + [3.0], + ] def test_read_rerun_robotics_mode_without_recording_skips_recording_entries( @@ -413,11 +416,66 @@ def test_read_rerun_robotics_mode_without_recording_skips_recording_entries( row = cast( Any, - mdr.read_rerun(str(rrd), output="robotics", fps=30.0).take(1)[0], + mdr.read_rerun( + str(rrd), + output="robotics", + include_recording=False, + fps=30.0, + ).take(1)[0], ) assert "rerun" not in row - assert row["frames"].num_rows == 3 + assert "frames" not in row + assert row.actions.to_pylist() == [[1.0], [2.0], [3.0]] + + +def test_read_rerun_robotics_mode_exposes_top_level_fields_to_primitives( + tmp_path: Path, +) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + selected = cast( + Any, + mdr.read_rerun(str(rrd), output="robotics", fps=30.0) + .select("action") + .take(1)[0], + ) + dropped = cast( + Any, + mdr.read_rerun(str(rrd), output="robotics", fps=30.0) + .drop("observation.state") + .take(1)[0], + ) + without_recording = cast( + Any, + mdr.read_rerun(str(rrd), output="robotics", fps=30.0).drop("rerun").take(1)[0], + ) + filtered = cast( + Any, + mdr.read_rerun(str(rrd), output="robotics", fps=30.0) + .filter(col("episode_id") == "episode-a") + .take(1)[0], + ) + projected_filtered = cast( + Any, + mdr.read_rerun(str(rrd), output="robotics", fps=30.0) + .drop("rerun") + .filter(col("episode_id") == "episode-a") + .take(1)[0], + ) + + assert selected["action"].to_pylist() == [[1.0], [2.0], [3.0]] + assert "rerun" not in selected + assert "observation.state" not in dropped + assert dropped.states is None + assert dropped["action"].to_pylist() == [[1.0], [2.0], [3.0]] + assert "rerun" not in without_recording + assert without_recording.actions.to_pylist() == [[1.0], [2.0], [3.0]] + assert "rerun" in filtered + assert filtered.actions.to_pylist() == [[1.0], [2.0], [3.0]] + assert "rerun" not in projected_filtered + assert projected_filtered.actions.to_pylist() == [[1.0], [2.0], [3.0]] def test_read_rerun_robotics_mode_with_recording_includes_recording_payload( @@ -431,7 +489,6 @@ def test_read_rerun_robotics_mode_with_recording_includes_recording_payload( mdr.read_rerun( str(rrd), output="robotics", - include_recording=True, timelines=("frame",), fps=30.0, ).take(1)[0], @@ -441,7 +498,8 @@ def test_read_rerun_robotics_mode_with_recording_includes_recording_payload( assert recording.recording_id == "episode-a" assert list(recording.tables) == ["frame"] assert recording.tables["frame"].num_rows == 3 - assert row["frames"].num_rows == 3 + assert "frames" not in row + assert row.actions.to_pylist() == [[1.0], [2.0], [3.0]] def test_read_rerun_describe_includes_robotics_metadata(tmp_path: Path) -> None: @@ -516,8 +574,8 @@ def test_read_rerun_robotics_mode_respects_explicit_selections( ).take(1)[0], ) - assert row["frames"].column("action").to_pylist() == [[30.0, 10.0], [40.0, 20.0]] - assert row["frames"].column("observation.state").to_pylist() == [ + assert row.actions.to_pylist() == [[30.0, 10.0], [40.0, 20.0]] + assert row.states.to_pylist() == [ [3.0, 1.0], [4.0, 2.0], ] @@ -579,8 +637,8 @@ def test_read_rerun_robotics_mode_with_explicit_timeline_uses_table_metadata( ).take(1)[0], ) - assert row["frames"].column("action").to_pylist() == [[30.0, 10.0], [40.0, 20.0]] - assert row["frames"].column("observation.state").to_pylist() == [ + assert row.actions.to_pylist() == [[30.0, 10.0], [40.0, 20.0]] + assert row.states.to_pylist() == [ [3.0, 1.0], [4.0, 2.0], ] From 8e59397be64193e974c01bb2b37c473820b749dd Mon Sep 17 00:00:00 2001 From: guipenedo Date: Wed, 17 Jun 2026 23:40:13 +0200 Subject: [PATCH 3/4] Clean up Rerun robotics side data --- docs/reading-data/rerun.md | 8 +++++--- src/refiner/pipeline/_rerun_io.py | 3 +++ src/refiner/pipeline/pipeline.py | 10 ++++++---- src/refiner/pipeline/sources/readers/rerun.py | 8 +++++--- src/refiner/robotics/tabular.py | 3 +-- tests/readers/test_rerun_reader.py | 12 ++++++++++++ 6 files changed, 32 insertions(+), 12 deletions(-) diff --git a/docs/reading-data/rerun.md b/docs/reading-data/rerun.md index 9c1b8981..8c54a7d1 100644 --- a/docs/reading-data/rerun.md +++ b/docs/reading-data/rerun.md @@ -73,9 +73,11 @@ the top-level `action` vector, scalar components under `/observation/state/**` into top-level `observation.state`, and encoded images under `/cam/**` into top-level video sources such as `cam.top`. -Rows also include the original `rerun` recording sidecar by default. Set -`include_recording=False` only when you want a lightweight robotics projection -and do not need the source Rerun structure later. +Rows also include a `rerun` recording sidecar by default. If `contents` is +omitted, that sidecar is built from the full recording view for the primary +timeline, not just the robotics prefixes. Set `include_recording=False` only +when you want a lightweight robotics projection and do not need the source +Rerun structure later. Use explicit selections when vector order or camera names matter: diff --git a/src/refiner/pipeline/_rerun_io.py b/src/refiner/pipeline/_rerun_io.py index 5388c364..8d64a8ea 100644 --- a/src/refiner/pipeline/_rerun_io.py +++ b/src/refiner/pipeline/_rerun_io.py @@ -5,6 +5,7 @@ import os import tempfile from pathlib import Path +from typing import ClassVar from typing import cast from refiner.io import DataFile @@ -64,6 +65,8 @@ def __setstate__(self, state: dict[str, object]) -> None: class RerunRecording: """Columnar Rerun recording data loaded from one RRD segment.""" + __refiner_side_data__: ClassVar[bool] = True + segment_id: str source_path: str tables: Mapping[str, Tabular] diff --git a/src/refiner/pipeline/pipeline.py b/src/refiner/pipeline/pipeline.py index 87b93435..1e37837f 100644 --- a/src/refiner/pipeline/pipeline.py +++ b/src/refiner/pipeline/pipeline.py @@ -1265,10 +1265,12 @@ def read_rerun( ``materialize_tables=False`` when you only need source chunk metadata. With ``output="robotics"``, the reader emits robotics episode rows directly, with top-level ``action``, ``observation.state``, and video fields derived - from configurable Rerun entity prefixes. The original ``rerun`` sidecar is - included by default; pass ``include_recording=False`` for a lightweight - projection. Pass ``actions``, ``states``, or ``videos`` to pin exact entity - paths and output order instead of using prefix-derived defaults. + from configurable Rerun entity prefixes. A ``rerun`` recording sidecar for + the primary timeline is included by default; when ``contents`` is omitted, + that sidecar uses the full recording view rather than only the robotics + prefixes. Pass ``include_recording=False`` for a lightweight projection. + Pass ``actions``, ``states``, or ``videos`` to pin exact entity paths and + output order instead of using prefix-derived defaults. """ return RefinerPipeline( source=RerunReader( diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index 85bf7f08..8d0d1050 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -523,8 +523,10 @@ def _robotics_row( timelines: Sequence[str], ) -> Row: timeline = self._primary_timeline(timelines) - contents = self._robotics_contents() - content_view = view.filter_contents(contents) + if self.include_recording: + content_view = self._view_for_contents(view) + else: + content_view = view.filter_contents(self._robotics_contents()) table = _collect_table( content_view.reader( index=timeline, @@ -551,7 +553,7 @@ def _robotics_row( ), application_id=application_id, recording_id=recording_id, - contents=tuple(contents), + contents=self.contents, timelines=(timeline,), include_static=self.include_static, use_source_chunks=False, diff --git a/src/refiner/robotics/tabular.py b/src/refiner/robotics/tabular.py index d4ee1b75..2ffd4f2c 100644 --- a/src/refiner/robotics/tabular.py +++ b/src/refiner/robotics/tabular.py @@ -2,7 +2,6 @@ from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass -from dataclasses import is_dataclass from typing import Any, cast import pyarrow as pa @@ -167,7 +166,7 @@ def _row_index_key(rows: Sequence[Row]) -> str: def _is_side_data_value(value: Any) -> bool: return value is not _MISSING and ( isinstance(value, Tabular | VideoSource | pa.Array | pa.ChunkedArray) - or (not isinstance(value, type) and is_dataclass(value)) + or bool(getattr(value, "__refiner_side_data__", False)) ) diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index c9aa6092..1b486431 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -40,6 +40,16 @@ def _tiny_rrd(path: Path) -> None: ) +def _rerun_entity_paths(table: Any) -> set[str]: + paths: set[str] = set() + for field in table.schema: + metadata = field.metadata or {} + value = metadata.get(b"rerun:entity_path") + if value is not None: + paths.add(value.decode("utf-8")) + return paths + + def _sparse_rrd(path: Path) -> None: import rerun as rr @@ -498,6 +508,8 @@ def test_read_rerun_robotics_mode_with_recording_includes_recording_payload( assert recording.recording_id == "episode-a" assert list(recording.tables) == ["frame"] assert recording.tables["frame"].num_rows == 3 + assert "/action_extra/y" in _rerun_entity_paths(recording.tables["frame"].table) + assert recording.contents is None assert "frames" not in row assert row.actions.to_pylist() == [[1.0], [2.0], [3.0]] From 210dcbeeacfc81cf0fb191e7258ef5c02ac93ead Mon Sep 17 00:00:00 2001 From: guipenedo Date: Wed, 17 Jun 2026 23:47:29 +0200 Subject: [PATCH 4/4] Simplify Rerun robotics rows --- docs/reading-data/rerun.md | 5 +- src/refiner/pipeline/pipeline.py | 7 +- src/refiner/pipeline/sources/readers/rerun.py | 181 ++++++++++-------- tests/readers/test_rerun_reader.py | 34 ++-- 4 files changed, 125 insertions(+), 102 deletions(-) diff --git a/docs/reading-data/rerun.md b/docs/reading-data/rerun.md index 8c54a7d1..e9a8d20a 100644 --- a/docs/reading-data/rerun.md +++ b/docs/reading-data/rerun.md @@ -75,9 +75,8 @@ under `/cam/**` into top-level video sources such as `cam.top`. Rows also include a `rerun` recording sidecar by default. If `contents` is omitted, that sidecar is built from the full recording view for the primary -timeline, not just the robotics prefixes. Set `include_recording=False` only -when you want a lightweight robotics projection and do not need the source -Rerun structure later. +timeline, not just the robotics prefixes. Project it away with `.drop("rerun")` +when downstream stages do not need the source Rerun structure. Use explicit selections when vector order or camera names matter: diff --git a/src/refiner/pipeline/pipeline.py b/src/refiner/pipeline/pipeline.py index 1e37837f..0e8e4e16 100644 --- a/src/refiner/pipeline/pipeline.py +++ b/src/refiner/pipeline/pipeline.py @@ -1246,7 +1246,6 @@ def read_rerun( primary_timeline: str | None = None, include_static: bool = True, materialize_tables: bool = True, - include_recording: bool | None = None, fill_latest_at: bool = False, action_prefix: str = DEFAULT_RERUN_ACTION_PREFIX, state_prefix: str = DEFAULT_RERUN_STATE_PREFIX, @@ -1268,9 +1267,8 @@ def read_rerun( from configurable Rerun entity prefixes. A ``rerun`` recording sidecar for the primary timeline is included by default; when ``contents`` is omitted, that sidecar uses the full recording view rather than only the robotics - prefixes. Pass ``include_recording=False`` for a lightweight projection. - Pass ``actions``, ``states``, or ``videos`` to pin exact entity paths and - output order instead of using prefix-derived defaults. + prefixes. Pass ``actions``, ``states``, or ``videos`` to pin exact entity + paths and output order instead of using prefix-derived defaults. """ return RefinerPipeline( source=RerunReader( @@ -1287,7 +1285,6 @@ def read_rerun( primary_timeline=primary_timeline, include_static=include_static, materialize_tables=materialize_tables, - include_recording=include_recording, fill_latest_at=fill_latest_at, action_prefix=action_prefix, state_prefix=state_prefix, diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index 8d0d1050..ed2561a6 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -113,7 +113,6 @@ def __init__( primary_timeline: str | None = None, include_static: bool = True, materialize_tables: bool = True, - include_recording: bool | None = None, fill_latest_at: bool = False, action_prefix: str = DEFAULT_RERUN_ACTION_PREFIX, state_prefix: str = DEFAULT_RERUN_STATE_PREFIX, @@ -130,10 +129,6 @@ def __init__( raise ValueError( "materialize_tables=False is only supported for recording output" ) - if output == "recording" and include_recording is False: - raise ValueError( - "include_recording=False is only supported for robotics output" - ) if output == "recording": _reject_recording_robotics_options( primary_timeline=primary_timeline, @@ -168,9 +163,6 @@ def __init__( self.include_static = include_static self.materialize_tables = materialize_tables self.use_source_chunks = self.timelines is None - self.include_recording = ( - True if include_recording is None else include_recording - ) self.fill_latest_at = fill_latest_at self.action_prefix = _normalize_entity_prefix(action_prefix) self.state_prefix = _normalize_entity_prefix(state_prefix) @@ -222,7 +214,6 @@ def describe(self) -> dict[str, Any]: "timelines": self.timelines, "include_static": self.include_static, "materialize_tables": self.materialize_tables, - "include_recording": self.include_recording, "fill_latest_at": self.fill_latest_at, } ) @@ -357,11 +348,7 @@ def _read_dataset( store_entries: Sequence[Any] | None = None, ) -> Iterator[SourceUnit]: if store_entries is None: - store_entries = ( - _recording_entries(local_path) - if self.output == "recording" or self.include_recording - else [] - ) + store_entries = _recording_entries(local_path) source_recording_count = len(store_entries) if store_entries else None entries_by_recording_id = {entry.recording_id: entry for entry in store_entries} timelines = self.timelines @@ -492,24 +479,6 @@ def _view_for_contents(self, dataset_or_view: Any) -> Any: return dataset_or_view return dataset_or_view.filter_contents(self.contents) - def _robotics_contents(self) -> tuple[str, ...]: - if self.contents is not None: - return self.contents - contents: list[str] = [] - if self.actions_explicit: - contents.extend(self.actions.values()) - else: - contents.append(_prefix_contents(self.action_prefix)) - if self.states_explicit: - contents.extend(self.states.values()) - else: - contents.append(_prefix_contents(self.state_prefix)) - if self.videos_explicit: - contents.extend(self.videos.values()) - else: - contents.append(_prefix_contents(self.camera_prefix)) - return tuple(dict.fromkeys(contents)) - def _robotics_row( self, view: Any, @@ -523,10 +492,7 @@ def _robotics_row( timelines: Sequence[str], ) -> Row: timeline = self._primary_timeline(timelines) - if self.include_recording: - content_view = self._view_for_contents(view) - else: - content_view = view.filter_contents(self._robotics_contents()) + content_view = self._view_for_contents(view) table = _collect_table( content_view.reader( index=timeline, @@ -536,28 +502,25 @@ def _robotics_row( row: dict[str, Any] = { "episode_id": segment_id, } - if self.include_recording: - static = ( - _collect_table(content_view.reader(index=None)) - if self.include_static - else None - ) - row["rerun"] = RerunRecording( - segment_id=segment_id, - source_path=source_path, - tables={timeline: Tabular(table)}, - static=Tabular(static) if static is not None else None, - source_file=source_file, - local_source=( - local_source if self._retain_batch_local_sources() else None - ), - application_id=application_id, - recording_id=recording_id, - contents=self.contents, - timelines=(timeline,), - include_static=self.include_static, - use_source_chunks=False, - ) + static = ( + _collect_table(content_view.reader(index=None)) + if self.include_static + else None + ) + row["rerun"] = RerunRecording( + segment_id=segment_id, + source_path=source_path, + tables={timeline: Tabular(table)}, + static=Tabular(static) if static is not None else None, + source_file=source_file, + local_source=local_source if self._retain_batch_local_sources() else None, + application_id=application_id, + recording_id=recording_id, + contents=self.contents, + timelines=(timeline,), + include_static=self.include_static, + use_source_chunks=False, + ) if self.fps is not None: row["fps"] = self.fps if self.robot_type is not None: @@ -583,14 +546,6 @@ def _robotics_row( if self.states_explicit else _prefixed_columns(scalar_columns, self.state_prefix) ) - if action_columns: - row["action"] = _list_column( - _singleton_scalar_matrix(table, action_columns) - ) - if state_columns: - row["observation.state"] = _list_column( - _singleton_scalar_matrix(table, state_columns) - ) image_columns = component_columns.get("EncodedImage:blob", {}) camera_columns = ( @@ -605,11 +560,6 @@ def _robotics_row( for name, column in camera_columns.items(): values = table.column(column).combine_chunks() _require_dense_encoded_images(values, video_name=name) - row[name] = VideoFrameSequence( - lambda values=values: _iter_encoded_images(values), - fps=self.fps or 30.0, - frame_count=len(values), - ) if self.file_path_column is not None: row[self.file_path_column] = source_path @@ -620,7 +570,16 @@ def _robotics_row( action_key="action", state_key="observation.state", video_keys={name: name for name in camera_columns}, - )(DictRow(row)) + )( + _RerunRoboticsSourceRow( + row, + timeline=timeline, + action_columns=action_columns, + state_columns=state_columns, + camera_columns=camera_columns, + fps=self.fps or 30.0, + ) + ) def _retain_batch_local_sources(self) -> bool: return ( @@ -630,6 +589,80 @@ def _retain_batch_local_sources(self) -> bool: ) +class _RerunRoboticsSourceRow(Row): + def __init__( + self, + data: Mapping[str, Any], + *, + timeline: str, + action_columns: Sequence[tuple[str, str]], + state_columns: Sequence[tuple[str, str]], + camera_columns: Mapping[str, str], + fps: float, + ) -> None: + self._data = dict(data) + self._timeline = timeline + self._action_columns = tuple(action_columns) + self._state_columns = tuple(state_columns) + self._camera_columns = dict(camera_columns) + self._fps = float(fps) + self._cache: dict[str, Any] = {} + + def __getitem__(self, key: str) -> Any: + if key == "action" and self._action_columns: + return self._cached( + key, + lambda: _list_column( + _singleton_scalar_matrix(self._table(), self._action_columns) + ), + ) + if key == "observation.state" and self._state_columns: + return self._cached( + key, + lambda: _list_column( + _singleton_scalar_matrix(self._table(), self._state_columns) + ), + ) + if key in self._camera_columns: + return self._cached( + key, lambda: self._video(key, self._camera_columns[key]) + ) + return self._data[key] + + def __iter__(self) -> Iterator[str]: + yield from self._data + if self._action_columns and "action" not in self._data: + yield "action" + if self._state_columns and "observation.state" not in self._data: + yield "observation.state" + for name in self._camera_columns: + if name not in self._data: + yield name + + def __len__(self) -> int: + return sum(1 for _ in self) + + def _recording(self) -> RerunRecording: + return cast(RerunRecording, self._data["rerun"]) + + def _table(self) -> pa.Table: + return self._recording().tables[self._timeline].table + + def _cached(self, key: str, factory: Any) -> Any: + if key not in self._cache: + self._cache[key] = factory() + return self._cache[key] + + def _video(self, name: str, column: str) -> VideoFrameSequence: + values = self._table().column(column).combine_chunks() + _require_dense_encoded_images(values, video_name=name) + return VideoFrameSequence( + lambda values=values: _iter_encoded_images(values), + fps=self._fps, + frame_count=len(values), + ) + + def _contents(contents: str | Sequence[str] | None) -> tuple[str, ...] | None: if contents is None: return None @@ -646,10 +679,6 @@ def _normalize_entity_prefix(value: str) -> str: return "/" if not stripped else "/" + stripped -def _prefix_contents(prefix: str) -> str: - return "/**" if prefix == "/" else f"{prefix}/**" - - def _selection_map( value: PathSelection | None, *, diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index 1b486431..0a8257ac 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -11,6 +11,7 @@ from refiner.pipeline import Row from refiner.pipeline.data.row import DictRow from refiner.pipeline.expressions import col +from refiner.pipeline.sources.readers import rerun as rerun_reader_module from refiner.pipeline.sources.readers.rerun import RerunReader from refiner.robotics.row import RoboticsRow @@ -306,11 +307,6 @@ def test_read_rerun_rejects_ignored_mode_options(tmp_path: Path) -> None: ): mdr.read_rerun(str(rrd), output="robotics", materialize_tables=False) - with pytest.raises( - ValueError, - match="include_recording=False is only supported for robotics output", - ): - mdr.read_rerun(str(rrd), output="recording", include_recording=False) with pytest.raises( ValueError, match="Rerun recording output does not use robotics options: primary_timeline", @@ -410,33 +406,35 @@ def test_read_rerun_robotics_mode_converts_to_robot_row(tmp_path: Path) -> None: ] -def test_read_rerun_robotics_mode_without_recording_skips_recording_entries( +def test_read_rerun_robotics_mode_derives_vectors_lazily_from_recording( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: rrd = tmp_path / "tiny.rrd" _tiny_rrd(rrd) + calls = 0 + real_matrix = rerun_reader_module._singleton_scalar_matrix + + def wrapped_matrix(*args: Any, **kwargs: Any) -> Any: + nonlocal calls + calls += 1 + return real_matrix(*args, **kwargs) monkeypatch.setattr( - "refiner.pipeline.sources.readers.rerun._recording_entries", - lambda *args, **kwargs: pytest.fail( - "robotics rows without recording payload do not need store metadata" - ), + rerun_reader_module, + "_singleton_scalar_matrix", + wrapped_matrix, ) row = cast( Any, - mdr.read_rerun( - str(rrd), - output="robotics", - include_recording=False, - fps=30.0, - ).take(1)[0], + mdr.read_rerun(str(rrd), output="robotics", fps=30.0).take(1)[0], ) - assert "rerun" not in row - assert "frames" not in row + assert calls == 0 assert row.actions.to_pylist() == [[1.0], [2.0], [3.0]] + assert row.actions.to_pylist() == [[1.0], [2.0], [3.0]] + assert calls == 1 def test_read_rerun_robotics_mode_exposes_top_level_fields_to_primitives(