[python] Group LeRobot video frame decoding by payload

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

@@ -1547,16 +1547,18 @@ def _identity(values):
 
 def _decode_video_rows(row_groups, collators):
     for collator in collators:
+        targets = []
+        input_rows = []
         for rows in row_groups:
-            indices = [
-                index for index, row in rows.items()
-                if collator.video_column in row
-            ]
-            if not indices:
-                continue
-            decoded = collator([rows[index] for index in indices])
-            for index, row in zip(indices, decoded):
-                rows[index] = row
+            for index, row in rows.items():
+                if collator.video_column in row:
+                    targets.append((rows, index))
+                    input_rows.append(row)
+        if not input_rows:
+            continue
+        decoded = collator(input_rows)
+        for (rows, index), row in zip(targets, decoded):
+            rows[index] = row
 
 
 def _normalize_index(index, size):

paimon-python/pypaimon/multimodal/video.py

@@ -33,6 +33,10 @@ class VideoFrameCollator:
     keeps Paimon independent of a particular video codec library while allowing
     PyAV, TorchCodec, or an application decoder to be plugged in.
 
+    Batch inputs are grouped by physical video and decoded in ascending frame
+    order to avoid unnecessary decoder seeks. Rows are restored to their input
+    order before ``collate_fn`` is called.
+
     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.
@@ -82,7 +86,7 @@ def __call__(self, rows):
         self._ensure_process_local_cache()
         single_row = isinstance(rows, Mapping)
         input_rows = [rows] if single_row else list(rows)
-        decoded_rows = [self._decode_row(row) for row in input_rows]
+        decoded_rows = self._decode_rows(input_rows)
         if single_row:
             return decoded_rows[0]
         return self._collate(decoded_rows)
@@ -108,7 +112,29 @@ def __getstate__(self):
         state["_owner_pid"] = None
         return state
 
-    def _decode_row(self, row):
+    def _decode_rows(self, rows):
+        decoded = [None] * len(rows)
+        grouped = OrderedDict()
+        for position, row in enumerate(rows):
+            output, descriptor = self._prepare_row(row)
+            if descriptor is None:
+                decoded[position] = output
+                continue
+            grouped.setdefault(descriptor.payload_descriptor, []).append(
+                (descriptor.frame_index, position, output)
+            )
+
+        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
+                )
+                decoded[position] = output
+        return decoded
+
+    def _prepare_row(self, row):
         if not isinstance(row, Mapping):
             raise ValueError("VideoFrameCollator expects row dictionaries.")
         if self.video_column not in row:
@@ -120,7 +146,7 @@ def _decode_row(self, row):
         output = dict(row)
         if raw is None:
             output[self.output_column] = None
-            return output
+            return output, None
         if hasattr(raw, "as_py"):
             raw = raw.as_py()
         if not VideoFrameDescriptor.is_video_frame_descriptor(raw):
@@ -138,11 +164,7 @@ def _decode_row(self, row):
                 "VideoFrameDescriptor without trailing bytes."
                 % self.video_column
             )
-        decoder = self._decoder(descriptor.payload_descriptor)
-        output[self.output_column] = self.decode_fn(
-            decoder, descriptor.frame_index, output
-        )
-        return output
+        return output, descriptor
 
     def _decoder(self, descriptor):
         resource = self._decoders.pop(descriptor, None)

paimon-python/pypaimon/tests/multimodal_video_test.py

@@ -23,6 +23,7 @@
 
 from pypaimon.common.file_io import FileIO
 from pypaimon.multimodal import VideoFrameCollator
+from pypaimon.multimodal.lerobot.dataset import _decode_video_rows
 from pypaimon.table.row.blob import VideoFrameDescriptor
 
 
@@ -94,6 +95,84 @@ def factory(stream):
         )
         self.assertEqual(descriptors[0], result[0]["video"])
 
+    def test_groups_and_sorts_frames_while_restoring_row_order(self):
+        descriptors = {
+            (video, frame): self._descriptor(
+                "episode-%s.mp4" % video,
+                ("video-%s" % video).encode(),
+                frame,
+            )
+            for video in ("one", "two")
+            for frame in (0, 1, 2, 3)
+        }
+        rows = [
+            {"request": "a", "video": descriptors["one", 3]},
+            {"request": "b", "video": descriptors["two", 2]},
+            {"request": "c", "video": descriptors["one", 1]},
+            {"request": "d", "video": descriptors["two", 0]},
+        ]
+        calls = []
+
+        def decode(decoder, frame, row):
+            video, decoded_frame = decoder.decode(frame)
+            calls.append((video, decoded_frame))
+            return row["request"], video, decoded_frame
+
+        collator = VideoFrameCollator(
+            self.table,
+            video_column="video",
+            decoder_factory=lambda stream: _Decoder(stream, []),
+            decode_fn=decode,
+            collate_fn=lambda decoded_rows: decoded_rows,
+        )
+        try:
+            result = collator(rows)
+        finally:
+            collator.close()
+
+        self.assertEqual(
+            [
+                (b"video-one", 1),
+                (b"video-one", 3),
+                (b"video-two", 0),
+                (b"video-two", 2),
+            ],
+            calls,
+        )
+        self.assertEqual(["a", "b", "c", "d"], [
+            row["request"] for row in result
+        ])
+        self.assertEqual(
+            [3, 2, 1, 0],
+            [row["frame"][2] for row in result],
+        )
+
+    def test_decodes_dataset_row_groups_in_one_batch(self):
+        row_groups = [
+            {4: {"request": "base", "video": b"base"}},
+            {1: {"request": "delta", "video": b"delta"}},
+        ]
+
+        class Collator:
+            video_column = "video"
+
+            def __init__(self):
+                self.calls = []
+
+            def __call__(self, rows):
+                self.calls.append([row["request"] for row in rows])
+                return [
+                    dict(row, video="decoded-" + row["request"])
+                    for row in rows
+                ]
+
+        collator = Collator()
+        _decode_video_rows(row_groups, [collator])
+
+        self.assertEqual([["base", "delta"]], collator.calls)
+        self.assertEqual("decoded-base", row_groups[0][4]["video"])
+        self.assertEqual("decoded-delta", row_groups[1][1]["video"])
+
     def test_evicts_least_recently_used_decoder(self):
         descriptors = [
             self._descriptor("episode-%d.mp4" % index, bytes([index]), index)