From 137fb165aaaef11b75519427e8ce6d5a834667c6 Mon Sep 17 00:00:00 2001 From: Yann Date: Mon, 12 Oct 2026 16:38:53 +0800 Subject: [PATCH 1/2] feat(python): decode video frames in contiguous windows Preserve video frame descriptors through row selection, coalesce frame decoding across a batch, and expose backend selection for training windows. Co-Authored-By: Codex AI-Model: gpt-6-astra Co-Authored-By: Codex Co-Authored-By: Codex AI-Contributed/Feature: 181/181 AI-Contributed/UT: 103/103 --- docs/docs/pypaimon/multimodal-reading.md | 25 +++- docs/docs/pypaimon/pytorch.md | 10 +- docs/docs/pypaimon/video.md | 27 +++++ paimon-python/pypaimon/multimodal/query.py | 6 +- .../pypaimon/multimodal/window_dataset.py | 113 ++++++++++++++---- .../tests/contiguous_window_dataset_test.py | 103 +++++++++++++++- 6 files changed, 249 insertions(+), 35 deletions(-) diff --git a/docs/docs/pypaimon/multimodal-reading.md b/docs/docs/pypaimon/multimodal-reading.md index e4b02f70dc68..c53b9951a911 100644 --- a/docs/docs/pypaimon/multimodal-reading.md +++ b/docs/docs/pypaimon/multimodal-reading.md @@ -315,9 +315,28 @@ window. These calls preserve their single `is_pad` mask and previous shapes. `anchor_columns` and `tail` require `window_size`; do not combine `tail` and `boundary`. New applications should use field offsets and per-column masks. -Columns configured by `video-frame-field` are rejected: a window read would drop -the `frame_index` and other metadata carried by their `VideoFrameDescriptor` -values. Read those columns with `to_torch()` instead. +Columns configured by `video-frame-field` keep their descriptors until all rows +for an item or batch have been selected. The Dataset groups descriptors by +physical video, decodes each requested frame once, and restores the window +order. Install `pypaimon[torch,video]` for PyAV, or install TorchCodec +separately. Use `video_backend="pyav"` or `"torchcodec"` to require one +backend, or omit it to select an available one: + +```python +video_windows = frames.scan().to_contiguous_window_dataset( + columns=["state", "video", "action"], + frame_offsets={"video": [-1, 0], "action": list(range(16))}, + group_key="episode_index", + order_key="frame_index", + boundary="pad", + column_transforms={"video": torch.stack}, + video_backend="pyav", +) +``` + +Without a video column transform, each window contains a list of `CHW` tensors. +The `torch.stack` transform above produces `TCHW`; DataLoader adds the batch +dimension. Decoder sessions and range caches are process-local and bounded. ### Distributed BLOB processing with Ray diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md index 320b824c85d8..3915b5ab2254 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/docs/docs/pypaimon/pytorch.md @@ -268,9 +268,13 @@ the sampler's `set_epoch(epoch)` each epoch. Its padding or drop behavior still applies when the sample count is not divisible by the number of ranks. Dataset version pinning does not save sampler progress or random augmentation state. -Columns configured by `video-frame-field` are rejected: a window read would drop -the `frame_index` and other metadata carried by their `VideoFrameDescriptor` -values. Read those columns with `to_torch()` instead. +Columns configured by `video-frame-field` are decoded after the requested rows +have been coalesced. Their descriptor metadata remains intact, overlapping +windows reuse the same decoded rows, and `column_transforms={"video": +torch.stack}` converts the raw list of `CHW` frames to `TCHW`. Install +`pypaimon[torch,video]` for PyAV, or install TorchCodec separately. Select the +backend with `video_backend="pyav"` or `"torchcodec"`; omitting the option +selects an available backend. ## File Format Metadata Cache Reusable PyArrow Dataset metadata is cached across reads. Configure its estimated diff --git a/docs/docs/pypaimon/video.md b/docs/docs/pypaimon/video.md index 01e3521e809d..30b347755014 100644 --- a/docs/docs/pypaimon/video.md +++ b/docs/docs/pypaimon/video.md @@ -269,3 +269,30 @@ Define the decoder class and function at module scope when DataLoader uses the `shuffle=False`; row-buffer shuffle can turn monotonic frame access into seeks and require more open decoder sessions. Paimon's worker sharding is split-level, so it does not duplicate a split across workers. + +### Read video training windows + +`ContiguousWindowDataset` decodes `video-frame-field` columns directly. It +keeps descriptors through row selection, groups overlapping requests by +physical video, and decodes each selected row once before assembling windows: + +```python +video_windows = frames.scan().to_contiguous_window_dataset( + columns=["state", "video", "action"], + frame_offsets={"video": [-2, -1, 0], "action": list(range(16))}, + group_key="episode_id", + order_key="frame_index", + boundary="pad", + column_transforms={"video": torch.stack}, + video_backend="pyav", +) + +loader = DataLoader(video_windows, batch_size=32, num_workers=4) +for batch in loader: + # video is [N, T, C, H, W]. + train(batch["video"], batch["state"], batch["action"]) +``` + +Install `pypaimon[torch,video]` for the PyAV path, or install TorchCodec +separately. Omit `video_backend` to select an available decoder. The decoder +and range cache stay local to each DataLoader worker. diff --git a/paimon-python/pypaimon/multimodal/query.py b/paimon-python/pypaimon/multimodal/query.py index 195c05e9acf8..13aa615b8054 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/paimon-python/pypaimon/multimodal/query.py @@ -183,7 +183,8 @@ def to_contiguous_window_dataset( delta_timestamps=None, fps=None, tolerance_s=None, - boundary=None): + boundary=None, + video_backend=None): """Build snapshot-pinned training windows from this single-table scan. The Dataset indexes only ``group_key``, ``order_key``, and Paimon row @@ -220,6 +221,8 @@ def to_contiguous_window_dataset( defaulting to 1e-4. Only valid with delta_timestamps. boundary: Handling at both group ends: drop (default), pad, or error. Drop retains only anchors valid for every requested offset. + video_backend: Decoder for video-frame-field columns: ``pyav``, + ``torchcodec``, or ``None`` to select the available backend. Returns: A snapshot-pinned ``ContiguousWindowDataset``. See that class for @@ -244,6 +247,7 @@ def to_contiguous_window_dataset( fps=fps, tolerance_s=tolerance_s, boundary=boundary, + video_backend=video_backend, ) def to_ray( diff --git a/paimon-python/pypaimon/multimodal/window_dataset.py b/paimon-python/pypaimon/multimodal/window_dataset.py index 7b158eba8354..7d6c5f647507 100644 --- a/paimon-python/pypaimon/multimodal/window_dataset.py +++ b/paimon-python/pypaimon/multimodal/window_dataset.py @@ -20,6 +20,7 @@ import copy import math import operator +from functools import partial import numpy as np import pyarrow as pa @@ -30,6 +31,7 @@ from pypaimon.common.options.core_options import CoreOptions from pypaimon.multimodal.blob_read import fetch_blob_bodies from pypaimon.multimodal.query import ScanQuery, _PreFilterQuery +from pypaimon.multimodal.video import VideoFrameCollator from pypaimon.read.datasource.torch_dataset import ( SplitRangeIndex, row_ranges_for_split, @@ -76,8 +78,8 @@ class ContiguousWindowDataset(Dataset): every column uses ``[0]``. ``stride`` schedules anchors from each group's first row; boundary handling never shifts that schedule. ``blob_parallelism`` controls concurrent BLOB reads for each item or batch. - Video frame columns are not supported yet, because a window read would drop - the frame metadata carried by their descriptors. + Video frame columns retain their descriptors until the requested rows are + grouped and decoded, so overlapping windows reuse each physical frame read. """ _BOUNDARY_POLICIES = ("drop", "pad", "error") @@ -101,7 +103,8 @@ def __init__( delta_timestamps=None, fps=None, tolerance_s=None, - boundary=None): + boundary=None, + video_backend=None): if not isinstance(query, ScanQuery) or isinstance(query, _PreFilterQuery): raise TypeError( "ContiguousWindowDataset is only supported on scan(), " @@ -135,7 +138,14 @@ def __init__( raise ValueError("group_key and order_key must not be is_pad.") self.columns = _columns( query, columns, self.group_key, self.order_key) - _reject_video_columns(query._table, self.columns) + self._video_columns = [ + name for name in self.columns + if name in query._table.options.video_frame_fields() + ] + if video_backend not in (None, "torchcodec", "pyav"): + raise ValueError( + "video_backend must be None, 'torchcodec', or 'pyav'.") + self.video_backend = video_backend self.anchor_columns = _anchor_columns(anchor_columns, self.columns) self.frame_offsets = _frame_offsets( self.columns, frame_offsets, delta_timestamps, fps, tolerance_s, @@ -170,6 +180,19 @@ def __init__( self.snapshot_id = snapshot_id self._table = _pin_table(query._table, snapshot_id) self._plans = {} + self._video_collators = { + name: VideoFrameCollator( + self._table, + video_column=name, + decoder_factory=partial( + _open_default_video_decoder, backend=video_backend), + decode_batch_fn=_decode_video_frames, + output_column=name, + collate_fn=list, + range_parallelism=self.blob_parallelism, + ) + for name in self._video_columns + } self._build_index(index) @classmethod @@ -223,6 +246,24 @@ def __getstate__(self): state["_plans"] = {} return state + def close(self): + """Close video decoders opened by this Dataset process.""" + first_error = None + for collator in self._video_collators.values(): + try: + collator.close() + except Exception as error: + if first_error is None: + first_error = error + if first_error is not None: + raise first_error + + def __del__(self): + try: + self.close() + except Exception: + pass + def _resolve_window(self, index): index = operator.index(index) if index < 0: @@ -402,11 +443,16 @@ def _read_rows(self, row_ids, columns=None): Returns: A list of row dictionaries aligned one-for-one with ``row_ids``. - The internal ``_ROW_ID`` field is removed, and BLOB descriptors are - resolved to their bodies. + The internal ``_ROW_ID`` field is removed. Ordinary BLOB + descriptors are resolved to their bodies, while video descriptors + are decoded to frames. """ columns = self.columns if columns is None else columns rows = self._plan_for(columns).read(row_ids) + for name in columns: + collator = self._video_collators.get(name) + if collator is not None: + rows = collator(rows) row_id_column = SpecialFields.ROW_ID.name by_row_id = {} for row in rows: @@ -423,7 +469,12 @@ def _plan_for(self, columns): plan = self._plans.get(key) if plan is None: plan = _PinnedRowIdPlan( - self._table, columns, self.blob_parallelism, self.snapshot_id) + self._table, + columns, + self.blob_parallelism, + self.snapshot_id, + self._video_columns, + ) self._plans[key] = plan return plan @@ -436,7 +487,13 @@ class _PinnedRowIdPlan: and batch reads never revisit the snapshot's manifests. """ - def __init__(self, table, columns, blob_parallelism, snapshot_id): + def __init__( + self, + table, + columns, + blob_parallelism, + snapshot_id, + video_columns=()): self._blob_parallelism = blob_parallelism self._snapshot_id = snapshot_id self._blob_columns = [ @@ -444,13 +501,18 @@ def __init__(self, table, columns, blob_parallelism, snapshot_id): if field.name in columns and is_blob_file_type(field.type) ] + self._body_blob_columns = [ + name for name in self._blob_columns if name not in video_columns + ] self._map_blob_columns = { field.name for field in table.fields - if field.name in self._blob_columns and is_map_blob_type(field.type) + if field.name in self._body_blob_columns + and is_map_blob_type(field.type) } self._array_blob_columns = { field.name for field in table.fields - if field.name in self._blob_columns and is_array_blob_type(field.type) + if field.name in self._body_blob_columns + and is_array_blob_type(field.type) } blob_column_set = set(self._blob_columns) self._projection = ( @@ -495,21 +557,21 @@ def read(self, row_ids): arrow = arrow.filter(pc.is_in( row_id_column, value_set=pa.array(requested, type=row_id_column.type))) - if not self._blob_columns: + if not self._body_blob_columns: return arrow.to_pylist() bodies = fetch_blob_bodies( self._table.file_io, - arrow.select(self._blob_columns).to_pydict(), - self._blob_columns, + arrow.select(self._body_blob_columns).to_pydict(), + self._body_blob_columns, self._blob_parallelism, self._map_blob_columns, self._array_blob_columns) - blob_column_set = set(self._blob_columns) + blob_column_set = set(self._body_blob_columns) rows = arrow.select([ name for name in arrow.column_names if name not in blob_column_set ]).to_pylist() - for name in self._blob_columns: + for name in self._body_blob_columns: values = bodies[name] if len(values) != len(rows): raise RuntimeError( @@ -642,15 +704,18 @@ def _require_tag_snapshot(table, tag_name, snapshot_id): "built from snapshot %s." % (tag_name, resolved, snapshot_id)) -def _reject_video_columns(table, columns): - video_columns = [ - name for name in columns if name in table.options.video_frame_fields() - ] - if video_columns: - raise ValueError( - "ContiguousWindowDataset does not support video frame columns %s " - "yet: a window read would drop their frame metadata, such as " - "frame_index." % video_columns) +def _open_default_video_decoder(stream, backend=None): + # Keep the codec implementation shared with PaimonLeRobotDataset while the + # generic collator owns grouping, range reads, and process-local caching. + from pypaimon.multimodal.lerobot.dataset import _open_video_decoder + return _open_video_decoder(stream, backend=backend) + + +def _decode_video_frames(decoder, frame_indices, unused_rows): + get_frames_at = getattr(decoder, "get_frames_at", None) + if get_frames_at is not None: + return get_frames_at(indices=frame_indices).data + return [decoder[index] for index in frame_indices] def _frame_offsets(columns, offsets, timestamps, fps, tolerance_s, diff --git a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py index 0fd9d218cbce..994c4d2bd45b 100644 --- a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py +++ b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py @@ -21,6 +21,7 @@ import shutil import tempfile import unittest +from types import SimpleNamespace from unittest.mock import patch import pyarrow as pa @@ -892,7 +893,7 @@ def test_rejects_nan_group_keys_that_never_compare_equal(self): window_size=2, columns=["value"], group_key="episode", order_key="step") - def test_rejects_video_frame_columns_that_would_lose_frame_metadata(self): + def test_decodes_video_windows_and_reuses_overlapping_frames(self): table = self.conn.create_table( "videos", schema=pa.schema([ @@ -905,12 +906,100 @@ def test_rejects_video_frame_columns_that_would_lose_frame_metadata(self): "blob-as-descriptor": "true", }), ) + video_path = os.path.join(self.temp_dir, "episode-a.mp4") + with open(video_path, "wb") as output: + output.write(b"encoded-video") + table.add_video( + pmm.Blob.from_local(video_path), + [ + {"episode": "episode-a", "step": step} + for step in range(3) + ], + first_frame=7, + ) + decoded = [] - with self.assertRaisesRegex( - ValueError, "video frame columns.*frame_index"): - table.scan().to_contiguous_window_dataset( + class Decoder: + + def get_frames_at(self, *, indices): + decoded.append(tuple(indices)) + return SimpleNamespace(data=[torch.tensor(index) for index in indices]) + + with patch( + "pypaimon.multimodal.lerobot.dataset._open_video_decoder", + return_value=Decoder()): + dataset = table.scan().to_contiguous_window_dataset( window_size=2, columns=["video"], group_key="episode", order_key="step") + try: + samples = dataset.__getitems__([0, 1]) + finally: + close = getattr(dataset, "close", None) + if close is not None: + close() + + self.assertEqual([(7, 8, 9)], decoded) + self.assertEqual( + [[7, 8], [8, 9]], + [[int(frame) for frame in sample["video"]] for sample in samples], + ) + + def test_decodes_encoded_video_windows_with_multiple_workers(self): + try: + import av # noqa: F401 + from PIL import Image + except ImportError: + self.skipTest("PyAV and Pillow are required for video decoding") + + table = self.conn.create_table( + "encoded_videos", + schema=pa.schema([ + pa.field("episode", pa.string(), nullable=False), + pa.field("step", pa.int32(), nullable=False), + pa.field("video", pa.large_binary()), + ]), + options=dict(_TABLE_OPTIONS, **{ + "video-frame-field": "video", + "blob-as-descriptor": "true", + }), + ) + table.add_images_as_video( + [ + Image.new("RGB", (8, 6), color) + for color in ("red", "green", "blue") + ], + [ + {"episode": "episode-a", "step": step} + for step in range(3) + ], + fps=10, + codec="mpeg4", + pixel_format="yuv420p", + gop_size=1, + codec_options={"qscale": "3"}, + ) + dataset = table.scan().to_contiguous_window_dataset( + window_size=2, + columns=["video"], + group_key="episode", + order_key="step", + column_transforms={"video": torch.stack}, + video_backend="pyav", + ) + try: + batches = list(torch.utils.data.DataLoader( + dataset, batch_size=2, num_workers=2, shuffle=False)) + finally: + dataset.close() + + self.assertEqual(1, len(batches)) + self.assertEqual((2, 2, 3, 6, 8), tuple(batches[0]["video"].shape)) + first = batches[0]["video"][0, 0].float().mean(dim=(1, 2)) + last = batches[0]["video"][1, 1].float().mean(dim=(1, 2)) + self.assertGreater(first[0], first[1] + 100) + self.assertGreater(first[0], first[2] + 100) + self.assertGreater(last[2], last[0] + 100) + self.assertGreater(last[2], last[1] + 100) def test_validates_configuration_and_scan_only_contract(self): table = self._table() @@ -949,6 +1038,12 @@ def test_validates_configuration_and_scan_only_contract(self): window_size=2, columns=["value"], group_key="episode", order_key="step") + with self.assertRaisesRegex(ValueError, "video_backend must be"): + table.scan().to_contiguous_window_dataset( + window_size=2, columns=["value"], + group_key="episode", order_key="step", + video_backend="unsupported") + if __name__ == "__main__": unittest.main() From 4f8a68292e38ae61f642c294bdd2510542b2e358 Mon Sep 17 00:00:00 2001 From: Yann Date: Tue, 13 Oct 2026 10:36:13 +0800 Subject: [PATCH 2/2] fix(python): compact decoded video frame storage Clone frames split from a batched decoder tensor so later window copies retain only each frame's storage. Co-Authored-By: Codex AI-Model: gpt-6-astra Co-Authored-By: Codex AI-Contributed/Feature: 5/5 AI-Contributed/UT: 8/8 --- paimon-python/pypaimon/multimodal/window_dataset.py | 5 ++++- .../pypaimon/tests/contiguous_window_dataset_test.py | 8 +++++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/window_dataset.py b/paimon-python/pypaimon/multimodal/window_dataset.py index 7d6c5f647507..8e1dd443df6a 100644 --- a/paimon-python/pypaimon/multimodal/window_dataset.py +++ b/paimon-python/pypaimon/multimodal/window_dataset.py @@ -714,7 +714,10 @@ def _open_default_video_decoder(stream, backend=None): def _decode_video_frames(decoder, frame_indices, unused_rows): get_frames_at = getattr(decoder, "get_frames_at", None) if get_frames_at is not None: - return get_frames_at(indices=frame_indices).data + frames = get_frames_at(indices=frame_indices).data + # Iterating a batched tensor returns views backed by the whole batch. + # Give each row compact storage before window assembly deep-copies it. + return [frame.clone() for frame in frames] return [decoder[index] for index in frame_indices] diff --git a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py index 994c4d2bd45b..0b5cb96dd514 100644 --- a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py +++ b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py @@ -923,7 +923,7 @@ class Decoder: def get_frames_at(self, *, indices): decoded.append(tuple(indices)) - return SimpleNamespace(data=[torch.tensor(index) for index in indices]) + return SimpleNamespace(data=torch.tensor(indices)) with patch( "pypaimon.multimodal.lerobot.dataset._open_video_decoder", @@ -943,6 +943,12 @@ def get_frames_at(self, *, indices): [[7, 8], [8, 9]], [[int(frame) for frame in sample["video"]] for sample in samples], ) + for sample in samples: + for frame in sample["video"]: + self.assertEqual( + frame.numel() * frame.element_size(), + frame.untyped_storage().nbytes(), + ) def test_decodes_encoded_video_windows_with_multiple_workers(self): try: