[codex] Extract Rerun reader

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)

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) |

docs/reading-data/rerun.md

@@ -0,0 +1,111 @@
+---
+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 robotics episode rows directly:
+
+```python
+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 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 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. 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:
+
+```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 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
+
+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.

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]
 ```

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]",

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,

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
 
@@ -297,31 +298,53 @@ 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
-    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[Block, Sequence[VectorizedOp]]:
+        if row_projection_ops:
+            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):
+                if not row_remaining_ops:
+                    return projected, ()
+                raise
+        return _tabular_from_rows(batch, schema=input_schema), ops
+
     def _chunk_rows_for_budget() -> int:
         if (
             max_vectorized_block_bytes is None
@@ -332,23 +355,24 @@ 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:
             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
                 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:
@@ -368,7 +392,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
@@ -381,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:
@@ -420,7 +444,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 +481,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:

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,