[python] Batch TorchCodec frame decoding for LeRobot datasets

paimon-python/pypaimon/multimodal/lerobot/dataset.py

@@ -269,7 +269,7 @@ 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,
             )
@@ -1539,8 +1539,11 @@ 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:
+        return get_frames_at(indices=frame_indices).data
+    return [decoder[index] for index in frame_indices]
 
 
 def _identity(values):

paimon-python/pypaimon/multimodal/video.py

@@ -37,6 +37,11 @@ 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.
+    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
     for decoders that already return batched objects.
@@ -48,16 +53,22 @@ def __init__(
             *,
             video_column,
             decoder_factory,
-            decode_fn,
+            decode_fn=None,
             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_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)
@@ -76,6 +87,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 +138,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
 

paimon-python/pypaimon/tests/multimodal_lerobot_test.py

@@ -47,6 +47,7 @@
 from pypaimon.multimodal.lerobot.dataset import (
     _PyAVVideoDecoder,
     _arrow_rows,
+    _decode_video_frames,
     _decode_video_rows,
     _image_tensor,
     _index_names,
@@ -387,6 +388,168 @@ 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_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",