From c57e8ea5d0f3d6fd560907ae15bb1f42eb762a9b Mon Sep 17 00:00:00 2001 From: QuakeWang Date: Tue, 15 Sep 2026 15:36:07 +0800 Subject: [PATCH 1/2] [python] Batch TorchCodec frame decoding for LeRobot datasets Decode each payload group with TorchCodec get_frames_at instead of repeated single-frame calls. Preserve sample order, bounded decoder reuse, and PyAV/custom callback behavior. Signed-off-by: QuakeWang --- .../pypaimon/multimodal/lerobot/dataset.py | 8 + paimon-python/pypaimon/multimodal/video.py | 30 +++- .../pypaimon/tests/multimodal_lerobot_test.py | 164 ++++++++++++++++++ .../pypaimon/tests/multimodal_video_test.py | 86 ++++++++- 4 files changed, 282 insertions(+), 6 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index b9220df4cbca..497f3e1a8a3c 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -268,6 +268,7 @@ def _set_frame_rows( decoder_factory=partial( _open_video_decoder, backend=self.video_backend), decode_fn=_decode_video_frame, + decode_batch_fn=_decode_video_frames, output_column=key, collate_fn=_identity, ) @@ -1541,6 +1542,13 @@ def _decode_video_frame(decoder, frame_index, unused_row): return decoder[frame_index] +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 _identity(values): return values diff --git a/paimon-python/pypaimon/multimodal/video.py b/paimon-python/pypaimon/multimodal/video.py index c027a91e41f1..e6e365d8878e 100644 --- a/paimon-python/pypaimon/multimodal/video.py +++ b/paimon-python/pypaimon/multimodal/video.py @@ -37,6 +37,10 @@ class VideoFrameCollator: order to avoid unnecessary decoder seeks. Rows are restored to their input order before ``collate_fn`` is called. + Optional ``decode_batch_fn`` replaces ``decode_fn`` for each payload group. + It receives the cached decoder, sorted frame indices, and corresponding row + dictionaries, and returns a sequence with one frame per row in that order. + The cache is process-local and keyed by physical video payload identity. ``collate_fn`` defaults to PyTorch's ``default_collate`` and may be replaced for decoders that already return batched objects. @@ -51,13 +55,16 @@ def __init__( decode_fn, output_column="frame", max_open_videos=8, - collate_fn=None): + collate_fn=None, + decode_batch_fn=None): if not video_column: raise ValueError("video_column is required.") if not callable(decoder_factory): raise ValueError("decoder_factory must be callable.") if not callable(decode_fn): raise ValueError("decode_fn must be callable.") + if decode_batch_fn is not None and not callable(decode_batch_fn): + raise ValueError("decode_batch_fn must be callable or None.") if ( isinstance(max_open_videos, bool) or not isinstance(max_open_videos, int) @@ -76,6 +83,7 @@ def __init__( self.video_column = video_column self.decoder_factory = decoder_factory self.decode_fn = decode_fn + self.decode_batch_fn = decode_batch_fn self.output_column = output_column self.max_open_videos = max_open_videos self.collate_fn = collate_fn @@ -126,11 +134,23 @@ def _decode_rows(self, rows): for payload, frames in grouped.items(): decoder = self._decoder(payload) - for frame_index, position, output in sorted( - frames, key=lambda frame: frame[0]): - output[self.output_column] = self.decode_fn( - decoder, frame_index, output + frames.sort(key=lambda frame: frame[0]) + if self.decode_batch_fn is not None: + values = self.decode_batch_fn( + decoder, + [index for index, _, _ in frames], + [output for _, _, output in frames], + ) + if len(values) != len(frames): + raise ValueError( + "decode_batch_fn must return one frame per row.") + else: + values = ( + self.decode_fn(decoder, index, output) + for index, _, output in frames ) + for (_, position, output), value in zip(frames, values): + output[self.output_column] = value decoded[position] = output return decoded diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index a9790b4e1607..dd8d02a99438 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -47,6 +47,7 @@ from pypaimon.multimodal.lerobot.dataset import ( _PyAVVideoDecoder, _arrow_rows, + _decode_video_frames, _image_tensor, _index_names, _open_video_decoder, @@ -331,6 +332,169 @@ def test_default_video_backend_falls_back_on_os_error(self): with self.assertRaises(OSError): _open_video_decoder(stream, backend="torchcodec") + def test_video_batches_include_delta_frames_and_preserve_backends(self): + try: + import torch + except ImportError as error: + self.skipTest(str(error)) + + metadata = { + "repo_id": "test/video-batches", + "info": { + "codebase_version": "v3.0", "fps": 10, + "total_frames": 3, "total_episodes": 1, "total_tasks": 1, + "features": { + "index": {"dtype": "int64", "shape": [1]}, + "episode_index": {"dtype": "int64", "shape": [1]}, + "frame_index": {"dtype": "int64", "shape": [1]}, + "timestamp": {"dtype": "float32", "shape": [1]}, + "task_index": {"dtype": "int64", "shape": [1]}, + "camera": {"dtype": "video", "shape": [2, 2, 3], + "names": ["height", "width", "channels"]}, + }, + }, + "episodes": [{"episode_index": 0, "dataset_from_index": 0, + "dataset_to_index": 3, "length": 3, + "tasks": ["pick"]}], + "tasks": ["pick"], + } + + class Reader(pmm.PaimonDatasetReader): + + def read_indices(self, indices, columns): + return pa.Table.from_pylist([{ + "index": index, "episode_index": 0, + "frame_index": index, "timestamp": index / 10, + "task_index": 0, + "camera": pmm.VideoFrameDescriptor( + "file:///shared.video", 0, 5, index + 3).serialize(), + } for index in indices], schema=self.schema).select(columns) + + for backend, batch in (("torchcodec", True), (None, True), + ("pyav", False), (None, False)): + with self.subTest(backend=backend, batch=batch): + calls = [] + + class Decoder: + def __getitem__(self, index): + calls.append(index) + return torch.full((3, 2, 2), index, dtype=torch.uint8) + + class BatchDecoder(Decoder): + def __getitem__(self, index): + raise AssertionError("unexpected single-frame call") + + def get_frames_at(self, *, indices): + calls.append(indices) + return SimpleNamespace(data=torch.stack([ + torch.full((3, 2, 2), i, dtype=torch.uint8) + for i in indices + ])) + + decoder = BatchDecoder() if batch else Decoder() + module = "pypaimon.multimodal.lerobot.dataset." + with patch( + module + "_open_torchcodec_decoder", + return_value=decoder, + side_effect=None if batch else OSError("unavailable"), + ) as open_torchcodec, patch( + module + "_PyAVVideoDecoder", return_value=decoder, + ) as open_pyav: + file_io = SimpleNamespace( + new_input_stream=lambda path: io.BytesIO(b"video")) + reader = Reader( + metadata, file_io=file_io, video_backend=backend, + delta_timestamps={"camera": [-0.1, 0.0, 0.1]}, + ) + try: + last, first, duplicate = reader.get_items([2, 0, 2]) + self.assertEqual([[3, 4, 5]] if batch else [3, 4, 5], calls) + self.assertEqual([3, 3, 2, 2], list(last["camera"].shape)) + torch.testing.assert_close( + last["camera"][:, 0, 0, 0], + torch.tensor([4, 5, 5], dtype=torch.float32) / 255) + torch.testing.assert_close( + first["camera"][:, 0, 0, 0], + torch.tensor([3, 3, 4], dtype=torch.float32) / 255) + self.assertEqual([False, False, True], + last["camera_is_pad"].tolist()) + self.assertEqual([True, False, False], + first["camera_is_pad"].tolist()) + last["camera"].zero_() + self.assertGreater(float(duplicate["camera"].sum()), 0) + reader.get_items([1]) + self.assertEqual( + [[3, 4, 5]] * 2 if batch else [3, 4, 5] * 2, calls) + self.assertEqual(0 if backend == "pyav" else 1, + open_torchcodec.call_count) + self.assertEqual(0 if batch else 1, open_pyav.call_count) + finally: + reader.close() + + def test_video_batch_decode_propagates_errors(self): + decoder = Mock() + decoder.get_frames_at.side_effect = RuntimeError("decode failed") + with self.assertRaisesRegex(RuntimeError, "decode failed"): + _decode_video_frames(decoder, [0, 1], [{}, {}]) + decoder.get_frames_at.assert_called_once_with(indices=[0, 1]) + + @unittest.skipUnless( + av is not None and importlib.util.find_spec("torchcodec") is not None, + "PyAV and TorchCodec are required for batch video decoding", + ) + def test_torchcodec_batch_matches_single_frame_decoding(self): + import torch + try: + from torchcodec.decoders import VideoDecoder + except (ImportError, OSError, RuntimeError) as error: + self.skipTest(str(error)) + + output = io.BytesIO() + with av.open(output, mode="w", format="mp4") as container: + stream = container.add_stream("libx264", rate=10) + stream.width = stream.height = 16 + stream.pix_fmt = "yuv420p" + stream.gop_size = 4 + stream.codec_context.max_b_frames = 2 + for index in range(12): + frame = av.VideoFrame.from_ndarray( + np.full((16, 16, 3), index * 16, dtype=np.uint8), + format="rgb24") + for packet in stream.encode(frame): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + + payload = output.getvalue() + decoder = _open_video_decoder(io.BytesIO(payload), backend="torchcodec") + self.assertIsInstance(decoder, VideoDecoder) + indices = [9, 1, 9, 5] + expected = torch.stack([decoder[index] for index in indices]) + file_io = SimpleNamespace( + new_input_stream=lambda path: io.BytesIO(payload)) + collator = pmm.VideoFrameCollator( + SimpleNamespace(file_io=file_io), video_column="video", + decoder_factory=lambda source: decoder, + decode_fn=lambda *args: self.fail("unexpected single-frame call"), + decode_batch_fn=_decode_video_frames, + collate_fn=lambda rows: rows, + ) + try: + with patch.object( + decoder, "get_frames_at", wraps=decoder.get_frames_at, + ) as decode_batch: + result = collator([{ + "video": pmm.VideoFrameDescriptor( + "file:///episode.mp4", 0, len(payload), index).serialize(), + } for index in indices]) + decode_batch.assert_called_once_with(indices=[1, 5, 9, 9]) + actual = torch.stack([row["frame"] for row in result]) + self.assertEqual(torch.uint8, actual.dtype) + self.assertEqual((4, 3, 16, 16), tuple(actual.shape)) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + finally: + collator.close() + def test_dataset_requires_supported_python(self): with patch( "pypaimon.multimodal.lerobot.dataset.sys.version_info", diff --git a/paimon-python/pypaimon/tests/multimodal_video_test.py b/paimon-python/pypaimon/tests/multimodal_video_test.py index 8aaedfccc548..27d25a4d4603 100644 --- a/paimon-python/pypaimon/tests/multimodal_video_test.py +++ b/paimon-python/pypaimon/tests/multimodal_video_test.py @@ -71,7 +71,10 @@ def test_reuses_decoder_for_rows_with_same_descriptor(self): def factory(stream): factory_calls.append("open") - return _Decoder(stream, factory_calls) + decoder = _Decoder(stream, factory_calls) + decoder.get_frames_at = lambda indices: self.fail( + "must preserve the custom single-frame callback") + return decoder collator = VideoFrameCollator( self.table, @@ -173,6 +176,87 @@ def __call__(self, rows): self.assertEqual("decoded-base", row_groups[0][4]["video"]) self.assertEqual("decoded-delta", row_groups[1][1]["video"]) + def test_batch_decode_groups_payload_ranges_and_restores_order(self): + path = os.path.join(self.temp_dir.name, "shared.video") + with open(path, "wb") as output: + output.write(b"firstsecond") + calls = [] + created = [] + + def factory(stream): + decoder = _Decoder(stream, []) + created.append(decoder) + return decoder + + def decode_batch(decoder, indices, rows): + calls.append((decoder.decode(0)[0], indices)) + return [ + (row["request"], decoder.decode(index)) + for index, row in zip(indices, rows) + ] + + def row(request, offset, length, index): + return {"request": request, "video": VideoFrameDescriptor( + path, offset, length, index).serialize()} + + collator = VideoFrameCollator( + self.table, + video_column="video", + decoder_factory=factory, + decode_fn=lambda *args: self.fail("unexpected single decode"), + decode_batch_fn=decode_batch, + max_open_videos=2, + collate_fn=lambda rows: rows, + ) + rows = [ + row("a", 0, 5, 3), row("b", 5, 6, 2), + row("c", 0, 5, 1), row("d", 0, 5, 3), + {"request": "e", "video": None}, row("f", 0, 11, 0), + ] + try: + result = collator(rows) + self.assertEqual( + [(b"first", [1, 3, 3]), (b"second", [2]), + (b"firstsecond", [0])], calls) + self.assertEqual([r["request"] for r in rows], + [r["request"] for r in result]) + self.assertEqual( + [("a", (b"first", 3)), ("b", (b"second", 2)), + ("c", (b"first", 1)), ("d", (b"first", 3)), + None, ("f", (b"firstsecond", 0))], + [r["frame"] for r in result]) + self.assertTrue(all("frame" not in r for r in rows)) + self.assertEqual(3, len(created)) + self.assertTrue(created[0].closed) + self.assertEqual(2, len(collator._decoders)) + self.assertEqual( + ("g", (b"second", 1)), + collator(row("g", 5, 6, 1))["frame"]) + self.assertEqual(3, len(created)) + self.assertEqual([], collator([])) + self.assertEqual(4, len(calls)) + finally: + collator.close() + self.assertTrue(all(decoder.closed for decoder in created)) + + def test_batch_decode_rejects_wrong_result_count(self): + descriptor = self._descriptor("episode.mp4", b"video", 0) + for count in (0, 2): + with self.subTest(count=count): + collator = VideoFrameCollator( + self.table, + video_column="video", + decoder_factory=lambda stream: _Decoder(stream, []), + decode_fn=lambda *args: self.fail("unexpected fallback"), + decode_batch_fn=lambda *args: [None] * count, + collate_fn=lambda rows: rows, + ) + try: + with self.assertRaisesRegex(ValueError, "one frame per row"): + collator([{"video": descriptor}]) + finally: + collator.close() + def test_evicts_least_recently_used_decoder(self): descriptors = [ self._descriptor("episode-%d.mp4" % index, bytes([index]), index) From 31bc385b54a19c2dfe37185dc15f11f4bc57f233 Mon Sep 17 00:00:00 2001 From: QuakeWang Date: Tue, 15 Sep 2026 16:54:50 +0800 Subject: [PATCH 2/2] [python] Allow batch-only video decode callbacks Make decode_fn optional when decode_batch_fn is provided, while validating that at least one callable is supplied. Remove the redundant LeRobot single-frame callback and cover missing and invalid callbacks. Signed-off-by: QuakeWang --- .../pypaimon/multimodal/lerobot/dataset.py | 5 -- paimon-python/pypaimon/multimodal/video.py | 10 ++-- .../pypaimon/tests/multimodal_lerobot_test.py | 1 - .../pypaimon/tests/multimodal_video_test.py | 46 +++++++++++++++++++ 4 files changed, 53 insertions(+), 9 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 497f3e1a8a3c..792e246f44e2 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -267,7 +267,6 @@ def _set_frame_rows( video_column=key, decoder_factory=partial( _open_video_decoder, backend=self.video_backend), - decode_fn=_decode_video_frame, decode_batch_fn=_decode_video_frames, output_column=key, collate_fn=_identity, @@ -1538,10 +1537,6 @@ def close(self): self._container.close() -def _decode_video_frame(decoder, frame_index, unused_row): - return decoder[frame_index] - - 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: diff --git a/paimon-python/pypaimon/multimodal/video.py b/paimon-python/pypaimon/multimodal/video.py index e6e365d8878e..ac2e7039e83c 100644 --- a/paimon-python/pypaimon/multimodal/video.py +++ b/paimon-python/pypaimon/multimodal/video.py @@ -40,6 +40,7 @@ class VideoFrameCollator: Optional ``decode_batch_fn`` replaces ``decode_fn`` for each payload group. It receives the cached decoder, sorted frame indices, and corresponding row dictionaries, and returns a sequence with one frame per row in that order. + At least one of ``decode_fn`` or ``decode_batch_fn`` must be provided. The cache is process-local and keyed by physical video payload identity. ``collate_fn`` defaults to PyTorch's ``default_collate`` and may be replaced @@ -52,7 +53,7 @@ def __init__( *, video_column, decoder_factory, - decode_fn, + decode_fn=None, output_column="frame", max_open_videos=8, collate_fn=None, @@ -61,10 +62,13 @@ def __init__( raise ValueError("video_column is required.") if not callable(decoder_factory): raise ValueError("decoder_factory must be callable.") - if not callable(decode_fn): - raise ValueError("decode_fn must be callable.") + if decode_fn is not None and not callable(decode_fn): + raise ValueError("decode_fn must be callable or None.") if decode_batch_fn is not None and not callable(decode_batch_fn): raise ValueError("decode_batch_fn must be callable or None.") + if decode_fn is None and decode_batch_fn is None: + raise ValueError( + "At least one of decode_fn or decode_batch_fn is required.") if ( isinstance(max_open_videos, bool) or not isinstance(max_open_videos, int) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index dd8d02a99438..70fb1020b55c 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -475,7 +475,6 @@ def test_torchcodec_batch_matches_single_frame_decoding(self): collator = pmm.VideoFrameCollator( SimpleNamespace(file_io=file_io), video_column="video", decoder_factory=lambda source: decoder, - decode_fn=lambda *args: self.fail("unexpected single-frame call"), decode_batch_fn=_decode_video_frames, collate_fn=lambda rows: rows, ) diff --git a/paimon-python/pypaimon/tests/multimodal_video_test.py b/paimon-python/pypaimon/tests/multimodal_video_test.py index 27d25a4d4603..a0a5f34861f0 100644 --- a/paimon-python/pypaimon/tests/multimodal_video_test.py +++ b/paimon-python/pypaimon/tests/multimodal_video_test.py @@ -62,6 +62,52 @@ def setUp(self): def tearDown(self): self.temp_dir.cleanup() + def test_accepts_batch_callback_without_single_frame_callback(self): + descriptors = [ + self._descriptor("episode.mp4", b"video", index) + for index in (2, 0) + ] + for kwargs in ({}, {"decode_fn": None}): + with self.subTest(kwargs=kwargs): + collator = VideoFrameCollator( + self.table, + video_column="video", + decoder_factory=lambda stream: _Decoder(stream, []), + decode_batch_fn=lambda decoder, indices, rows: [ + decoder.decode(index) for index in indices + ], + collate_fn=lambda rows: rows, + **kwargs, + ) + try: + result = collator([{"video": value} for value in descriptors]) + self.assertEqual([(b"video", 2), (b"video", 0)], + [row["frame"] for row in result]) + finally: + collator.close() + + def test_rejects_missing_or_invalid_decode_callbacks(self): + callback = lambda *args: None + cases = [ + ({}, "At least one"), + ({"decode_fn": None, "decode_batch_fn": None}, "At least one"), + ({"decode_fn": False}, "decode_fn must be callable"), + ({"decode_fn": False, "decode_batch_fn": callback}, + "decode_fn must be callable"), + ({"decode_batch_fn": False}, "decode_batch_fn must be callable"), + ({"decode_fn": callback, "decode_batch_fn": False}, + "decode_batch_fn must be callable"), + ] + for kwargs, message in cases: + with self.subTest(kwargs=kwargs): + with self.assertRaisesRegex(ValueError, message): + VideoFrameCollator( + self.table, + video_column="video", + decoder_factory=lambda stream: _Decoder(stream, []), + **kwargs, + ) + def test_reuses_decoder_for_rows_with_same_descriptor(self): descriptors = [ self._descriptor("episode-1.mp4", b"video-one", frame)